Update app.py

app.py

@@ -33,55 +33,57 @@ graph = StateGraph(AgentState)
 
 # Model Invocation
 def invoke_model(state):
-    question = state['messages'][-1].content
+    """
+    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]
     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["function"]["name"] == "multiply":
-            arguments = json.loads(tool_call["function"]["arguments"])
+        if tool_call.get("function", {}).get("name") == "multiply":
+            arguments = json.loads(tool_call.get("function").get("arguments"))
             result = multiply.invoke(arguments)
-            tool_results.append(
-                AIMessage(
-                    content=f"The result of multiplying {arguments['first_number']} and {arguments['second_number']} is {result}.",
-                    additional_kwargs={"tool_call_id": tool_call["id"]}
-                )
-            )
-    return {"messages": tool_results or [AIMessage(content="No valid tool input provided.")]}
+            # 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.")]}
 
 
 graph.add_node("tool", invoke_tool)
+graph.add_edge("tool", END)
+graph.set_entry_point("agent")
 
-# Explicit Router Node
+# Router Logic
 def router(state):
     """
-    Route to tool if tool calls exist; otherwise END the workflow.
+    Decide whether to invoke a tool or return the response.
     """
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    if tool_calls:
-        return {"messages": [AIMessage(content="Routing to tool...")]}
-    else:
-        return {"messages": [AIMessage(content=state['messages'][-1].content)]}
+    return "tool" if tool_calls else END
 
-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)
-
-# Compile the graph
-graph.set_entry_point("agent")
+graph.add_conditional_edges("agent", router, {"tool": "tool", END: END})
 app_graph = graph.compile()
 
-# Save graph visualization with xray for visibility
+# Save graph visualization as an image
 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
     graph_viz = app_graph.get_graph(xray=True)
     tmpfile.write(graph_viz.draw_mermaid_png())
@@ -89,14 +91,21 @@ 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")
-    first_number = st.number_input("First Number", value=0, step=1)
-    second_number = st.number_input("Second Number", value=0, step=1)
+    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)
 
     if st.button("Multiply"):
         question = f"What is {first_number} * {second_number}?"
@@ -106,16 +115,22 @@ 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"):
-        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}")
+        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.")
 
 # Sidebar for References
 st.sidebar.title("References")