Create agents/tools.py

agents/tools.py

@@ -0,0 +1,149 @@
+"""
+Tools for the pharmaceutical data management agents.
+"""
+
+import time
+from typing import Dict, Any, Callable
+
+def tool_list_tables(db):
+    """Tool to list available tables in the database."""
+    def _list_tables(args):
+        return db.get_tables()
+    
+    return {
+        "name": "list_tables",
+        "description": "List available tables in the database, categorized by pipeline stage",
+        "parameters": {
+            "type": "object",
+            "properties": {},
+            "required": []
+        },
+        "function": _list_tables
+    }
+
+def tool_describe_table(db):
+    """Tool to describe a table schema."""
+    def _describe_table(args):
+        table_name = args.get("table_name")
+        if not table_name:
+            return {"error": "Missing table_name parameter"}
+        
+        query = f"DESCRIBE {table_name}"
+        return db.execute_query(query)
+    
+    return {
+        "name": "describe_table",
+        "description": "Get the schema of a specific table",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "table_name": {
+                    "type": "string",
+                    "description": "Name of the table to describe"
+                }
+            },
+            "required": ["table_name"]
+        },
+        "function": _describe_table
+    }
+
+def tool_sample_table(db):
+    """Tool to get a sample of data from a table."""
+    def _sample_table(args):
+        table_name = args.get("table_name")
+        rows = args.get("rows", 5)
+        
+        if not table_name:
+            return {"error": "Missing table_name parameter"}
+        
+        return db.get_table_sample(table_name, rows)
+    
+    return {
+        "name": "sample_table",
+        "description": "Get a sample of rows from a specific table",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "table_name": {
+                    "type": "string",
+                    "description": "Name of the table to sample"
+                },
+                "rows": {
+                    "type": "integer",
+                    "description": "Number of rows to return (default: 5)"
+                }
+            },
+            "required": ["table_name"]
+        },
+        "function": _sample_table
+    }
+
+def tool_execute_query(db):
+    """Tool to execute a SQL query."""
+    def _execute_query(args):
+        query = args.get("query")
+        
+        if not query:
+            return {"error": "Missing query parameter"}
+        
+        return db.execute_query(query)
+    
+    return {
+        "name": "execute_query",
+        "description": "Execute a SQL query on the database",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "SQL query to execute"
+                }
+            },
+            "required": ["query"]
+        },
+        "function": _execute_query
+    }
+
+def tool_get_confidence(state_provider: Callable[[], Dict[str, Any]]):
+    """
+    Tool to calculate confidence score for the current plan.
+    Takes a function that returns the current state to ensure we have the latest.
+    """
+    def _get_confidence(args):
+        area = args.get("area", "overall")
+        state = state_provider()
+        user_intent = state.get("user_intent", {})
+        pipeline_plan = state.get("pipeline_plan", {})
+        
+        # This would implement a real confidence scoring system
+        # For demo, we'll return simulated confidence scores
+        
+        completeness = len(user_intent) / 5  # Simulate based on intent completeness
+        clarity = 0.7 if "target_tables" in pipeline_plan else 0.3
+        feasibility = 0.85  # High by default for demo
+        
+        if area == "intent":
+            return {"confidence": round(completeness * 100, 1), "area": "intent"}
+        elif area == "plan":
+            return {"confidence": round(clarity * 100, 1), "area": "plan"}
+        elif area == "feasibility":
+            return {"confidence": round(feasibility * 100, 1), "area": "feasibility"}
+        else:
+            overall = (completeness + clarity + feasibility) / 3
+            return {"confidence": round(overall * 100, 1), "area": "overall"}
+    
+    return {
+        "name": "get_confidence",
+        "description": "Calculate confidence score for the current plan or specific area",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "area": {
+                    "type": "string",
+                    "description": "Area to get confidence for (intent, plan, feasibility, or overall)",
+                    "enum": ["intent", "plan", "feasibility", "overall"]
+                }
+            }
+        },
+        "function": _get_confidence
+    }