55 lines
1.5 KiB
Python
55 lines
1.5 KiB
Python
"""Tools for the RAG agent.
|
|
|
|
This module defines two LangChain tools that interact with the
|
|
``KnowledgeBase`` defined in :mod:`src.vector_store`.
|
|
|
|
The tools are decorated with ``@tool`` from ``langchain.tools`` so that
|
|
the agent can invoke them automatically.
|
|
"""
|
|
|
|
from langchain.tools import tool
|
|
from .vector_store import kb
|
|
|
|
@tool("search_knowledge_base")
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
"""Search the local knowledge base.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
The search query.
|
|
max_results: int, optional
|
|
Limit of results to return.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
A formatted string with the search results.
|
|
"""
|
|
results = kb.search(query, limit=max_results)
|
|
if not results:
|
|
return "No relevant documents found."
|
|
lines = []
|
|
for i, res in enumerate(results, 1):
|
|
title = res["metadata"].get("title", "Untitled")
|
|
lines.append(f"{i}. Title: {title}\nContent: {res['page_content']}\n")
|
|
return "\n".join(lines)
|
|
|
|
@tool("add_to_knowledge_base")
|
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
"""Add a new document to the knowledge base.
|
|
|
|
Parameters
|
|
----------
|
|
content: str
|
|
The full text of the document.
|
|
title: str
|
|
A short title for the document.
|
|
|
|
Returns
|
|
-------
|
|
str
|
|
Confirmation message.
|
|
"""
|
|
kb.add_document(title=title, content=content)
|
|
return f"Document '{title}' added to the knowledge base." |