Update agents/executor.py

agents/executor.py

@@ -5,7 +5,7 @@ This agent executes SQL queries and reports on the results.
 
 import time
 import json
-from typing import Dict, Any, List, Union
+from typing import Dict, Any, List
 from agents.state import AgentState
 from anthropic import Anthropic
 from anthropic.types import MessageParam
@@ -22,15 +22,11 @@ def executor_agent(anthropic_client: Anthropic, db, state: AgentState) -> AgentS
     Returns:
         Updated state
     """
-    # Ensure messages exist and current_agent is a list or string
+    # Ensure messages exist and current_agent is a list
     messages = state.get("messages", [])
-    current_agent = state.get("current_agent", "executor_agent")
+    current_agent = state.get("current_agent", ["executor_agent"])
     sql_queries = state.get("sql_queries", [])
     
-    # Ensure current_agent is a list
-    if isinstance(current_agent, str):
-        current_agent = [current_agent]
-    
     # Add agent-specific instructions
     system_message = """
     You are an AI assistant specializing in executing and validating pharmaceutical data pipelines.
@@ -111,8 +107,7 @@ def executor_agent(anthropic_client: Anthropic, db, state: AgentState) -> AgentS
         # Update status
         new_state["status"] = "complete"
         # Append "complete" to current_agent
-        current_agent.append("complete")
-        new_state["current_agent"] = current_agent
+        new_state["current_agent"].append("complete")
         
         return new_state