Update graph/workflow.py

graph/workflow.py

@@ -134,21 +134,48 @@ def create_agent_graph(anthropic_client, db):
     # Set the entry point to understanding_agent 
     workflow.add_edge(START, "understanding_agent")
     
-    # Modify conditional routing to handle the new Annotated current_agent type
-    def route_understanding(state: AgentState) -> Union[str, None]:
-        """Route from understanding agent based on the last added current_agent value."""
-        current_agents = state['current_agent']
-        return current_agents[-1] if current_agents else "understanding_agent"
-    
-    def route_planning(state: AgentState) -> Union[str, None]:
-        """Route from planning agent based on the last added current_agent value."""
-        current_agents = state['current_agent']
-        return current_agents[-1] if current_agents else "planning_agent"
-    
-    def route_sql_generator(state: AgentState) -> Union[str, None]:
-        """Route from SQL generator agent based on the last added current_agent value."""
-        current_agents = state['current_agent']
-        return current_agents[-1] if current_agents else "sql_generator_agent"
+    # Modify routing functions to handle str or list current_agent
+    def route_understanding(state: AgentState):
+        """Route from understanding agent."""
+        current_agent = state['current_agent']
+        # If it's a list, take the last element
+        if isinstance(current_agent, list):
+            current_agent = current_agent[-1]
+        
+        # Routing logic
+        routing_map = {
+            "understanding_agent": "understanding_agent",
+            "planning_agent": "planning_agent"
+        }
+        return routing_map.get(current_agent, "understanding_agent")
+    
+    def route_planning(state: AgentState):
+        """Route from planning agent."""
+        current_agent = state['current_agent']
+        # If it's a list, take the last element
+        if isinstance(current_agent, list):
+            current_agent = current_agent[-1]
+        
+        # Routing logic
+        routing_map = {
+            "planning_agent": "planning_agent",
+            "sql_generator_agent": "sql_generator_agent"
+        }
+        return routing_map.get(current_agent, "planning_agent")
+    
+    def route_sql_generator(state: AgentState):
+        """Route from SQL generator agent."""
+        current_agent = state['current_agent']
+        # If it's a list, take the last element
+        if isinstance(current_agent, list):
+            current_agent = current_agent[-1]
+        
+        # Routing logic
+        routing_map = {
+            "sql_generator_agent": "sql_generator_agent",
+            "executor_agent": "executor_agent"
+        }
+        return routing_map.get(current_agent, "sql_generator_agent")
     
     workflow.add_conditional_edges(
         "understanding_agent",