Update app.py

app.py

@@ -1,13 +1,12 @@
 import os
 import streamlit as st
 import json
-import time
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import tool
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langgraph.graph import StateGraph, END
 from typing import TypedDict, Annotated, Sequence
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import BaseMessage
 import operator
 import networkx as nx
 import matplotlib.pyplot as plt
@@ -38,17 +37,18 @@ def search(query: str):
 
 tools = [search, multiply]
 tool_map = {tool.name: tool for tool in tools}
+
 model_with_tools = model.bind_tools(tools)
 
 # Define Agent State class
 class AgentState(TypedDict):
-    messages: Annotated[Sequence[HumanMessage], operator.add]
+    messages: Annotated[Sequence[BaseMessage], operator.add]
 
 # Define workflow nodes
 def invoke_model(state):
     messages = state['messages']
     question = messages[-1]
-    return {"messages": [model_with_tools.invoke(question.content)]}
+    return {"messages": [model_with_tools.invoke(question)]}
 
 def invoke_tool(state):
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
@@ -63,15 +63,14 @@ def invoke_tool(state):
     selected_tool = tool_details.get("function").get("name")
     st.sidebar.write(f"Selected tool: {selected_tool}")
 
-    # Add human-in-the-loop for all tools
-    if st.session_state['human_loop']:
-        response = st.sidebar.radio(f"Proceed with tool '{selected_tool}'?", ["Yes", "No"], index=0)
-        if response == "No":
-            raise ValueError(f"Execution of '{selected_tool}' was canceled.")
-
-    # Invoke tool and return response as AIMessage
+    if selected_tool == "search":
+        if 'human_loop' in st.session_state and st.session_state['human_loop']:
+            response = st.sidebar.radio("Proceed with web search?", ["Yes", "No"])
+            if response == "No":
+                raise ValueError("User canceled the search tool execution.")
+    
     response = tool_map[selected_tool].invoke(json.loads(tool_details.get("function").get("arguments")))
-    return {"messages": [AIMessage(content=str(response))]}
+    return {"messages": [response]}
 
 def router(state):
     tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
@@ -90,7 +89,7 @@ graph.set_entry_point("agent")
 compiled_app = graph.compile()
 
 # Function to render graph with NetworkX
-def render_graph_nx():
+def render_graph_nx(graph):
     G = nx.DiGraph()
     G.add_edge("agent", "tool", label="invoke tool")
     G.add_edge("agent", "end", label="end condition")
@@ -106,43 +105,29 @@ def render_graph_nx():
 
 # Streamlit UI
 st.title("LLM Tool Workflow Demo")
-st.write("Compare results **with and without human intervention** in the workflow.")
+st.write("This app demonstrates LLM-based tool usage with and without human intervention.")
 
-# Sidebar for configuration
+# Sidebar for options
 st.sidebar.header("Configuration")
-st.session_state['human_loop'] = st.sidebar.checkbox("Enable Human-in-the-Loop", value=False)
+st.session_state['human_loop'] = st.sidebar.checkbox("Enable Human-in-the-Loop (For Search)", value=False)
 
 # Input prompt
 prompt = st.text_input("Enter your question:", "What is 24 * 365?")
 if st.button("Run Workflow"):
     st.subheader("Execution Results")
-
-    # Run workflow without and with human-in-the-loop
-    for human_loop in [False, True]:
-        st.markdown(f"### {'With' if human_loop else 'Without'} Human-in-the-Loop")
-        st.session_state['human_loop'] = human_loop
+    try:
         intermediate_outputs = []
-        start_time = time.time()
-
-        try:
-            # Collect intermediate outputs
-            for s in compiled_app.stream({"messages": [HumanMessage(content=prompt)]}):
-                intermediate_outputs.append(s)
-
-            # Extract and display the final response
-            final_output = intermediate_outputs[-1]
-            messages = final_output.get('messages', [])
-            if messages:
-                response_content = messages[-1].content
-                st.write("Response:", response_content)
-            else:
-                st.write("Response: No content generated.")
-        except Exception as e:
-            st.error(f"Error: {e}")
-
-        st.write(f"Execution Time: {time.time() - start_time:.2f} seconds")
-        st.write("---")
-
-    # Display Workflow Graph
+        for s in compiled_app.stream({"messages": [prompt]}):
+            intermediate_outputs.append(s)
+            st.write("Response:", list(s.values())[0])
+            st.write("---")
+
+        st.sidebar.write("### Intermediate Outputs")
+        for i, output in enumerate(intermediate_outputs):
+            st.sidebar.write(f"Step {i+1}: {output}")
+    except Exception as e:
+        st.error(f"Error occurred: {e}")
+
+    # Display Graph
     st.subheader("Workflow Graph")
-    render_graph_nx()
+    render_graph_nx(graph)