agents/tools.py

"""
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
    }