init

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

@@ -0,0 +1,98 @@
+"""
+Tool definitions and schemas for the Agentic RAG system.
+This module contains all tool implementations and their corresponding schemas.
+"""
+
+import logging
+from typing import Dict, Any, List
+from langchain_core.tools import tool
+
+from ..retrieval.retrieval import AgenticRetrieval
+
+logger = logging.getLogger(__name__)
+
+
+# Tool Definitions using @tool decorator (following LangGraph best practices)
+@tool
+async def retrieve_standard_regulation(query: str) -> Dict[str, Any]:
+    """Search for attributes/metadata of China standards and regulations in automobile/manufacturing industry"""
+    async with AgenticRetrieval() as retrieval:
+        try:
+            result = await retrieval.retrieve_standard_regulation(
+                query=query
+            )
+            return {
+                "tool_name": "retrieve_standard_regulation",
+                "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"Retrieval error: {e}")
+            return {"error": str(e), "results_count": 0, "results": []}
+
+
+@tool  
+async def retrieve_doc_chunk_standard_regulation(query: str) -> Dict[str, Any]:
+    """Search for detailed document content chunks of China standards and regulations in automobile/manufacturing industry"""
+    async with AgenticRetrieval() as retrieval:
+        try:
+            result = await retrieval.retrieve_doc_chunk_standard_regulation(
+                query=query
+            )
+            return {
+                "tool_name": "retrieve_doc_chunk_standard_regulation",
+                "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"Doc chunk retrieval error: {e}")
+            return {"error": str(e), "results_count": 0, "results": []}
+
+
+# Available tools list
+tools = [retrieve_standard_regulation, retrieve_doc_chunk_standard_regulation]
+
+
+def get_tool_schemas() -> List[Dict[str, Any]]:
+    """
+    Generate tool schemas for LLM function calling.
+    
+    Returns:
+        List of tool schemas in OpenAI function calling format
+    """
+    tools.append();
+
+    tool_schemas = []
+    for tool in 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_tools_by_name() -> Dict[str, Any]:
+    """
+    Create a mapping of tool names to tool functions.
+    
+    Returns:
+        Dictionary mapping tool names to tool functions
+    """
+    return {tool.name: tool for tool in tools}