55 lines
2.3 KiB
Python
55 lines
2.3 KiB
Python
from pathlib import Path
|
||
from uuid import uuid4
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from qdrant_client import QdrantClient
|
||
from qdrant_client.http.models import Distance, VectorParams
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_core.documents import Document
|
||
|
||
class KnowledgeBase:
|
||
"""A simple wrapper around Qdrant for storing and searching documents."""
|
||
def __init__(self, collection_name: str = "rag_kb", persist_path: str = "qdrant_data"):
|
||
# Initialize Qdrant client (in‑memory or on‑disk)
|
||
self.client = QdrantClient(path=persist_path)
|
||
# Create collection with cosine distance and 3072‑dim vectors (nomic‑embed‑text output)
|
||
self.client.create_collection(
|
||
collection_name=collection_name,
|
||
vectors_config=VectorParams(size=3072, distance=Distance.COSINE),
|
||
)
|
||
# Create vector store wrapper
|
||
self.vector_store = QdrantVectorStore(
|
||
client=self.client,
|
||
collection_name=collection_name,
|
||
embedding=OllamaEmbeddings(model="nomic-embed-text"),
|
||
)
|
||
# Text splitter for chunking documents
|
||
self.splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
||
|
||
def add_document(self, title: str, content: str):
|
||
"""Add a document to the knowledge base.
|
||
|
||
The content is split into chunks, each chunk is embedded and stored.
|
||
"""
|
||
# Create Document objects with metadata
|
||
docs = self.splitter.create_documents([content], metadata={"title": title})
|
||
# Add to vector store
|
||
self.vector_store.add_documents(docs)
|
||
|
||
def search(self, query: str, limit: int = 10):
|
||
"""Search the knowledge base for the most relevant documents.
|
||
|
||
Returns a list of dictionaries containing page_content, metadata and score.
|
||
"""
|
||
results = self.vector_store.similarity_search_with_score(query, k=limit)
|
||
return [
|
||
{
|
||
"page_content": doc.page_content,
|
||
"metadata": doc.metadata,
|
||
"score": score,
|
||
}
|
||
for doc, score in results
|
||
]
|
||
|
||
# Global singleton instance used by tools and the agent
|
||
kb = KnowledgeBase() |