39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
"""
|
|
Tools for the RAG agent.
|
|
|
|
Two tools: search_knowledge_base and add_to_knowledge_base.
|
|
"""
|
|
|
|
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 knowledge base.
|
|
|
|
Returns a formatted string of results.
|
|
"""
|
|
hits = vector_store.search(query, k=max_results)
|
|
if not hits:
|
|
return "No relevant documents found."
|
|
lines: List[str] = []
|
|
for i, hit in enumerate(hits, 1):
|
|
title = hit["metadata"].get("title", f"doc_{hit['id']}")
|
|
snippet = hit["document"][:200]
|
|
lines.append(f"{i}. {title}: {snippet}...")
|
|
return "\n".join(lines)
|
|
|
|
@tool("add_to_knowledge_base")
|
|
def add_to_knowledge_base(content: str, title: str = "document") -> str:
|
|
"""Add a document to the knowledge base.
|
|
|
|
Splits content into chunks and stores each with metadata.
|
|
"""
|
|
chunks = split_text(content)
|
|
for idx, chunk in enumerate(chunks):
|
|
doc_id = f"{title}_{idx}"
|
|
vector_store.add_document(doc_id=doc_id, text=chunk, metadata={"title": title})
|
|
return f"Added {len(chunks)} chunks from '{title}'."
|