Update app.py

app.py

@@ -46,16 +46,29 @@ graph.add_node("tool", invoke_tool)
 graph.add_edge("tool", END)
 graph.set_entry_point("agent")
 
+# Define the router function
 def router(state):
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    if len(tool_calls):  # Check if tool_calls exist
-        return "multiply"  # Route to the tool node
+    if len(tool_calls):
+        return "tools"
     else:
-        return "end"  # Route to the end node
-
-
-graph.add_conditional_edges("agent", router, {"multiply": "tool", "end": END})
-app_graph = graph.compile()
+        return "end"
+
+# Add graph nodes and edges
+app_graph = StateGraph(MessagesState)
+app_graph.add_node("assistant", assistant)
+app_graph.add_node("tools", ToolNode(tools))
+
+# Add conditional routing logic
+app_graph.add_edge(START, "assistant")
+app_graph.add_conditional_edges("assistant", router, {
+    "tools": "tools",
+    "end": END
+})
+app_graph.add_edge("tools", "assistant")
+
+# Compile the graph
+react_graph = app_graph.compile()
 
 # Save graph visualization as an image
 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile: