Update graph/workflow.py

graph/workflow.py

@@ -2,10 +2,10 @@
 LangGraph workflow for pharmaceutical data management agents.
 """
 
-from langgraph.graph import StateGraph, END
+from langgraph.graph import StateGraph, END, START
 from langchain.tools import StructuredTool
 from langchain_core.tools import tool
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Callable
 
 from agents.state import AgentState
 from agents.understanding import understanding_agent
@@ -19,8 +19,8 @@ from agents.tools import (
     tool_execute_query,
     tool_get_confidence
 )
+from agents.utils.logging import log_agent_activity
 
-# Fix for the workflow.py file
 def create_agent_graph(anthropic_client, db):
     """
     Create the agent workflow graph.
@@ -30,7 +30,7 @@ def create_agent_graph(anthropic_client, db):
         db: The database connection
     
     Returns:
-        Compiled LangGraph workflow
+        Compiled LangGraph workflow and state update function
     """
     # Wrap the agents with the anthropic client
     understanding = lambda state: understanding_agent(anthropic_client, state)
@@ -47,8 +47,8 @@ def create_agent_graph(anthropic_client, db):
     }
     
     # Create a function to access the current state for the confidence tool
-    state_provider = lambda: state_dict
     state_dict = {}  # This will be updated in the Streamlit app
+    state_provider = lambda: state_dict
     
     # Create tools node with database-related tools
     # Convert our custom tools to LangChain StructuredTool format
@@ -110,16 +110,14 @@ def create_agent_graph(anthropic_client, db):
     lc_tools.append(get_confidence)
     
     # Create Tool Agent Node
-    # In the updated LangGraph, we don't use ToolNode directly
-    # We'll create a tool handler function instead
+    # We'll create a tool handler function
     def tool_handler(state: Dict[str, Any]) -> Dict[str, Any]:
         """Handle tool calls from the agent workflow."""
+        log_agent_activity("TOOL", state)
         # This would normally implement logic to determine which tool to call
-        # based on the agent's request, but for this demo we'll use a simpler approach
         # In production, you would parse agent messages to identify tool calls
         
-        # Just return the state unmodified - tools are actually called via the LangChain integration
-        # in each agent's implementation
+        # Just return the state unmodified - tools are actually called via the agents
         return state
     
     nodes["tools"] = tool_handler
@@ -131,15 +129,16 @@ def create_agent_graph(anthropic_client, db):
     for name, node in nodes.items():
         workflow.add_node(name, node)
     
-    # Define the edges (transitions between agents)
-    # First, create conditional routing based on the current_agent field
+    # Set the entry point to understanding_agent (FIX APPLIED HERE)
+    workflow.add_edge(START, "understanding_agent")
+    
+    # Define conditional routing based on the current_agent field
     workflow.add_conditional_edges(
         "understanding_agent",
         lambda x: x["current_agent"],
         {
             "understanding_agent": "understanding_agent",
-            "planning_agent": "planning_agent",
-            "tools": "tools"  # Add tools as a possible destination
+            "planning_agent": "planning_agent"
         }
     )
     
@@ -148,8 +147,7 @@ def create_agent_graph(anthropic_client, db):
         lambda x: x["current_agent"],
         {
             "planning_agent": "planning_agent",
-            "sql_generator_agent": "sql_generator_agent",
-            "tools": "tools"  # Add tools as a possible destination
+            "sql_generator_agent": "sql_generator_agent"
         }
     )
     
@@ -158,27 +156,17 @@ def create_agent_graph(anthropic_client, db):
         lambda x: x["current_agent"],
         {
             "sql_generator_agent": "sql_generator_agent",
-            "executor_agent": "executor_agent",
-            "tools": "tools"  # Add tools as a possible destination
+            "executor_agent": "executor_agent"
         }
     )
     
     # Executor agent finishes the workflow
     workflow.add_edge("executor_agent", END)
     
-    # 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"
-        }
-    )
+    # 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)
     
     # Compile the workflow
     app = workflow.compile()