Update src/vector_store.py

src/vector_store.py

@@ -1,80 +1,92 @@
-"""
-Vector store module using Qdrant and Ollama embeddings.
+"""Vector store implementation using Qdrant and Ollama embeddings.
+
+This module provides a KnowledgeBase class that wraps a QdrantVectorStore and exposes
+methods for adding documents and searching the knowledge base.
 """
 
 from pathlib import Path
-from typing import List, Dict
+from typing import List, Dict, Any
 
 from langchain_ollama import OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
+from qdrant_client import QdrantClient
 
-# Configuration
-QDRANT_HOST = "localhost"
-QDRANT_PORT = 6333
-COLLECTION_NAME = "knowledge_base"
+class KnowledgeBase:
+    """A simple wrapper around QdrantVectorStore.
 
-# 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.
+    Parameters
+    ----------
+    collection_name: str
+        Name of the Qdrant collection to use.
+    host: str, optional
+        Qdrant host address. Defaults to ``localhost``.
+    port: int, optional
+        Qdrant port. Defaults to ``6333``.
     """
-    # 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},
-            }
+
+    def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
+        self.collection_name = collection_name
+        self.client = QdrantClient(host=host, port=port)
+        # Ensure the collection exists
+        self.client.recreate_collection(
+            collection_name=self.collection_name,
+            vectors_config={"size": 1024, "distance": "Cosine"},
         )
-    # 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.
+        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
+        self.vector_store = QdrantVectorStore(
+            client=self.client,
+            collection_name=self.collection_name,
+            embeddings=self.embeddings,
+        )
+        self.splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
 
-    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
-    ]
+    def add_document(self, content: str, title: str) -> None:
+        """Add a document to the knowledge base.
 
-# Utility: load documents from a directory
+        The document is split into chunks, embedded, and stored in Qdrant.
+        """
+        chunks = self.splitter.split_text(content)
+        documents: List[Document] = []
+        for idx, chunk in enumerate(chunks):
+            meta = {"title": title, "chunk_index": idx}
+            documents.append(Document(page_content=chunk, metadata=meta))
+        self.vector_store.add_documents(documents)
 
-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)
+    def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+        """Search the knowledge base for the most relevant chunks.
 
-# Example usage (commented out)
+        Returns a list of dictionaries containing the chunk content, title, and score.
+        """
+        results = self.vector_store.similarity_search_with_score(query, k=max_results)
+        output = []
+        for doc, score in results:
+            output.append({
+                "content": doc.page_content,
+                "title": doc.metadata.get("title"),
+                "chunk_index": doc.metadata.get("chunk_index"),
+                "score": score,
+            })
+        return output
+
+    def load_from_directory(self, directory: str) -> None:
+        """Load all .txt files from a directory into the knowledge base.
+
+        Parameters
+        ----------
+        directory: str
+            Path to the directory containing text files.
+        """
+        path = Path(directory)
+        for file_path in path.rglob("*.txt"):
+            title = file_path.stem
+            content = file_path.read_text(encoding="utf-8")
+            self.add_document(content, title)
+
+# Example usage (uncomment for quick test)
 # if __name__ == "__main__":
-#     load_documents_from_dir("./docs")
-#     print(search_documents("what is langchain", 3))
+#     kb = KnowledgeBase()
+#     kb.load_from_directory("data")
+#     print(kb.search("What is Python?", 3))