diff --git a/tools.py b/tools.py index 7342633..8d5bdd6 100644 --- a/tools.py +++ b/tools.py @@ -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"] \ No newline at end of file