Create agents/executor.py

agents/executor.py

@@ -0,0 +1,133 @@
+"""
+Executor agent for pharmaceutical data management.
+This agent executes SQL queries and reports on the results.
+"""
+
+import time
+import json
+from typing import Dict, Any, List
+from langgraph.graph.message import add_messages
+from anthropic.types import MessageParam
+
+def executor_agent(anthropic_client, db, state: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Agent that executes the SQL queries and reports results.
+    
+    Args:
+        anthropic_client: The Anthropic client for calling Claude API
+        db: The database connection
+        state: Current state of the agent workflow
+    
+    Returns:
+        Updated state
+    """
+    # Get current messages and SQL queries
+    messages = state["messages"]
+    sql_queries = state["sql_queries"]
+    
+    # Add agent-specific instructions
+    system_message = """
+    You are an AI assistant specializing in executing and validating pharmaceutical data pipelines.
+    Your job is to execute the SQL queries and report on the results.
+    
+    For each query:
+    1. Execute it and capture the results
+    2. Validate the data (check for issues, errors, unexpected results)
+    3. Provide a clear summary of what was accomplished
+    
+    Tag your complete response with EXECUTION_COMPLETE when all queries are processed.
+    """
+    
+    # Execute the queries
+    execution_results = _execute_queries(db, sql_queries)
+    
+    # Calculate overall success
+    success_rate = sum(1 for r in execution_results if r["success"]) / len(execution_results) if execution_results else 0
+    
+    # Prepare context for Claude
+    execution_context = json.dumps(execution_results, indent=2)
+    
+    # Prepare prompt for Claude
+    prompt_messages = [
+        MessageParam(role="system", content=system_message),
+        *[MessageParam(role=m["role"], content=m["content"]) for m in messages],
+        MessageParam(role="user", content=f"Here are the execution results of the SQL queries. Please analyze and report on them.\n\n{execution_context}")
+    ]
+    
+    # Call Claude API
+    response = anthropic_client.messages.create(
+        model="claude-3-7-sonnet-20250219",
+        messages=prompt_messages,
+        max_tokens=2000
+    )
+    
+    # Extract the response
+    agent_response = response.content[0].text
+    
+    # Update state with execution results
+    new_state = state.copy()
+    new_state["execution_results"] = {
+        "queries_executed": len(execution_results),
+        "success_rate": success_rate,
+        "details": execution_results,
+        "summary": agent_response.replace("EXECUTION_COMPLETE", "").strip(),
+        "completed_at": time.time()
+    }
+    
+    # Calculate confidence scores
+    confidence_scores = {
+        "intent_understanding": 0.95,  # High by default for demo
+        "plan_quality": 0.85 if success_rate > 0.8 else 0.6,
+        "execution_success": success_rate,
+        "overall": (0.95 + (0.85 if success_rate > 0.8 else 0.6) + success_rate) / 3
+    }
+    new_state["confidence_scores"] = confidence_scores
+    
+    # Update status
+    new_state["status"] = "complete"
+    new_state["current_agent"] = "complete"
+    
+    # Add agent's response to messages
+    new_messages = add_messages(state, [
+        {"role": "assistant", "content": agent_response.replace("EXECUTION_COMPLETE", "").strip()}
+    ])
+    new_state["messages"] = new_messages
+    
+    return new_state
+
+def _execute_queries(db, sql_queries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Execute SQL queries against the database.
+    
+    Args:
+        db: The database connection
+        sql_queries: List of SQL queries to execute
+    
+    Returns:
+        List of execution results
+    """
+    execution_results = []
+    for query in sql_queries:
+        try:
+            # Execute the query
+            result = db.execute_query(query["sql"])
+            
+            # Add to results
+            execution_results.append({
+                "query_name": query["name"],
+                "success": "error" not in result,
+                "result_summary": "Successfully processed data" if "error" not in result else result["error"],
+                "row_count": len(result.get("data", [])) if "data" in result else 0,
+                "executed_at": time.time()
+            })
+        except Exception as e:
+            # Handle execution errors
+            execution_results.append({
+                "query_name": query["name"],
+                "success": False,
+                "result_summary": f"Error: {str(e)}",
+                "row_count": 0,
+                "executed_at": time.time()
+            })
+    
+    return execution_results