Update ui/agent_workflow.py

ui/agent_workflow.py

@@ -1,4 +1,3 @@
-# agent_workflow.py
 """
 Agent workflow UI component for the pharmaceutical data management agent.
 """
@@ -14,17 +13,8 @@ def render_workflow_tab(session_state):
     """
     st.subheader("Agent Workflow Visualization")
     
-    # Display current agent and status
-    current_agent_raw = session_state.conversation.get("current_agent", "understanding_agent")
-    
-    # Handle both string and list types for current_agent
-    if isinstance(current_agent_raw, list):
-        # Get the last agent in the list if it's not empty
-        current_agent = current_agent_raw[-1] if current_agent_raw else "understanding_agent"
-    else:
-        # If it's already a string, use it directly
-        current_agent = current_agent_raw
-        
+    # Display current agent and status - ensure it's consistent
+    current_agent = session_state.conversation.get("current_agent", "understanding_agent")
     status = session_state.conversation.get("status", "planning")
     
     st.markdown(f"**Current State:** {status.title()}")