Update rag_tools.py

This commit is contained in:
2026-06-04 20:02:50 +00:00
parent 0c4400b542
commit 06be1bf2e3
+21 -34
View File
@@ -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."