Add qdrant_store.py

qdrant_store.py

@@ -0,0 +1,90 @@
+"""
+Qdrant vector store wrapper for adding and searching documents.
+"""
+
+from pathlib import Path
+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
+
+# Configuration
+QDRANT_HOST = "localhost"
+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
+        Full text of the document.
+    title: str
+        Title or identifier for the document.
+    """
+    # 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 search(query: str, max_results: int = 5) -> List[Document]:
+    """Semantic search in the knowledge base.
+
+    Parameters
+    ----------
+    query: str
+        Search query.
+    max_results: int
+        Number of top results to return.
+    """
+    return vector_store.similarity_search(query, k=max_results)
+
+
+def load_directory(directory: str) -> None:
+    """Load all text files from a directory into the vector store.
+
+    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",
+]