38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
from typing import Any
|
|
|
|
from app.rag.retriever import retrieve_document_chunks
|
|
from app.tools._dispatch import NOT_HANDLED, ToolContext
|
|
|
|
TOOL_NAMES = frozenset({"search_documents"})
|
|
|
|
TOOL_DEFINITIONS: list[dict[str, Any]] = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "search_documents",
|
|
"description": "Семантический поиск по загруженным документам (RAG).",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Поисковый запрос"},
|
|
"limit": {"type": "integer", "description": "Макс. фрагментов"},
|
|
},
|
|
"required": ["query"],
|
|
},
|
|
},
|
|
},
|
|
]
|
|
|
|
|
|
async def execute(name: str, arguments: dict[str, Any], ctx: ToolContext) -> Any:
|
|
if name not in TOOL_NAMES:
|
|
return NOT_HANDLED
|
|
|
|
if name == "search_documents":
|
|
return await retrieve_document_chunks(
|
|
arguments.get("query", ""),
|
|
user_id=ctx.user_id,
|
|
top_k=int(arguments.get("limit") or 6),
|
|
)
|
|
return NOT_HANDLED
|