Update tools.py

This commit is contained in:
2026-06-04 20:03:08 +00:00
parent 60021947c8
commit 51c07404cf
+25 -29
View File
@@ -1,51 +1,47 @@
"""Tools for the RAG agent.
"""Tool definitions for the RAG agent.
Two tools are defined:
1. search_knowledge_base performs semantic search in the vector store.
2. add_to_knowledge_base adds a new document to the vector store.
The tools are simple wrappers around the vector store functions defined in
`vector_store.py`. They are decorated with `@tool` from LangChain so that the
agent can call them.
"""
from typing import List
from typing import Any
from langchain.tools import tool
from langchain.schema import Document
from vector_store import store
from .vector_store import add_document, search
@tool("search_knowledge_base")
async def search_knowledge_base(query: str, max_results: int = 5) -> List[Document]:
"""Search the knowledge base for relevant chunks.
def search_knowledge_base(query: str, max_results: int = 5) -> Any:
"""Semantic search in the knowledge base.
Parameters
----------
query: str
The search query.
max_results: int
Number of top results to return.
Returns
-------
List[Document]
List of documents returned by Qdrant similarity search.
The user query.
max_results: int, optional
Number of top results to return. Defaults to 5.
"""
return store.search(query, max_results)
results = search(query, max_results)
# Convert results to a readable string
formatted = "\n".join(
f"{i+1}. [{res['metadata'].get('title', 'Unknown')}] {res['content'][:200]}"
for i, res in enumerate(results)
)
return formatted if formatted else "No relevant documents found."
@tool("add_to_knowledge_base")
async def add_to_knowledge_base(content: str, title: str) -> str:
def add_to_knowledge_base(content: str, title: str) -> Any:
"""Add a new document to the knowledge base.
Parameters
----------
content: str
Full text of the document.
The full text of the document.
title: str
Title or identifier for the document.
Returns
-------
str
Confirmation message.
A short title that will be stored as metadata.
"""
store.add_document(content, title)
add_document(content, title)
return f"Document '{title}' added to the knowledge base."
__all__ = ["search_knowledge_base", "add_to_knowledge_base"]