Files
task-6a02e23da6fe2e4ac16acf65/rag_tools.py
T
2026-06-04 20:02:50 +00:00

35 lines
1.2 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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, Dict
from langchain.tools import tool
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]:
"""Return the top *max_results* relevant chunks for *query*.
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.
"""
return search_text(query, k=max_results)
@tool("add_to_knowledge_base")
async def add_to_knowledge_base(content: str, title: str) -> str:
"""Add *content* under *title* to the knowledge base.
Returns a confirmation string.
"""
add_document(title, content)
return f"Document '{title}' added to knowledge base."