Update app.py

app.py

@@ -6,7 +6,7 @@ import tempfile
 from typing import TypedDict, Annotated, Sequence
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import tool
-from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.messages import HumanMessage, ToolMessage
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from langgraph.graph import StateGraph, END
 
@@ -36,19 +36,10 @@ 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 the model's response with tool_calls, if any
     return {"messages": [response]}
 
 graph.add_node("agent", invoke_model)
@@ -60,30 +51,33 @@ def invoke_tool(state):
     """
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
     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)
-            return {"messages": [AIMessage(content=f"Tool Result: {result}")]}
-    return {"messages": [AIMessage(content="No valid tool input provided.")]}
+            return {"messages": [ToolMessage(content=str(result))]}
+    return {"messages": [ToolMessage(content="No valid tool input provided.")]}
 
 graph.add_node("tool", invoke_tool)
-graph.add_edge("tool", END)
-graph.set_entry_point("agent")
 
-# Router Logic
+# Router Node: Manual Addition
 def router(state):
-    """
-    Decide whether to invoke a tool or return the response.
-    """
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
-    return "tool" if tool_calls else END
+    return {"messages": [ToolMessage(content="Routing...")]}  # Dummy message for router node
 
-graph.add_conditional_edges("agent", router, {"tool": "tool", END: END})
+graph.add_node("router", router)
+
+# Add explicit edges from agent to router, and router to tool/END
+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")
 app_graph = graph.compile()
 
-# Save graph visualization as an image
+# Save graph visualization with router explicitly included
 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
-    graph_viz = app_graph.get_graph(xray=True)
+    graph_viz = app_graph.get_graph(xray=True)  # Ensures detailed visualization
     tmpfile.write(graph_viz.draw_mermaid_png())
     graph_image_path = tmpfile.name
 
@@ -93,6 +87,7 @@ st.title("Simple Tool Calling Demo")
 # Display the workflow graph
 st.image(graph_image_path, caption="Workflow Visualization")
 
+# Tabbed Interface
 tab1, tab2 = st.tabs(["Try Multiplication", "Ask General Queries"])
 
 # Multiplication Tool Tab
@@ -121,7 +116,6 @@ with tab2:
     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)
@@ -133,3 +127,4 @@ with tab2:
 # Sidebar for References
 st.sidebar.title("References")
 st.sidebar.markdown("1. [LangGraph Tool Calling](https://github.com/aritrasen87/LLM_RAG_Model_Deployment/blob/main/LangGraph_02_ToolCalling.ipynb)")
+