42 lines
1.6 KiB
Python
42 lines
1.6 KiB
Python
from langchain.tools import tool
|
|
from qdrant_client import QdrantClient
|
|
from qdrant_client.http import models
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
# Initialize global store
|
|
client = QdrantClient(url="http://localhost:6333")
|
|
collection_name = "rag_collection"
|
|
# Ensure collection exists
|
|
if collection_name not in client.get_collections().collections:
|
|
client.recreate_collection(
|
|
collection_name=collection_name,
|
|
vectors_config=models.VectorParams(size=384, distance=models.Distance.COSINE),
|
|
)
|
|
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
|
|
|
|
@tool("search_knowledge_base")
|
|
def search_knowledge_base(query: str, max_results: int = 5):
|
|
"""Semantic search in the knowledge base."""
|
|
query_vec = embeddings.embed_query(query)
|
|
results = client.search(
|
|
collection_name=collection_name,
|
|
query_vector=query_vec,
|
|
limit=max_results,
|
|
with_payload=True,
|
|
)
|
|
return [r.payload for r in results]
|
|
|
|
@tool("add_to_knowledge_base")
|
|
def add_to_knowledge_base(content: str, title: str):
|
|
"""Add a document to the knowledge base."""
|
|
chunks = splitter.split_text(content)
|
|
vectors = embeddings.embed_documents(chunks)
|
|
points = []
|
|
for i, (chunk, vec) in enumerate(zip(chunks, vectors)):
|
|
points.append(models.PointStruct(id=i, vector=vec, payload={"title": title, "chunk": chunk}))
|
|
client.upsert(collection_name=collection_name, points=models.Batch(points=points))
|
|
return f"Added {len(chunks)} chunks to the knowledge base."
|