Update app.py

app.py

@@ -1,6 +1,7 @@
 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
@@ -37,7 +38,6 @@ 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
@@ -63,12 +63,12 @@ def invoke_tool(state):
     selected_tool = tool_details.get("function").get("name")
     st.sidebar.write(f"Selected tool: {selected_tool}")
 
-    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.")
-    
+    # 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.")
+
     response = tool_map[selected_tool].invoke(json.loads(tool_details.get("function").get("arguments")))
     return {"messages": [response]}
 
@@ -89,7 +89,7 @@ graph.set_entry_point("agent")
 compiled_app = graph.compile()
 
 # Function to render graph with NetworkX
-def render_graph_nx(graph):
+def render_graph_nx():
     G = nx.DiGraph()
     G.add_edge("agent", "tool", label="invoke tool")
     G.add_edge("agent", "end", label="end condition")
@@ -104,30 +104,44 @@ def render_graph_nx(graph):
     st.pyplot(plt)
 
 # Streamlit UI
-st.title("LLM Tool Workflow Demo")
-st.write("This app demonstrates LLM-based tool usage with and without human intervention.")
+st.title("Blah Blah Demo")
+st.write("Compare results **with and without human intervention** in the workflow.")
 
-# Sidebar for options
+# Sidebar for configuration
 st.sidebar.header("Configuration")
-st.session_state['human_loop'] = st.sidebar.checkbox("Enable Human-in-the-Loop (For Search)", value=False)
+st.session_state['human_loop'] = st.sidebar.checkbox("Enable Human-in-the-Loop", value=False)
 
-# Input prompt
+# Input and comparison mode
 prompt = st.text_input("Enter your question:", "What is 24 * 365?")
 if st.button("Run Workflow"):
     st.subheader("Execution Results")
+
+    # Without human-in-the-loop
+    st.markdown("### Without Human-in-the-Loop")
+    st.session_state['human_loop'] = False
+    start_time = time.time()
     try:
         intermediate_outputs = []
         for s in compiled_app.stream({"messages": [prompt]}):
             intermediate_outputs.append(s)
-            st.write("Response:", list(s.values())[0])
-            st.write("---")
+        st.write("Response:", intermediate_outputs[-1]['messages'][0])
+    except Exception as e:
+        st.error(f"Error: {e}")
+    st.write(f"Execution Time: {time.time() - start_time:.2f} seconds")
 
-        st.sidebar.write("### Intermediate Outputs")
-        for i, output in enumerate(intermediate_outputs):
-            st.sidebar.write(f"Step {i+1}: {output}")
+    # With human-in-the-loop
+    st.markdown("### With Human-in-the-Loop")
+    st.session_state['human_loop'] = True
+    start_time = time.time()
+    try:
+        intermediate_outputs = []
+        for s in compiled_app.stream({"messages": [prompt]}):
+            intermediate_outputs.append(s)
+        st.write("Response:", intermediate_outputs[-1]['messages'][0])
     except Exception as e:
-        st.error(f"Error occurred: {e}")
+        st.error(f"Error: {e}")
+    st.write(f"Execution Time: {time.time() - start_time:.2f} seconds")
 
-    # Display Graph
+    # Display Workflow Graph
     st.subheader("Workflow Graph")
-    render_graph_nx(graph)
+    render_graph_nx()