Add src/vector_store.py

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))