From 3e970602fd1ce5ccd1425e3b92a74a0b73ccc942 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 22:57:42 +0000 Subject: [PATCH] Update src/vector_store.py --- src/vector_store.py | 138 ++++++++++++++++++++++++-------------------- 1 file changed, 75 insertions(+), 63 deletions(-) diff --git a/src/vector_store.py b/src/vector_store.py index 85106f6..b739054 100644 --- a/src/vector_store.py +++ b/src/vector_store.py @@ -1,80 +1,92 @@ -""" -Vector store module using Qdrant and Ollama embeddings. +"""Vector store implementation using Qdrant and Ollama embeddings. + +This module provides a KnowledgeBase class that wraps a QdrantVectorStore and exposes +methods for adding documents and searching the knowledge base. """ from pathlib import Path -from typing import List, Dict +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 qdrant_client import QdrantClient -# Configuration -QDRANT_HOST = "localhost" -QDRANT_PORT = 6333 -COLLECTION_NAME = "knowledge_base" +class KnowledgeBase: + """A simple wrapper around QdrantVectorStore. -# Initialize embeddings and vector store -embeddings = OllamaEmbeddings(model="nomic-embed-text") -vector_store = QdrantVectorStore( - url=f"http://{QDRANT_HOST}:{QDRANT_PORT}", - collection_name=COLLECTION_NAME, - embeddings=embeddings, -) - -# Ensure collection exists -vector_store._ensure_collection_exists() - -# Text splitter -text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) - -def add_document(content: str, title: str) -> None: - """Add a document to the vector store. - - The document is split into chunks, embeddings are generated via Ollama, and - each chunk is stored with metadata containing the title. + Parameters + ---------- + collection_name: str + Name of the Qdrant collection to use. + host: str, optional + Qdrant host address. Defaults to ``localhost``. + port: int, optional + Qdrant port. Defaults to ``6333``. """ - # Split content into chunks - chunks = text_splitter.split_text(content) - # Prepare documents with metadata - documents = [] - for i, chunk in enumerate(chunks): - documents.append( - { - "content": chunk, - "metadata": {"title": title, "chunk_index": i}, - } + + def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333): + self.collection_name = collection_name + self.client = QdrantClient(host=host, port=port) + # Ensure the collection exists + self.client.recreate_collection( + collection_name=self.collection_name, + vectors_config={"size": 1024, "distance": "Cosine"}, ) - # Add to vector store - vector_store.add_documents(documents) -def search_documents(query: str, max_results: int = 5) -> List[Dict]: - """Semantic search in the vector store. + 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) - Returns a list of dictionaries containing the matched chunk and its metadata. - """ - results = vector_store.similarity_search_with_score(query, k=max_results) - # similarity_search_with_score returns list of tuples (Document, score) - return [ - { - "content": doc.page_content, - "metadata": doc.metadata, - "score": score, - } - for doc, score in results - ] + def add_document(self, content: str, title: str) -> None: + """Add a document to the knowledge base. -# Utility: load documents from a directory + The document is split into chunks, embedded, and stored in Qdrant. + """ + chunks = self.splitter.split_text(content) + documents: List[Document] = [] + for idx, chunk in enumerate(chunks): + meta = {"title": title, "chunk_index": idx} + documents.append(Document(page_content=chunk, metadata=meta)) + self.vector_store.add_documents(documents) -def load_documents_from_dir(directory: str) -> None: - """Load all .txt files from a directory and add them to the vector store.""" - path = Path(directory) - for file_path in path.rglob("*.txt"): - title = file_path.stem - content = file_path.read_text(encoding="utf-8") - add_document(content, title) + def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]: + """Search the knowledge base for the most relevant chunks. -# Example usage (commented out) + Returns a list of dictionaries containing the chunk content, title, and score. + """ + results = self.vector_store.similarity_search_with_score(query, k=max_results) + output = [] + for doc, score in results: + 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: + """Load all .txt files from a directory into the knowledge base. + + Parameters + ---------- + 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__": -# load_documents_from_dir("./docs") -# print(search_documents("what is langchain", 3)) \ No newline at end of file +# kb = KnowledgeBase() +# kb.load_from_directory("data") +# print(kb.search("What is Python?", 3)) \ No newline at end of file