Building a Streamlit UI for My LangGraph AI Agent

Editor
9 Min Read


In my previous article, I how I built a LangGraph-based AI agent to automate a 15-minute customer service booking session.

The agent handles the entire booking process like a real customer service representative. It’s a LangGraph-based agent that orchestrates the following operations:

  • Responds to customer queries and understands their needs.
  • Calculates the price for the service and informs the customer.
  • Handles the customer’s acceptance or rejection.
  • Proposes optimized time slots.
  • Confirms and records the appointment.

In the first version of the agent, I did not focus much on the UI/UX part. I just built a Python CLI to test the functionality of the agent.

The customer service agent ran entirely in the terminal. The CLI worked well for testing but it was not the best way to demonstrate a customer-facing booking experience.

The full source code of this project is available on GitHub at customer-service-agent. Feel free to clone the repo and test it yourself.

In this article, we will build a clean, interactive Streamlit UI on top of the existing LangGraph agent.

A quick note on the terminology: Throughout this article, I use agent and graph interchangeably. In LangGraph, the agent architecture is defined and executed as a compiled state graph object so they essentially mean the same thing in this article.

User interface for the agent

In terms of implementation, Streamlit is not very different from a Python CLI. Both serve as a wrapper for the LangGraph agent. Streamlit is, of course, much more user friendly and looks more appealing.

The CLI interface collected input, invoked the graph, and printed the response. The Streamlit page will do the same but it will also render structured information extracted from the graph state such as current booking details, price quote, and acceptance buttons.

The architecture still lies in the same application. Streamlit only presents the state to the user and sends user actions back to the agent.

Streamlit page

Since we’re using poetry for dependency management, we can install streamlit using:

poetry add streamlit

This updates both pyproject.toml and poetry.lock files. Then we create a new file streamlit_app.py .

We start by importing the graph builder, models, and observibility utilities.

from __future__ import annotations

import os
from datetime import datetime
from typing import Any
from uuid import uuid4

import streamlit as st
from dotenv import load_dotenv
from langchain_core.messages import AIMessage, HumanMessage
from langchain_openai import ChatOpenAI

from customer_service_agent.graph import build_graph
from customer_service_agent.models import (
    AgentState,
    BookingDetails,
    TimeOption,
)
from customer_service_agent.observability import (
    create_langfuse_handler,
    flush_langfuse,
    graph_config,
)

The agent graph does not include any Streamlit-specific logic. This separation is important because it allows us to run the graph from a CLI, an API, WhatsApp, or another frontend later.

The graph expects an Agent State to be initialized ( graph = StateGraph(AgentState) ) so we add the following in streamlit_app.py :

INITIAL_STATE: AgentState = {
    "messages": [],
    "booking_details": BookingDetails(),
    "calculated_price": None,
    "time_options": [],
    "selected_slot": None,
    "status": "gathering_info",
}

After the first turn (i.e. first customer message), LangGraph’s checkpointer retains the state.

Streamlit reruns the entire Python script whenever user interacts with a widget (e.g. sends a chat message, clicks a button, selects an appointment) so we cannot use local variables. To preserve the conversation across Streamlit runs, we need to use session_state .

We can initialize the session as follows:

def initialize_session() -> None:
    if "graph" in st.session_state:
        return

    llm = ChatOpenAI(
        model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
        temperature=0,
    )

    handler = create_langfuse_handler()

    st.session_state.graph = build_graph(llm)
    st.session_state.handler = handler
    st.session_state.config = graph_config(
        str(uuid4()),
        handler,
    )
    st.session_state.agent_state = INITIAL_STATE.copy()
    st.session_state.started = False

The function first checks if the graph has already been created for this browser session (if "graph" in st.session_state ). Without this check, every Streamlit rerun would replace the graph.

The graph_config function generates a UUID to be used as thread_id , which is required by LangGraph to identify a conversation. If a new UUID were generated on every Streamlit rerun, LangGraph would see every message as belonging to a new conversation.

We also have the Langfuse tracing integration inside this function. The handler is saved so that the subsequent graph calls can use the same tracing configuration.

Then, we have the _invoke function for handling user input.

def _invoke(customer_text: str) -> None:
    """Submit one customer turn to the graph and retain its latest state."""
    graph_input: dict[str, Any] = {"messages": [HumanMessage(content=customer_text)]}
    if not st.session_state.started:
        graph_input.update(INITIAL_STATE)
        graph_input["messages"] = [HumanMessage(content=customer_text)]
        st.session_state.started = True

    try:
        result = st.session_state.graph.invoke(graph_input, config=st.session_state.config)
        st.session_state.agent_state = result
        flush_langfuse(st.session_state.handler)
    except Exception:
        st.session_state.started = bool(st.session_state.agent_state.get("messages"))
        st.error("The assistant could not process that request. Please try again.")

This function sends the customer action to the LangGraph agent and saves the resulting state for the Streamlit interface. The customer action can be a chat input or a button click (e.g. “Accept quote”).

The customer’s message is converted into a LangChain HumanMessage. The messages uses LangGraph’s add_messages reducer so new messages are added to the existing conversation instead of replacing it.

For the first message, we initialize the graph using the INITIAL_STATE defined earlier with empty booking details, scheduling options, price, and the initial status.

Then, whenever we receive a new custom action, we invoke the graph and update its state with the result. This is where the LLM calls happen:

result = st.session_state.graph.invoke(graph_input, config=st.session_state.config)
st.session_state.agent_state = result

This runs the customer message through the graph. The returned graph state (i.e. results ) is stored in Streamlit’s session so that the page can render the latest details of messages, booking summary, price, appointment options, and status.

Finally, we have some render functions (_render...() ) defined in streamlit_app.py to convert the current LangGraph state into visible Streamlit components.

def _render_messages(state: AgentState) -> None:
    if not state.get("messages"):
        with st.chat_message("assistant"):
            st.write(
                "Hi! I can help you book house or couch cleaning. "
                "Tell me what you need, including the size and service address."
            )
        return

    for message in state["messages"]:
        if isinstance(message, HumanMessage):
            role = "user"
        elif isinstance(message, AIMessage):
            role = "assistant"
        else:
            continue
        with st.chat_message(role):
            st.write(str(message.content))

For example, the _render_messages function displays the conversation history as Streamlit chat bubbles. It receives the conversation through the latest LangGraph state using state["messages"] . If the conversation has no messages yet, the function shows an initial greeting.

Let’s see how it works

We’ve gone over the streamlit_app.py to learn how the page is structured. It’s time to see it in action:

We can test it locally using the following command:

poetry run streamlit run customer_service_agent/streamlit_app.py

It will open up a page at http://localhost:8501/ . The page looks like this:

In order to test the agent, we need an OPENAI_API_KEY. It’ll cost you a few cents to test.

Here is a chat example:

I did not give the address and the agent asked for it as we’d expect. Let’s try the same but providing the address in the first message:

Since I have the address in the first message, the agent did not ask for it. I accepted the quote and the agent gave me three options to choose and then completed the booking:

There is a lot more we can do on the interface to make it more user friendly. But we now have a version that runs smoothly and with a nice and clean user interface.

I’m planning to improve the agent and add more functionality such as WhatsApp integration. It may even turned out to be a product that I can sell to some local businesses.

Stay tuned for what’s coming and thank you for reading!

Share this Article
Please enter CoinGecko Free Api Key to get this plugin works.