34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
"""
|
|
RAG tools for the agent.
|
|
|
|
search_knowledge_base and add_to_knowledge_base are decorated with @tool.
|
|
"""
|
|
import os
|
|
from typing import List, Dict
|
|
|
|
from langchain.tools import tool
|
|
from vector_store import vector_store
|
|
from chunker import split_text
|
|
|
|
@tool("Search knowledge base")
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
"""Semantic search in the vector store."""
|
|
results = vector_store.similarity_search(query, k=max_results)
|
|
if not results:
|
|
return "No relevant documents found."
|
|
out_lines = []
|
|
for i, res in enumerate(results, 1):
|
|
out_lines.append(f"{i}. {res['content'][:200]}... (distance: {res['distance']:.3f})")
|
|
return "\n".join(out_lines)
|
|
|
|
@tool("Add document to knowledge base")
|
|
def add_to_knowledge_base(content: str, title: str = "document") -> str:
|
|
"""Adds a text chunk to the vector store."""
|
|
# Split content into chunks
|
|
chunks = split_text(content)
|
|
docs = []
|
|
for idx, chunk in enumerate(chunks):
|
|
docs.append({"content": chunk, "metadata": {"title": title, "chunk_index": idx}})
|
|
vector_store.add_documents(docs)
|
|
return f"Added {len(chunks)} chunks to the knowledge base."
|