Update app.py

app.py

@@ -33,57 +33,55 @@ graph = StateGraph(AgentState)
 
 # Model Invocation
 def invoke_model(state):
-    """
-    Invoke the model and handle tool invocation logic.
-    """
-    # Extract the question as a string
-    question = state['messages'][-1].content if isinstance(state['messages'][-1], HumanMessage) else state['messages'][-1]
+    question = state['messages'][-1].content
     response = model_with_tools.invoke(question)
-
-    # If the response is plain text (no tool calls)
-    if isinstance(response, str):
-        return {"messages": [AIMessage(content=response)]}
-
-    # If no tool calls exist
-    if not response.additional_kwargs.get("tool_calls", []):
-        return {"messages": [AIMessage(content=response.content)]}
-
-    # If tool calls are present, return the full response
     return {"messages": [response]}
 
 graph.add_node("agent", invoke_model)
 
 # Tool Invocation
 def invoke_tool(state):
-    """
-    Invoke the 'multiply' tool if it's called by the model.
-    """
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
+    tool_results = []
+
     for tool_call in tool_calls:
-        if tool_call.get("function", {}).get("name") == "multiply":
-            arguments = json.loads(tool_call.get("function").get("arguments"))
+        if tool_call["function"]["name"] == "multiply":
+            arguments = json.loads(tool_call["function"]["arguments"])
             result = multiply.invoke(arguments)
-            # Append the result back to the state
-            return {"messages": state['messages'] + [AIMessage(content=f"Tool Result: {result}")]}
-    return {"messages": state['messages'] + [AIMessage(content="No valid tool input provided.")]}
-
+            tool_results.append(
+                AIMessage(content=f"Tool Result: {result}", additional_kwargs={"tool_call_id": tool_call["id"]})
+            )
+    return {"messages": tool_results or [AIMessage(content="No valid tool input provided.")]}
 
 graph.add_node("tool", invoke_tool)
-graph.add_edge("tool", END)
-graph.set_entry_point("agent")
 
-# Router Logic
+# Fixed Router Node
 def router(state):
     """
-    Decide whether to invoke a tool or return the response.
+    Route to tool if tool calls exist; otherwise END the workflow.
     """
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    return "tool" if tool_calls else END
+    if tool_calls:
+        return {"__state__": "tool"}  # Direct routing to tool
+    else:
+        return {"__state__": END}  # End workflow if no tool calls
+
+graph.add_node("router", router)
+
+# Add explicit edges
+graph.add_edge("agent", "router")
+graph.add_conditional_edges(
+    "router",
+    lambda state: "tool" if state['messages'][-1].additional_kwargs.get("tool_calls") else END,
+    {"tool": "tool", END: END}
+)
+graph.add_edge("tool", END)
 
-graph.add_conditional_edges("agent", router, {"tool": "tool", END: END})
+# Compile the graph
+graph.set_entry_point("agent")
 app_graph = graph.compile()
 
-# Save graph visualization as an image
+# Save graph visualization with xray for visibility
 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
     graph_viz = app_graph.get_graph(xray=True)
     tmpfile.write(graph_viz.draw_mermaid_png())
@@ -91,21 +89,14 @@ with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
 
 # Streamlit Interface
 st.title("Simple Tool Calling Demo")
-
-# Display the workflow graph
 st.image(graph_image_path, caption="Workflow Visualization")
 
 tab1, tab2 = st.tabs(["Try Multiplication", "Ask General Queries"])
 
-# Multiplication Tool Tab
 with tab1:
     st.subheader("Try Multiplication")
-    col1, col2 = st.columns(2)
-
-    with col1:
-        first_number = st.number_input("First Number", value=0, step=1)
-    with col2:
-        second_number = st.number_input("Second Number", value=0, step=1)
+    first_number = st.number_input("First Number", value=0, step=1)
+    second_number = st.number_input("Second Number", value=0, step=1)
 
     if st.button("Multiply"):
         question = f"What is {first_number} * {second_number}?"
@@ -115,22 +106,16 @@ with tab1:
         except Exception as e:
             st.error(f"Error: {e}")
 
-# General Query Tab
 with tab2:
     st.subheader("General Query")
     user_input = st.text_input("Enter your question here")
 
     if st.button("Submit"):
-        if user_input:
-            try:
-                # Pass the user input as a HumanMessage
-                result = app_graph.invoke({"messages": [HumanMessage(content=user_input)]})
-                st.write("Response:")
-                st.success(result['messages'][-1].content)
-            except Exception as e:
-                st.error(f"Error: {e}")
-        else:
-            st.warning("Please enter a valid input.")
+        try:
+            result = app_graph.invoke({"messages": [HumanMessage(content=user_input)]})
+            st.success(result['messages'][-1].content)
+        except Exception as e:
+            st.error(f"Error: {e}")
 
 # Sidebar for References
 st.sidebar.title("References")