Update app.py

app.py

@@ -82,21 +82,14 @@ def research_node(state):
     last_message = result["messages"][-1]
     content = last_message.content if hasattr(last_message, "content") else last_message
     
-    # Determine next step based on content
-    if "FINAL ANSWER" in content:
-        # If there's a final answer, we can end the workflow
-        goto = "end"
-    else:
-        # Otherwise, continue to the chart generator
-        goto = "chart_generator"
-    
     # Create an AIMessage with the researcher name
     if not isinstance(last_message, dict):
         result["messages"][-1] = AIMessage(content=content, name="researcher")
     else:
         result["messages"][-1]["name"] = "researcher"
     
-    return Command(update={"messages": result["messages"]}, goto=goto)
+    # In langgraph 0.0.41 style, return the updated state
+    return {"messages": result["messages"], "has_final_answer": "FINAL ANSWER" in content}
 
 # Chart generation phase
 def chart_node(state):
@@ -120,22 +113,25 @@ def chart_node(state):
     else:
         result["messages"][-1]["name"] = "chart_generator"
     
-    # Always go to end after chart generation
-    return Command(update={"messages": result["messages"]}, goto="end")
+    # Return the updated state
+    return {"messages": result["messages"], "has_final_answer": True}
+
+# Define the router function for the graph
+def router(state):
+    # Check if we have a final answer
+    return "chart_generator" if not state.get("has_final_answer", False) else "end"
 
-# End node that simply returns the state
+# Simplified end node that doesn't modify the state
 def end_node(state):
     return state
 
-# Build LangGraph - FIXED with proper node naming
-workflow = StateGraph(dict)  # Using dict for state in langgraph 0.0.41
+# Build a proper LangGraph for v0.0.41
+workflow = StateGraph({"messages": [], "has_final_answer": False})
 workflow.add_node("researcher", research_node)
 workflow.add_node("chart_generator", chart_node)
 workflow.add_node("end", end_node)
 workflow.set_entry_point("researcher")
-# Use a defined node name "end" instead of "__end__"
-workflow.add_edge("researcher", "chart_generator")
-workflow.add_edge("researcher", "end")
+workflow.add_conditional_edges("researcher", router)
 workflow.add_edge("chart_generator", "end")
 graph = workflow.compile()
 
@@ -145,12 +141,18 @@ def run_langgraph(user_input):
         # Create a human message
         human_message = HumanMessage(content=user_input)
         
-        # Stream the events
-        events = graph.stream({"messages": [human_message]})
+        # Run the workflow with proper initialization
+        events = graph.stream({"messages": [human_message], "has_final_answer": False})
         outputs = list(events)
         
         # Get the final message
-        final_message = outputs[-1]["messages"][-1]
+        final_state = outputs[-1]
+        final_messages = final_state.get("messages", [])
+        
+        if not final_messages:
+            return "Error: No messages in the final state.", None
+            
+        final_message = final_messages[-1]
         final_content = final_message.content if hasattr(final_message, "content") else final_message
         
         if isinstance(final_content, str) and "FINAL ANSWER" in final_content: