Update app.py

app.py

@@ -6,7 +6,7 @@ from langgraph.graph.message import add_messages
 from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import HumanMessage, ToolMessage, AIMessage
-from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode, tools_condition
 import os
 
 # Streamlit UI Header
@@ -31,52 +31,38 @@ class State(TypedDict):
 # Initialize LLM and Tools
 llm = ChatOpenAI(model="gpt-4o-mini")
 tool = TavilySearchResults(max_results=2)
-llm_with_tools = llm.bind_tools([tool])
+tools = [tool]
+llm_with_tools = llm.bind_tools(tools)
 
 # Agent Node
 def Agent(state: State):
-    st.sidebar.write("Agent Input State:", state["messages"])
+    st.sidebar.write("Agent received input:", state["messages"])
     response = llm_with_tools.invoke(state["messages"])
     st.sidebar.write("Agent Response:", response)
     return {"messages": [response]}
 
-# Tools Execution Node
-def ExecuteTools(state: State):
-    tool_calls = state["messages"][-1].tool_calls
-    responses = []
-
-    if tool_calls:
-        for call in tool_calls:
-            tool_name = call["name"]
-            args = call["args"]
-            st.sidebar.write("Tool Call Detected:", tool_name, args)
-
-            if tool_name == "tavily_search_results_json":
-                tool_response = tool.invoke({"query": args["query"]})
-                st.sidebar.write("Tool Response:", tool_response)
-                responses.append(ToolMessage(content=str(tool_response), tool_call_id=call["id"]))
-    return {"messages": responses}
-
-# Memory Checkpoint
+# Set up Graph
 memory = MemorySaver()
-
-# Build the Graph
 graph = StateGraph(State)
+
+# Add nodes
 graph.add_node("Agent", Agent)
-graph.add_node("ExecuteTools", ExecuteTools)
+tool_node = ToolNode(tools=[tool])
+graph.add_node("tools", tool_node)
 
-graph.add_conditional_edges("Agent", tools_condition, {"True": "ExecuteTools", "False": "Agent"})
-graph.add_edge("ExecuteTools", "Agent")
+# Add edges
+graph.add_conditional_edges("Agent", tools_condition)
+graph.add_edge("tools", "Agent")
 graph.set_entry_point("Agent")
 
-# Compile the Graph
-app = graph.compile(checkpointer=memory, interrupt_before=["ExecuteTools"])
+# Compile with Breakpoint
+app = graph.compile(checkpointer=memory, interrupt_before=["tools"])
 
 # Display Graph Visualization
 st.subheader("Graph Visualization")
 st.image(app.get_graph().draw_mermaid_png(), caption="Workflow Graph", use_container_width=True)
 
-# Run the Workflow
+# Input Section
 st.subheader("Run the Workflow")
 user_input = st.text_input("Enter a message to start the graph:", "Search for the weather in Uttar Pradesh")
 thread_id = st.text_input("Thread ID", "1")
@@ -88,13 +74,11 @@ if st.button("Execute Workflow"):
     st.write("### Execution Outputs")
     outputs = []
 
+    # Execute the workflow
     try:
-        # Stream the graph execution
         for event in app.stream(input_message, thread, stream_mode="values"):
-            output_message = event["messages"][-1]
-            st.code(output_message.content)
-            outputs.append(output_message.content)
-            st.sidebar.write("Intermediate State:", event["messages"])
+            st.code(event["messages"][-1].content)
+            outputs.append(event["messages"][-1].content)
 
         # Display Intermediate Outputs
         if outputs:
@@ -103,27 +87,28 @@ if st.button("Execute Workflow"):
                 st.write(f"**Step {idx}:**")
                 st.code(output)
         else:
-            st.warning("No outputs generated. Check the workflow or tool calls.")
+            st.warning("No outputs generated yet.")
 
-        # Snapshot of Current State
+        # Show State Snapshot
         st.subheader("Current State Snapshot")
         snapshot = app.get_state(thread)
         current_message = snapshot.values["messages"][-1]
         st.code(current_message.pretty_print())
 
-        # Manual Update for Interrupted State
+        # Handle Tool Calls with Manual Input
         if hasattr(current_message, "tool_calls") and current_message.tool_calls:
             tool_call_id = current_message.tool_calls[0]["id"]
-            manual_response = st.text_area("Manual Tool Response", "Enter response to continue execution...")
-            if st.button("Update State"):
+            st.warning("Execution paused before tool execution. Provide manual input to resume.")
+            manual_response = st.text_area("Manual Tool Response", "Enter the tool's response here...")
+            if st.button("Resume Execution"):
                 new_messages = [
                     ToolMessage(content=manual_response, tool_call_id=tool_call_id),
                     AIMessage(content=manual_response),
                 ]
                 app.update_state(thread, {"messages": new_messages})
-                st.success("State updated successfully!")
+                st.success("State updated! Rerun the workflow to continue.")
                 st.code(app.get_state(thread).values["messages"][-1].pretty_print())
         else:
-            st.warning("No tool calls detected to update the state.")
+            st.info("No tool calls detected at this step.")
     except Exception as e:
         st.error(f"Error during execution: {e}")