diff --git a/rag_tools.py b/rag_tools.py index 7ec0017..d5df0f1 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,48 +1,35 @@ -""" -Tools for the RAG agent. +"""Tools for the RAG agent. + +Two tools are exposed via the ``@tool`` decorator: +* ``search_knowledge_base`` – semantic search in the Qdrant vector store. +* ``add_to_knowledge_base`` – add a document to the vector store. """ -from typing import List +from typing import List, Dict from langchain.tools import tool -from qdrant_store import search, add_document + +from .vector_store import add_document, search_text + +# --------------------------------------------------------------------------- +# Tool definitions +# --------------------------------------------------------------------------- @tool("search_knowledge_base") async def search_knowledge_base(query: str, max_results: int = 5) -> List[str]: - """Perform semantic search in the knowledge base. + """Return the top *max_results* relevant chunks for *query*. - Parameters - ---------- - query: str - Search query. - max_results: int - Number of results to return. - - Returns - ------- - List[str] - List of retrieved document snippets. + The function is asynchronous because LangChain expects async tools when + the agent runs in an async context. The underlying vector store calls + are synchronous, so we simply wrap the result. """ - docs = search(query, max_results) - return [doc.page_content for doc in docs] + return search_text(query, k=max_results) @tool("add_to_knowledge_base") async def add_to_knowledge_base(content: str, title: str) -> str: - """Add a new document to the knowledge base. + """Add *content* under *title* to the knowledge base. - Parameters - ---------- - content: str - Full text of the document. - title: str - Title or identifier for the document. - - Returns - ------- - str - Confirmation message. + Returns a confirmation string. """ - add_document(content, title) - return f"Document '{title}' added successfully." - -__all__ = ["search_knowledge_base", "add_to_knowledge_base"] + add_document(title, content) + return f"Document '{title}' added to knowledge base." \ No newline at end of file