72 lines
2.7 KiB
Python
72 lines
2.7 KiB
Python
"""
|
||
LangChain tools for adding and searching a Qdrant knowledge base.
|
||
|
||
The module exposes two @tool‑decorated functions:
|
||
* add_to_knowledge_base(content: str, title: str = "document")
|
||
* search_knowledge_base(query: str, max_results: int = 5)
|
||
|
||
Both tools use an in‑memory Qdrant client and Ollama embeddings.
|
||
"""
|
||
import os
|
||
from typing import List
|
||
|
||
from langchain_core.documents import Document
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain_ollama import ChatOllama, OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from langchain.tools import tool
|
||
|
||
# --- Configuration -------------------------------------------------------
|
||
QDRANT_COLLECTION = "rag_agent"
|
||
EMBEDDING_MODEL = "nomic-embed-text"
|
||
LLM_MODEL = "llama3"
|
||
|
||
# Initialize embeddings and vector store (in‑memory Qdrant)
|
||
embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url="http://localhost:11434")
|
||
vector_store = QdrantVectorStore(client=None, collection_name=QDRANT_COLLECTION, embedding=embeddings)
|
||
|
||
# Ensure the collection exists (creates if not present)
|
||
try:
|
||
vector_store.client.create_collection(
|
||
name=QDRANT_COLLECTION,
|
||
vectors_config=dict(size=embeddings.embed_query("test").shape[0], distance="cosine"),
|
||
)
|
||
except Exception:
|
||
# Collection already exists – ignore
|
||
pass
|
||
|
||
# Text splitter for chunking documents
|
||
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||
|
||
@tool("Add a document to the knowledge base")
|
||
def add_to_knowledge_base(content: str, title: str = "document") -> str:
|
||
"""
|
||
Splits *content* into chunks, embeds them via Ollama and stores in Qdrant.
|
||
Returns a confirmation string with number of chunks added.
|
||
"""
|
||
# Split content
|
||
docs: List[Document] = splitter.split_text(content)
|
||
for i, chunk in enumerate(docs):
|
||
doc = Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
|
||
vector_store.add_documents([doc])
|
||
return f"Added {len(docs)} chunks from '{title}'."
|
||
|
||
@tool("Search the knowledge base")
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||
"""
|
||
Performs a semantic similarity search in Qdrant.
|
||
Returns a formatted string with top results and their metadata.
|
||
"""
|
||
docs = vector_store.similarity_search(query, k=max_results)
|
||
if not docs:
|
||
return "No relevant documents found."
|
||
lines: List[str] = []
|
||
for i, doc in enumerate(docs, 1):
|
||
title = doc.metadata.get("title", "unknown")
|
||
idx = doc.metadata.get("chunk_index", "?")
|
||
lines.append(f"{i}. [{title} - chunk {idx}]\n{doc.page_content[:200]}...")
|
||
return "\n\n".join(lines)
|
||
|
||
# Expose tool names for agent creation
|
||
TOOLS = [add_to_knowledge_base, search_knowledge_base]
|