51 lines
1.8 KiB
Python
51 lines
1.8 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). "
|
|
"Возвращает фрагменты с источниками: документ, страница, раздел, chunk_id."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {"type": "string", "description": "Поисковый запрос"},
|
|
"limit": {"type": "integer", "description": "Макс. фрагментов"},
|
|
"document_id": {
|
|
"type": "integer",
|
|
"description": "ID документа, если нужно искать только в одном документе",
|
|
},
|
|
},
|
|
"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":
|
|
document_id = arguments.get("document_id")
|
|
if document_id in ("", None):
|
|
document_id = None
|
|
elif document_id is not None:
|
|
document_id = int(document_id)
|
|
return await retrieve_document_chunks(
|
|
arguments.get("query", ""),
|
|
user_id=ctx.user_id,
|
|
top_k=int(arguments.get("limit") or 6),
|
|
document_id=document_id,
|
|
)
|
|
return NOT_HANDLED
|