Update agents/executor.py

agents/executor.py

@@ -1,3 +1,4 @@
+# executor_agent.py - Fixed version
 """
 Executor agent for pharmaceutical data management.
 This agent executes SQL queries and reports on the results.
@@ -25,7 +26,7 @@ def executor_agent(anthropic_client, db, state: Dict[str, Any]) -> Dict[str, Any
     messages = state.get("messages", [])
     sql_queries = state.get("sql_queries", [])
     
-    # Add agent-specific instructions as system parameter
+    # 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.
@@ -47,65 +48,75 @@ def executor_agent(anthropic_client, db, state: Dict[str, Any]) -> Dict[str, Any
     # Prepare context for Claude
     execution_context = json.dumps(execution_results, indent=2)
     
-    # Convert state messages to the format expected by Anthropic API
-    # Excluding any messages with role "system"
-    anthropic_messages = []
+    # Format messages for the Anthropic API
+    formatted_messages = []
     for msg in messages:
-        if msg["role"] != "system":  # Skip system messages
-            anthropic_messages.append(MessageParam(
+        if isinstance(msg, dict) and "role" in msg and "content" in msg:
+            formatted_messages.append(MessageParam(
                 role=msg["role"],
                 content=msg["content"]
             ))
     
     # Add final user message with context
-    anthropic_messages.append(
+    formatted_messages.append(
         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 with system message as separate parameter
-    response = anthropic_client.messages.create(
-        model="claude-3-7-sonnet-20250219",
-        system=system_message,
-        messages=anthropic_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
+    try:
+        # Call Claude API
+        response = anthropic_client.messages.create(
+            model="claude-3-7-sonnet-20250219",
+            system=system_message,
+            messages=formatted_messages,
+            max_tokens=2000
+        )
+        
+        # Extract the response
+        agent_response = response.content[0].text
+        
+        # Create a new state instead of modifying the original
+        new_state = state.copy()
+        
+        # Add agent's response to messages
+        new_message = {"role": "assistant", "content": agent_response.replace("EXECUTION_COMPLETE", "").strip()}
+        new_state["messages"] = messages + [new_message]
+        
+        # Update execution results
+        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"
+        
+        return new_state
+    
+    except Exception as e:
+        # Handle any errors
+        print(f"Error in executor_agent: {str(e)}")
+        # Return a valid state in case of error
+        error_state = state.copy()
+        error_state["messages"] = messages + [{"role": "assistant", "content": f"I encountered an error: {str(e)}"}]
+        error_state["current_agent"] = "executor_agent"
+        error_state["status"] = "error"
+        return error_state
 
 def _execute_queries(db, sql_queries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     """