Files
catonline_ai/vw-agentic-rag/service/graph/user_manual_tools.py
2025-09-26 17:15:54 +08:00

78 lines
2.3 KiB
Python

"""
User manual specific tools for the Agentic RAG system.
This module contains tools specifically for user manual retrieval and processing.
"""
import logging
from typing import Dict, Any, List
from langchain_core.tools import tool
from ..retrieval.retrieval import AgenticRetrieval
logger = logging.getLogger(__name__)
# User Manual Tools
@tool
async def retrieve_system_usermanual(query: str) -> Dict[str, Any]:
"""Search for document content chunks of user manual of this system(CATOnline)"""
async with AgenticRetrieval() as retrieval:
try:
result = await retrieval.retrieve_doc_chunk_user_manual(
query=query
)
return {
"tool_name": "retrieve_system_usermanual",
"results_count": len(result.results),
"results": result.results, # Already dict objects, no need for model_dump()
"took_ms": result.took_ms
}
except Exception as e:
logger.error(f"User manual retrieval error: {e}")
return {"error": str(e), "results_count": 0, "results": []}
# User manual tools list
user_manual_tools = [retrieve_system_usermanual]
def get_user_manual_tool_schemas() -> List[Dict[str, Any]]:
"""
Generate tool schemas for user manual tools.
Returns:
List of tool schemas in OpenAI function calling format
"""
tool_schemas = []
for tool in user_manual_tools:
schema = {
"type": "function",
"function": {
"name": tool.name,
"description": tool.description,
"parameters": {
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Search query for retrieving relevant information"
}
},
"required": ["query"]
}
}
}
tool_schemas.append(schema)
return tool_schemas
def get_user_manual_tools_by_name() -> Dict[str, Any]:
"""
Create a mapping of user manual tool names to tool functions.
Returns:
Dictionary mapping tool names to tool functions
"""
return {tool.name: tool for tool in user_manual_tools}