Add vector_store.py

vector_store.py

@@ -0,0 +1,93 @@
+"""Module for interacting with Qdrant vector store using Ollama embeddings.
+
+This module provides a simple wrapper around LangChain's QdrantVectorStore.
+It handles initialization, adding documents (with chunking), and semantic search.
+"""
+
+from pathlib import Path
+from typing import List
+
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
+class QdrantStore:
+    """A wrapper around QdrantVectorStore.
+
+    Parameters
+    ----------
+    collection_name: str
+        Name of the collection in Qdrant. Defaults to "knowledge_base".
+    url: str
+        URL of the Qdrant instance. Defaults to "http://localhost:6333".
+    """
+
+    def __init__(self, collection_name: str = "knowledge_base", url: str = "http://localhost:6333"):
+        self.collection_name = collection_name
+        self.url = url
+        # Use Ollama embeddings model
+        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
+        # Initialize an empty vector store (will create collection if not exists)
+        self.vector_store = QdrantVectorStore.from_texts(
+            [], self.embeddings, url=self.url, collection_name=self.collection_name
+        )
+
+    def _split_text(self, text: str) -> List[str]:
+        """Split a long text into manageable chunks.
+
+        Uses RecursiveCharacterTextSplitter with a chunk size of 1000 characters and
+        an overlap of 200 characters to preserve context.
+        """
+        splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        return splitter.split_text(text)
+
+    def add_document(self, content: str, title: str) -> None:
+        """Add a document to the vector store.
+
+        Parameters
+        ----------
+        content: str
+            Full text of the document.
+        title: str
+            Title or identifier for the document.
+        """
+        chunks = self._split_text(content)
+        # Prepare metadata for each chunk
+        metadatas = [{"title": title, "chunk_index": i} for i in range(len(chunks))]
+        self.vector_store.add_texts(chunks, metadatas)
+
+    def search(self, query: str, max_results: int = 5):
+        """Perform a semantic search.
+
+        Parameters
+        ----------
+        query: str
+            The search query.
+        max_results: int
+            Number of top results to return.
+
+        Returns
+        -------
+        List[Document]
+            List of LangChain Document objects.
+        """
+        return self.vector_store.similarity_search(query, k=max_results)
+
+    def load_documents_from_dir(self, directory: str) -> None:
+        """Load all .txt files from a directory into the vector store.
+
+        Parameters
+        ----------
+        directory: str
+            Path to the directory containing text files.
+        """
+        path = Path(directory)
+        for file_path in path.rglob("*.txt"):
+            with file_path.open("r", encoding="utf-8") as f:
+                content = f.read()
+            title = file_path.stem
+            self.add_document(content, title)
+
+# Singleton instance used by tools
+store = QdrantStore()