Update qdrant_store.py

qdrant_store.py

@@ -1,90 +1,50 @@
-"""
+"""Module for interacting with Qdrant vector store using Ollama embeddings.
-Qdrant vector store wrapper for adding and searching documents.
 """
 
-from pathlib import Path
+from typing import List, Dict, Any
-from typing import List, Optional
 
-from qdrant_client import QdrantClient
 from langchain_ollama import OllamaEmbeddings
-from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_qdrant import QdrantVectorStore
-from langchain.docstore.document import Document
+from qdrant_client import QdrantClient
 
-# Configuration
+class QdrantStore:
-QDRANT_HOST = "localhost"
+    """Wrapper around QdrantVectorStore.
-QDRANT_PORT = 6333
-COLLECTION_NAME = "knowledge_base"
-EMBEDDING_MODEL = "nomic-embed-text"
-
-# Initialize embedding model
-embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
-
-# Initialize Qdrant client
-client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
-
-# Ensure collection exists
-if COLLECTION_NAME not in client.get_collections().collections:
-    client.recreate_collection(collection_name=COLLECTION_NAME, vectors_config=client.get_default_vector_config())
-
-# Create vector store instance
-vector_store = QdrantVectorStore(
-    client=client,
-    collection_name=COLLECTION_NAME,
-    embeddings=embeddings,
-)
-
-# Text splitter
-splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-
-
-def add_document(content: str, title: str) -> None:
-    """Add a document to the vector store.
 
     Parameters
     ----------
-    content: str
+    collection_name: str
-        Full text of the document.
+        Name of the Qdrant collection.
-    title: str
+    host: str
-        Title or identifier for the document.
+        Qdrant host URL.
+    port: int
+        Qdrant port.
     """
-    # Split into chunks
-    chunks = splitter.split_text(content)
-    docs = [Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
-            for i, chunk in enumerate(chunks)]
-    # Add to vector store
-    vector_store.add_documents(docs)
 
+    def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
+        self.collection_name = collection_name
+        self.client = QdrantClient(host=host, port=port)
+        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
+        # Create collection if not exists
+        if collection_name not in self.client.get_collections().collections:
+            self.client.recreate_collection(collection_name=collection_name, vectors_config={"size": 512, "distance": "Cosine"})
+        self.store = QdrantVectorStore.from_existing_collection(
+            collection_name=collection_name,
+            embedding=self.embeddings,
+            client=self.client,
+        )
 
-def search(query: str, max_results: int = 5) -> List[Document]:
+    def add_documents(self, documents: List[str], titles: List[str] | None = None, metadatas: List[Dict[str, Any]] | None = None) -> None:
-    """Semantic search in the knowledge base.
+        """Add documents to the collection.
 
-    Parameters
+        Each document is added as a separate vector. If titles or metadatas are provided, they are attached.
-    ----------
-    query: str
-        Search query.
-    max_results: int
-        Number of top results to return.
         """
-    return vector_store.similarity_search(query, k=max_results)
+        if titles is None:
+            titles = [f"doc_{i}" for i in range(len(documents))]
+        if metadatas is None:
+            metadatas = [{} for _ in range(len(documents))]
+        self.store.add_texts(texts=documents, metadatas=metadatas, ids=titles)
 
-
+    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
-def load_directory(directory: str) -> None:
+        """Semantic search returning list of dicts with 'content' and 'metadata'."""
-    """Load all text files from a directory into the vector store.
+        results = self.store.similarity_search_with_score(query, k=k)
-
+        return [{"content": r[0].page_content, "metadata": r[0].metadata, "score": r[1]} for r in results]
-    Parameters
-    ----------
-    directory: str
-        Path to directory containing .txt files.
-    """
-    for file_path in Path(directory).rglob("*.txt"):
-        text = file_path.read_text(encoding="utf-8")
-        title = file_path.stem
-        add_document(text, title)
-
-# Expose public API
-__all__ = [
-    "add_document",
-    "search",
-    "load_directory",
-]