diff --git a/src/vector_store.py b/src/vector_store.py new file mode 100644 index 0000000..85106f6 --- /dev/null +++ b/src/vector_store.py @@ -0,0 +1,80 @@ +""" +Vector store module using Qdrant and Ollama embeddings. +""" + +from pathlib import Path +from typing import List, Dict + +from langchain_ollama import OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore +from langchain_text_splitters import RecursiveCharacterTextSplitter + +# Configuration +QDRANT_HOST = "localhost" +QDRANT_PORT = 6333 +COLLECTION_NAME = "knowledge_base" + +# Initialize embeddings and vector store +embeddings = OllamaEmbeddings(model="nomic-embed-text") +vector_store = QdrantVectorStore( + url=f"http://{QDRANT_HOST}:{QDRANT_PORT}", + collection_name=COLLECTION_NAME, + embeddings=embeddings, +) + +# Ensure collection exists +vector_store._ensure_collection_exists() + +# Text splitter +text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +def add_document(content: str, title: str) -> None: + """Add a document to the vector store. + + The document is split into chunks, embeddings are generated via Ollama, and + each chunk is stored with metadata containing the title. + """ + # Split content into chunks + chunks = text_splitter.split_text(content) + # Prepare documents with metadata + documents = [] + for i, chunk in enumerate(chunks): + documents.append( + { + "content": chunk, + "metadata": {"title": title, "chunk_index": i}, + } + ) + # Add to vector store + vector_store.add_documents(documents) + +def search_documents(query: str, max_results: int = 5) -> List[Dict]: + """Semantic search in the vector store. + + Returns a list of dictionaries containing the matched chunk and its metadata. + """ + results = vector_store.similarity_search_with_score(query, k=max_results) + # similarity_search_with_score returns list of tuples (Document, score) + return [ + { + "content": doc.page_content, + "metadata": doc.metadata, + "score": score, + } + for doc, score in results + ] + +# Utility: load documents from a directory + +def load_documents_from_dir(directory: str) -> None: + """Load all .txt files from a directory and add them to the vector store.""" + path = Path(directory) + for file_path in path.rglob("*.txt"): + title = file_path.stem + content = file_path.read_text(encoding="utf-8") + add_document(content, title) + +# Example usage (commented out) +# if __name__ == "__main__": +# load_documents_from_dir("./docs") +# print(search_documents("what is langchain", 3)) \ No newline at end of file