Update agents/sql_generator.py

agents/sql_generator.py

@@ -4,7 +4,7 @@ This agent converts a data pipeline plan into executable SQL queries.
 """
 
 import time
-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
@@ -20,15 +20,11 @@ def sql_generator_agent(anthropic_client: Anthropic, state: AgentState) -> Agent
     Returns:
         Updated state
     """
-    # Ensure messages exist and current_agent is a list or string
+    # Get current messages and pipeline plan
     messages = state.get("messages", [])
-    current_agent = state.get("current_agent", "sql_generator_agent")
+    current_agent = state.get("current_agent", ["sql_generator_agent"])
     pipeline_plan = state.get("pipeline_plan", {})
     
-    # 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 generating SQL for pharmaceutical data pipelines.
@@ -99,13 +95,11 @@ def sql_generator_agent(anthropic_client: Anthropic, state: AgentState) -> Agent
             sql_queries = _extract_sql_queries(agent_response)
             
             new_state["sql_queries"] = sql_queries
-            # Append next agent, ensuring it's represented as a string
-            current_agent.append("executor_agent")
-            new_state["current_agent"] = current_agent
+            # Append next agent
+            new_state["current_agent"].append("executor_agent")
         else:
             # Append current agent if SQL generation is not complete
-            current_agent.append("sql_generator_agent")
-            new_state["current_agent"] = current_agent
+            new_state["current_agent"].append("sql_generator_agent")
         
         return new_state