diff --git a/vector_store.py b/vector_store.py new file mode 100644 index 0000000..45603f7 --- /dev/null +++ b/vector_store.py @@ -0,0 +1,93 @@ +"""Vector store implementation using Qdrant. + +This module provides functions to create a Qdrant vector store backed by Ollama embeddings +and to load documents from a directory into the store. The store is persisted in a local +directory and can be reused across runs. +""" + +from pathlib import Path +from typing import List + +from langchain_qdrant import Qdrant +from langchain_ollama import OllamaEmbeddings +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.docstore.document import Document + +__all__ = ["create_vectorstore", "load_documents"] + +def create_vectorstore(persist_directory: str = "./qdrant_db") -> Qdrant: + """Create or load a Qdrant vector store. + + Parameters + ---------- + persist_directory: str + Path to the directory where Qdrant will store its data. The directory + will be created if it does not exist. + + Returns + ------- + Qdrant + A Qdrant vector store instance. + """ + # Ensure the directory exists + Path(persist_directory).mkdir(parents=True, exist_ok=True) + + # Use Ollama embeddings (nomic-embed-text) for semantic similarity + embeddings = OllamaEmbeddings(model="nomic-embed-text") + + # Qdrant can run in local mode when ``location`` is provided. + # The ``url`` is set to the default local address. + return Qdrant( + collection_name="documents", + embedding=embeddings, + url="http://localhost:6333", # Qdrant server address + location=persist_directory, + ) + +def _load_text_files(directory: str) -> List[str]: + """Recursively read all .txt and .md files from *directory*. + + Parameters + ---------- + directory: str + Root directory to search for documents. + + Returns + ------- + List[str] + List of file contents. + """ + texts: List[str] = [] + for path in Path(directory).rglob("*"): + if path.suffix.lower() in {".txt", ".md"}: + try: + with open(path, "r", encoding="utf-8") as f: + texts.append(f.read()) + except Exception as exc: # pragma: no cover - defensive + print(f"Could not read {path}: {exc}") + return texts + +def load_documents(directory: str, vectorstore: Qdrant) -> None: + """Load documents from *directory* into the provided *vectorstore*. + + The function performs chunking via :class:`RecursiveCharacterTextSplitter` + before adding the chunks to the vector store. + """ + raw_texts = _load_text_files(directory) + if not raw_texts: + print(f"No .txt or .md files found in {directory}") + return + + # Chunk each document into manageable pieces + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + chunks: List[str] = [] + for text in raw_texts: + chunks.extend(splitter.split_text(text)) + + # Convert to LangChain Document objects + documents = [Document(page_content=chunk) for chunk in chunks] + + # Add documents to Qdrant. The underlying Qdrant client will handle + # persistence automatically. + vectorstore.add_documents(documents) + print(f"Loaded {len(documents)} chunks into Qdrant.") \ No newline at end of file