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 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") @tool("search_knowledge_base")
async def search_knowledge_base(query: str, max_results: int = 5) -> List[str]: 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 The function is asynchronous because LangChain expects async tools when
---------- the agent runs in an async context. The underlying vector store calls
query: str are synchronous, so we simply wrap the result.
Search query.
max_results: int
Number of results to return.
Returns
-------
List[str]
List of retrieved document snippets.
""" """
docs = search(query, max_results) return search_text(query, k=max_results)
return [doc.page_content for doc in docs]
@tool("add_to_knowledge_base") @tool("add_to_knowledge_base")
async def add_to_knowledge_base(content: str, title: str) -> str: 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 Returns a confirmation string.
----------
content: str
Full text of the document.
title: str
Title or identifier for the document.
Returns
-------
str
Confirmation message.
""" """
add_document(content, title) add_document(title, content)
return f"Document '{title}' added successfully." return f"Document '{title}' added to knowledge base."
__all__ = ["search_knowledge_base", "add_to_knowledge_base"]