init

vw-agentic-rag/service/graph/user_manual_tools.py

@@ -0,0 +1,77 @@
+"""
+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}