diff --git a/qdrant_store.py b/qdrant_store.py index fe37047..96b42ad 100644 --- a/qdrant_store.py +++ b/qdrant_store.py @@ -1,90 +1,50 @@ -""" -Qdrant vector store wrapper for adding and searching documents. +"""Module for interacting with Qdrant vector store using Ollama embeddings. """ -from pathlib import Path -from typing import List, Optional +from typing import List, Dict, Any -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 -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. +class QdrantStore: + """Wrapper around QdrantVectorStore. Parameters ---------- - content: str - Full text of the document. - title: str - Title or identifier for the document. + collection_name: str + Name of the Qdrant collection. + host: str + 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]: - """Semantic search in the knowledge base. + def add_documents(self, documents: List[str], titles: List[str] | None = None, metadatas: List[Dict[str, Any]] | None = None) -> None: + """Add documents to the collection. - Parameters - ---------- - query: str - Search query. - max_results: int - Number of top results to return. - """ - return vector_store.similarity_search(query, k=max_results) + Each document is added as a separate vector. If titles or metadatas are provided, they are attached. + """ + 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 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", -] + def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]: + """Semantic search returning list of dicts with 'content' and 'metadata'.""" + 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]