Update src/vector_store.py

src/vector_store.py

@@ -1,92 +1,125 @@
-"""Vector store implementation using Qdrant and Ollama embeddings.
+"""Vector store and knowledge base implementation using Qdrant and Ollama embeddings.
 
-This module provides a KnowledgeBase class that wraps a QdrantVectorStore and exposes
+This module defines a `KnowledgeBase` class that manages a Qdrant collection, provides methods to add documents (with chunking) and perform semantic search.
-methods for adding documents and searching the knowledge base.
 """
 
+from __future__ import annotations
+
+import os
 from pathlib import Path
 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 langchain_text_splitters import RecursiveCharacterTextSplitter
 from qdrant_client import QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
+
+# Default configuration constants
+DEFAULT_COLLECTION_NAME = "knowledge_base"
+DEFAULT_VECTOR_SIZE = 3072  # size of nomic-embed-text embeddings
+DEFAULT_DISTANCE = Distance.COSINE
+DEFAULT_QDRANT_PATH = Path("./qdrant_data")
 
 class KnowledgeBase:
-    """A simple wrapper around QdrantVectorStore.
+    """A thin wrapper around QdrantVectorStore.
+
+    The class ensures that the collection is created only once and provides
+    convenient methods for adding documents and performing semantic search.
+    """
+
+    def __init__(
+        self,
+        collection_name: str = DEFAULT_COLLECTION_NAME,
+        host: str | None = None,
+        port: int | None = None,
+        path: str | None = None,
+        api_key: str | None = None,
+    ) -> None:
+        """Create or connect to a Qdrant collection.
 
         Parameters
         ----------
         collection_name: str
-        Name of the Qdrant collection to use.
+            Name of the Qdrant collection.
-    host: str, optional
+        host, port: str/int
-        Qdrant host address. Defaults to ``localhost``.
+            Optional host and port for a remote Qdrant instance.
-    port: int, optional
+        path: str
-        Qdrant port. Defaults to ``6333``.
+            Path for an on‑disk Qdrant instance (used in local mode).
+        api_key: str
+            API key for Qdrant Cloud.
         """
 
-    def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
+        # Determine client connection.
+        if host and port:
+            self.client = QdrantClient(url=f"{host}:{port}")
+        elif path:
+            self.client = QdrantClient(path=path)
+        else:
+            # Default to a persistent on‑disk client.
+            self.client = QdrantClient(path=str(DEFAULT_QDRANT_PATH))
+
         self.collection_name = collection_name
-        self.client = QdrantClient(host=host, port=port)
+
-        # Ensure the collection exists
+        # Create collection if it does not exist.
-        self.client.recreate_collection(
+        if collection_name not in self.client.get_collections().collections:
-            collection_name=self.collection_name,
+            self.client.create_collection(
-            vectors_config={"size": 1024, "distance": "Cosine"},
+                collection_name=collection_name,
+                vectors_config=VectorParams(size=DEFAULT_VECTOR_SIZE, distance=DEFAULT_DISTANCE),
             )
 
+        # Embedding model from Ollama.
         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)
 
-    def add_document(self, content: str, title: str) -> None:
+        # Vector store wrapper.
+        self.store = QdrantVectorStore(
+            client=self.client,
+            collection_name=collection_name,
+            embedding=self.embeddings,
+        )
+
+        # Text splitter for chunking.
+        self.splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500, chunk_overlap=50, length_function=len
+        )
+
+    # ---------------------------------------------------------------------
+    # Public API
+    # ---------------------------------------------------------------------
+
+    def add_document(self, title: str, content: str) -> None:
         """Add a document to the knowledge base.
 
-        The document is split into chunks, embedded, and stored in Qdrant.
+        The content is split into chunks, embedded, and stored.
         """
-        chunks = self.splitter.split_text(content)
+        # Split into Document objects with metadata.
-        documents: List[Document] = []
+        docs = self.splitter.create_documents([content])
-        for idx, chunk in enumerate(chunks):
+        for i, doc in enumerate(docs):
-            meta = {"title": title, "chunk_index": idx}
+            # Attach metadata: title and chunk index.
-            documents.append(Document(page_content=chunk, metadata=meta))
+            doc.metadata.update({"title": title, "chunk_index": i})
-        self.vector_store.add_documents(documents)
+        # Add to store.
+        self.store.add_documents(docs)
 
-    def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
+    def search(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
-        """Search the knowledge base for the most relevant chunks.
+        """Perform a semantic search and return results.
 
-        Returns a list of dictionaries containing the chunk content, title, and score.
+        Returns a list of dictionaries containing the chunk content and metadata.
         """
-        results = self.vector_store.similarity_search_with_score(query, k=max_results)
+        results = self.store.similarity_search(query, k=limit)
         output = []
-        for doc, score in results:
+        for doc in results:
-            output.append({
+            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:
+    def get_all_documents(self) -> List[Document]:
-        """Load all .txt files from a directory into the knowledge base.
+        """Return all documents stored in the collection."""
+        return self.store.get_all_documents()
 
-        Parameters
+# End of src/vector_store.py
-        ----------
-        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__":
-#     kb = KnowledgeBase()
-#     kb.load_from_directory("data")
-#     print(kb.search("What is Python?", 3))