Update src/vector_store.py
This commit is contained in:
+105
-72
@@ -1,92 +1,125 @@
|
||||
"""Vector store implementation using Qdrant and Ollama embeddings.
|
||||
"""Vector store and knowledge base 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.
|
||||
This module defines a `KnowledgeBase` class that manages a Qdrant collection, provides methods to add documents (with chunking) and perform semantic search.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
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 langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http.models import Distance, VectorParams
|
||||
|
||||
# Default configuration constants
|
||||
DEFAULT_COLLECTION_NAME = "knowledge_base"
|
||||
DEFAULT_VECTOR_SIZE = 3072 # size of nomic-embed-text embeddings
|
||||
DEFAULT_DISTANCE = Distance.COSINE
|
||||
DEFAULT_QDRANT_PATH = Path("./qdrant_data")
|
||||
|
||||
class KnowledgeBase:
|
||||
"""A simple wrapper around QdrantVectorStore.
|
||||
"""A thin wrapper around QdrantVectorStore.
|
||||
|
||||
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``.
|
||||
The class ensures that the collection is created only once and provides
|
||||
convenient methods for adding documents and performing semantic search.
|
||||
"""
|
||||
|
||||
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"},
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
def add_document(self, content: str, title: str) -> None:
|
||||
"""Add a document to the knowledge base.
|
||||
|
||||
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 search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
|
||||
"""Search the knowledge base for the most relevant chunks.
|
||||
|
||||
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.
|
||||
def __init__(
|
||||
self,
|
||||
collection_name: str = DEFAULT_COLLECTION_NAME,
|
||||
host: str | None = None,
|
||||
port: int | None = None,
|
||||
path: str | None = None,
|
||||
api_key: str | None = None,
|
||||
) -> None:
|
||||
"""Create or connect to a Qdrant collection.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
directory: str
|
||||
Path to the directory containing text files.
|
||||
collection_name: str
|
||||
Name of the Qdrant collection.
|
||||
host, port: str/int
|
||||
Optional host and port for a remote Qdrant instance.
|
||||
path: str
|
||||
Path for an on‑disk Qdrant instance (used in local mode).
|
||||
api_key: str
|
||||
API key for Qdrant Cloud.
|
||||
"""
|
||||
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__":
|
||||
# kb = KnowledgeBase()
|
||||
# kb.load_from_directory("data")
|
||||
# print(kb.search("What is Python?", 3))
|
||||
# Determine client connection.
|
||||
if host and port:
|
||||
self.client = QdrantClient(url=f"{host}:{port}")
|
||||
elif path:
|
||||
self.client = QdrantClient(path=path)
|
||||
else:
|
||||
# Default to a persistent on‑disk client.
|
||||
self.client = QdrantClient(path=str(DEFAULT_QDRANT_PATH))
|
||||
|
||||
self.collection_name = collection_name
|
||||
|
||||
# Create collection if it does not exist.
|
||||
if collection_name not in self.client.get_collections().collections:
|
||||
self.client.create_collection(
|
||||
collection_name=collection_name,
|
||||
vectors_config=VectorParams(size=DEFAULT_VECTOR_SIZE, distance=DEFAULT_DISTANCE),
|
||||
)
|
||||
|
||||
# Embedding model from Ollama.
|
||||
self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
|
||||
# Vector store wrapper.
|
||||
self.store = QdrantVectorStore(
|
||||
client=self.client,
|
||||
collection_name=collection_name,
|
||||
embedding=self.embeddings,
|
||||
)
|
||||
|
||||
# Text splitter for chunking.
|
||||
self.splitter = RecursiveCharacterTextSplitter(
|
||||
chunk_size=500, chunk_overlap=50, length_function=len
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------
|
||||
# Public API
|
||||
# ---------------------------------------------------------------------
|
||||
|
||||
def add_document(self, title: str, content: str) -> None:
|
||||
"""Add a document to the knowledge base.
|
||||
|
||||
The content is split into chunks, embedded, and stored.
|
||||
"""
|
||||
# Split into Document objects with metadata.
|
||||
docs = self.splitter.create_documents([content])
|
||||
for i, doc in enumerate(docs):
|
||||
# Attach metadata: title and chunk index.
|
||||
doc.metadata.update({"title": title, "chunk_index": i})
|
||||
# Add to store.
|
||||
self.store.add_documents(docs)
|
||||
|
||||
def search(self, query: str, limit: int = 5) -> List[Dict[str, Any]]:
|
||||
"""Perform a semantic search and return results.
|
||||
|
||||
Returns a list of dictionaries containing the chunk content and metadata.
|
||||
"""
|
||||
results = self.store.similarity_search(query, k=limit)
|
||||
output = []
|
||||
for doc in results:
|
||||
output.append(
|
||||
{
|
||||
"content": doc.page_content,
|
||||
"title": doc.metadata.get("title"),
|
||||
"chunk_index": doc.metadata.get("chunk_index"),
|
||||
}
|
||||
)
|
||||
return output
|
||||
|
||||
def get_all_documents(self) -> List[Document]:
|
||||
"""Return all documents stored in the collection."""
|
||||
return self.store.get_all_documents()
|
||||
|
||||
# End of src/vector_store.py
|
||||
Reference in New Issue
Block a user