Update graph/workflow.py

graph/workflow.py

@@ -20,6 +20,7 @@ from agents.tools import (
     tool_get_confidence
 )
 
+# Fix for the workflow.py file
 def create_agent_graph(anthropic_client, db):
     """
     Create the agent workflow graph.
@@ -137,7 +138,8 @@ def create_agent_graph(anthropic_client, db):
         lambda x: x["current_agent"],
         {
             "understanding_agent": "understanding_agent",
-            "planning_agent": "planning_agent"
+            "planning_agent": "planning_agent",
+            "tools": "tools"  # Add tools as a possible destination
         }
     )
     
@@ -146,7 +148,8 @@ def create_agent_graph(anthropic_client, db):
         lambda x: x["current_agent"],
         {
             "planning_agent": "planning_agent",
-            "sql_generator_agent": "sql_generator_agent"
+            "sql_generator_agent": "sql_generator_agent",
+            "tools": "tools"  # Add tools as a possible destination
         }
     )
     
@@ -155,17 +158,27 @@ def create_agent_graph(anthropic_client, db):
         lambda x: x["current_agent"],
         {
             "sql_generator_agent": "sql_generator_agent",
-            "executor_agent": "executor_agent"
+            "executor_agent": "executor_agent",
+            "tools": "tools"  # Add tools as a possible destination
         }
     )
     
     # Executor agent finishes the workflow
     workflow.add_edge("executor_agent", END)
     
-    # Add edges for tools - they can be called from any agent and return to that agent
-    for agent in ["understanding_agent", "planning_agent", "sql_generator_agent", "executor_agent"]:
-        workflow.add_edge(agent, "tools")
-        workflow.add_edge("tools", agent)
+    # Add edge from tools back to the calling agent
+    # This is a simplified approach - in a production system, you would
+    # need to track which agent called the tool and return to that agent
+    workflow.add_conditional_edges(
+        "tools",
+        lambda x: x["current_agent"],
+        {
+            "understanding_agent": "understanding_agent",
+            "planning_agent": "planning_agent",
+            "sql_generator_agent": "sql_generator_agent",
+            "executor_agent": "executor_agent"
+        }
+    )
     
     # Compile the workflow
     app = workflow.compile()