30 lines
1.5 KiB
Python
30 lines
1.5 KiB
Python
from langchain.tools import tool
|
|
from qdrant_client import QdrantClient
|
|
from langchain.embeddings.ollama import OllamaEmbeddings
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
import os
|
|
|
|
# Initialize Qdrant client and collection
|
|
qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333")
|
|
client = QdrantClient(url=qdrant_url)
|
|
collection_name = "knowledge_base"
|
|
if not client.has_collection(collection_name):
|
|
client.create_collection(name=collection_name, vectors_config={"size": 384, "distance": "Cosine"})
|
|
|
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
|
|
|
@tool("search_knowledge_base", description="Semantic search in knowledge base")
|
|
def search_knowledge_base(query: str, max_results: int = 5):
|
|
vector = embeddings.embed_query(query)
|
|
results = client.search(collection_name=collection_name, query_vector=vector, limit=max_results)
|
|
return [hit.payload["text"] for hit in results]
|
|
|
|
@tool("add_to_knowledge_base", description="Add document to knowledge base")
|
|
def add_to_knowledge_base(content: str, title: str):
|
|
docs = text_splitter.split_text(content)
|
|
vectors = embeddings.embed_documents(docs)
|
|
ids = [title + f"_{i}" for i in range(len(docs))]
|
|
client.upsert(collection_name=collection_name, points=[{"id": id_, "vector": vec, "payload": {"text": doc}} for id_, vec, doc in zip(ids, vectors, docs)])
|
|
return f"Added {len(docs)} chunks to knowledge base"
|