app.py

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 BaseMessage
import operator
import networkx as nx
import matplotlib.pyplot as plt

# Set API keys and validate credentials
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")

if not OPENAI_API_KEY or not TAVILY_API_KEY:
    st.error("API keys not found. Please set OPENAI_API_KEY and TAVILY_API_KEY as environment variables.")
    st.stop()

# Initialize OpenAI LLM
model = ChatOpenAI(temperature=0)

# Define Tools
@tool
def multiply(first_number: int, second_number: int) -> int:
    """Multiplies two integers together."""
    return first_number * second_number

@tool
def search(query: str):
    """Performs web search on the user query."""
    tavily = TavilySearchResults(max_results=1)
    result = tavily.invoke(query)
    return result

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[BaseMessage], operator.add]

# Define workflow nodes
def invoke_model(state):
    messages = state['messages']
    question = messages[-1]
    return {"messages": [model_with_tools.invoke(question)]}

def invoke_tool(state):
    tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
    tool_details = None

    for tool_call in tool_calls:
        tool_details = tool_call

    if tool_details is None:
        raise Exception("No tool input found.")

    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.")

    response = tool_map[selected_tool].invoke(json.loads(tool_details.get("function").get("arguments")))
    return {"messages": [response]}

def router(state):
    tool_calls = state['messages'][-1].additional_kwargs.get("tool_calls", [])
    if len(tool_calls):
        return "tool"
    else:
        return "end"

# Graph setup
graph = StateGraph(AgentState)
graph.add_node("agent", invoke_model)
graph.add_node("tool", invoke_tool)
graph.add_conditional_edges("agent", router, {"tool": "tool", "end": END})
graph.add_edge("tool", END)
graph.set_entry_point("agent")
compiled_app = graph.compile()

# Function to render graph with NetworkX
def render_graph_nx():
    G = nx.DiGraph()
    G.add_edge("agent", "tool", label="invoke tool")
    G.add_edge("agent", "end", label="end condition")
    G.add_edge("tool", "end", label="finish")

    pos = nx.spring_layout(G, seed=42)
    plt.figure(figsize=(8, 6))
    nx.draw(G, pos, with_labels=True, node_color="lightblue", node_size=3000, font_size=10, font_weight="bold")
    edge_labels = nx.get_edge_attributes(G, "label")
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=9)
    plt.title("Workflow Graph")
    st.pyplot(plt)

# Streamlit UI
st.title("Blah Blah Demo")
st.write("Compare results **with and without human intervention** in the workflow.")

# Sidebar for configuration
st.sidebar.header("Configuration")
st.session_state['human_loop'] = st.sidebar.checkbox("Enable Human-in-the-Loop", value=False)

# 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:", 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")

    # 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: {e}")
    st.write(f"Execution Time: {time.time() - start_time:.2f} seconds")

    # Display Workflow Graph
    st.subheader("Workflow Graph")
    render_graph_nx()