How to use multi agent LangGraph with Streamlit to stream the contents as soon as it is generated?

[ayuranjan]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

python ... 

  
    def agent_node(state, agent, name):
        result = agent.invoke(state)
        return {"messages": [HumanMessage(content=result["output"], name=name)]}
    
llm = ChatOpenAI(model="gpt-4-1106-preview")

    vision_language_agent = AgentUtils.create_agent( llm, toolsForVisionLanguageAgent(),AgentPrompts.vision_language_agent_prompt)
    vision_language_node = functools.partial(agent_node, agent=vision_language_agent, name="Vision_Language_Agent")

    video_agent = AgentUtils.create_agent( llm, toolsForVideoAgent(), AgentPrompts.video_agent_prompt)
    video_node = functools.partial(agent_node, agent=video_agent, name="Video_Agent")

    language_agent = AgentUtils.create_agent( llm, toolsForLanguageAgent(), AgentPrompts.language_agent_prompt)
    language_node = functools.partial(agent_node, agent=language_agent, name="Language_Agent")

    supervisor_agent = AgentUtils.create_team_supervisor( llm, AgentPrompts.supervisor_agent_prompt,["Language_Agent","Vision_Language_Agent", "Video_Agent"])


manager = MultiAgentGraphManager(vision_language_node, video_node, language_node, supervisor_agent)

    chatBot(manager.research_chain)



To Run this I did this -----> 



def without_steaming(research_chain, prompt, config):
    previous_message = None
    final_output = None

    for s in research_chain.stream({"messages": [prompt]}, config):
        if "__end__" not in s:
            if 'supervisor' in s:
                agent_name = s['supervisor']['next']
                if agent_name != 'FINISH':
                    st.write(f"Supervisor calling \"{agent_name}\"")
                    print(f"Supervisor calling \"{agent_name}\"")
            elif 'Language_Agent' in s:
                message = s['Language_Agent']['messages'][0].content
                st.write(f"Language Agent: {message}")
                print(f"Language Agent: {message}")
            elif 'Vision_Language_Agent' in s:
                message = s['Vision_Language_Agent']['messages'][0].content
                st.write(f"Vision Language Agent: {message}")
                print(f"Vision Language Agent: {message}")
            elif 'Video_Agent' in s:
                message = s['Video_Agent']['messages'][0].content
                st.write(f"Video Agent: {message}")
                print(f"Video Agent: {message}")

            # Keep track of the final output message
            if 'supervisor' in s and s['supervisor']['next'] == 'FINISH':


                final_output = previous_message
            previous_message = s


        if final_output:
            agent_name = list(final_output.keys())[0]
            final_message = final_output[agent_name]['messages'][0].content
            # st.markdown(f"Final Output: {final_message}")
            print(f"Final Output: {final_message}")
            
            

def chatBot(research_chain):
    config = {"configurable": {"thread_id": 42}}
    if prompt := st.chat_input():
        st.chat_message("user").write(prompt)
        with st.chat_message("assistant"):
            without_steaming(research_chain , prompt= prompt, config=config)

Description

I am trying to make a multi agent system with langgraph ( similar to Research Team mentioned in the documentation example here - https://github.com/langchain-ai/langgraph/blob/main/examples/multi_agent/hierarchical_agent_teams.ipynb -- create_agent and create_team_supervisor are taken from the example ).

manager.research_chain returns the compiled graph

My compiled graph looks like -

I read the streamline document as well as the documentation mentioned in langchain regarding using streamlit.

I am not able to stream the output( i.e I want to print the chunks as soon as It is generated but currently its waiting for the final result and then gives the whole result at once).

I tried this

for s in research_chain.stream({"messages": [prompt]}, config):
if "end" not in s:
st.write_stream(s)

But this prints the agent name only.

I tried this -

for s in research_chain.stream({"messages": [prompt]}, config):
if "end" not in s:
st.write_stream(s)

But this also waits for the output of each agent to get completed and then display the output of each agent at once.
Ideally I would want it to display the final contents of the agent in iteration / stream.

I also thought this could be due to the fact that, our child agent is running using .invoke() in the def agent_node but changing that to stream doesn't makes sense since these will be internal calls and we need the full result to decide on what to do next. But at least it would be better if the final result of each agent could be steamed so the user doesn't feel the lag since for some uses cases it would take time to complete them, so if the user knows about what is going on then it would be better. Any suggestion to improve this user experience is welcomed.

System Info

System Information

OS: Darwin
OS Version: Darwin Kernel Version 23.5.0: Wed May 1 20:14:38 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T6020
Python Version: 3.9.6 (default, Feb 3 2024, 15:58:27)
[Clang 15.0.0 (clang-1500.3.9.4)]

Package Information

langchain_core: 0.2.22
langchain: 0.2.10
langchain_community: 0.2.7
langsmith: 0.1.85
langchain_anthropic: 0.1.19
langchain_experimental: 0.0.62
langchain_openai: 0.1.17
langchain_text_splitters: 0.2.2
langchainhub: 0.1.20
langgraph: 0.1.7

[dosubot[bot]]

Hey @ayuranjan! I'm here to help with any questions or bugs you might have. Let's crack this code conundrum together! 🤖

To use multi-agent LangGraph with Streamlit to stream the contents as soon as they are generated, you can follow the example provided for using an agent with tools. Here is a sample code snippet that demonstrates how to set this up:

import streamlit as st
from langchain import hub
from langchain.agents import AgentExecutor, create_tool_calling_agent, load_tools
from langchain_openai import OpenAI
from langchain_community.callbacks.streamlit import StreamlitCallbackHandler
from langchain_core.prompts import ChatPromptTemplate

# Initialize the language model with streaming enabled
llm = OpenAI(temperature=0, streaming=True)

# Load the necessary tools
tools = load_tools(["ddg-search"])

# Define the prompt
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant."),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ]
)

# Create the agent
agent = create_tool_calling_agent(llm, tools, prompt)

# Create the agent executor
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# Streamlit app setup
if prompt := st.chat_input():
    st.chat_message("user").write(prompt)
    with st.chat_message("assistant"):
        # Initialize the Streamlit callback handler
        st_callback = StreamlitCallbackHandler(st.container())
        
        # Invoke the agent executor with the callback handler
        response = agent_executor.stream(
            {"input": prompt}, {"callbacks": [st_callback]}
        )
        
        # Display the response
        for chunk in response:
            st.write(chunk["output"])

Make sure to set the OPENAI_API_KEY for the above app code to run successfully. You can do this via Streamlit's secrets.toml or any other local environment management tool. This setup will allow you to stream the contents generated by the multi-agent LangGraph in real-time within a Streamlit app.

Additionally, LangChain provides methods like .stream() and .astream() for synchronous and asynchronous streaming, respectively. These methods return an iterator that you can consume with a simple for loop, allowing you to handle intermediate results before the final response is ready.

LangChain Expression Language (LCEL) can also be used to build chains that support streaming. For example:

from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template("tell me a joke about {topic}")
parser = StrOutputParser()
chain = prompt | model | parser

async for chunk in chain.astream({"topic": "parrot"}):
    print(chunk, end="|", flush=True)

This setup ensures that streaming works even when using parsers at the end of the chain. LCEL chains benefit from an automatic implementation of stream and astream, allowing streaming of the final output ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

[rossedman]

This response fixed many of the same problems I was having

[ayuranjan]

So you mentioned this

 response = agent_executor.stream(
            {"input": prompt}, {"callbacks": [st_callback]}
        )
        
        # Display the response
        for chunk in response:
            st.write(chunk["output"])

Which is equivalent to

# This is needed since its a langraph which requires the thread_id for memory 
config = {"configurable": {"thread_id": 42}, 
              "callback": [StreamlitCallbackHandler(st.container())]}


for chunk in research_chain.stream( {"messages": [prompt]},config = config) :
                st.write(chunk)

But this gives the result of the Nodes of the graph i.e First it returns the final result by Supervisor, then it returns the full final result from language agent. Now since I asked who is Cristiano Ronaldo, it took it some time for the language agent to return result. Once it completely returned the answer , it was displayed at once in streamlit. Ideally it should write in chunks so that it looks like someone is typing. But currently it waits and return the full answer at once from each agent. I believe this is due to the fact that child agents are called using invoke so it returns the final answer at once.

[JoshuaC215]

I have some working code for streaming langgraph results and streamlit here in case it's a useful reference:

https://github.com/JoshuaC215/agent-service-toolkit

(I was looking at this discussion during research for building it :) )

[shiv248]

@ayuranjan if you're still looking I'd check out this group of cookbooks of common UI/UX implementations for LangGraph with Streamlit here. github.com/shiv248/Streamlit-x-LangGraph-Cookbooks.

[joluoch]

hi guys, Im new in using langgraph and im having problems with streamlit and langgraph when it comes to handling user feedback. I would like some help on how I can make the ai wait at least for the user to read its response and provide some feedback before continuing with the graph. here is my langgraph code and my streamlit UI implementation. Guys please help me

langgragh.py

`import os
import asyncio
from typing import Literal
from typing_extensions import TypedDict
from langgraph.graph import MessagesState, END
from langgraph.types import Command
from db_queries import *
from langchain_core.messages import HumanMessage, AIMessage
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from Prompts import *
from langgraph.types import interrupt
from langgraph.checkpoint.memory import MemorySaver
import uuid
import time

os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY")
os.environ["LANGSMITH_API_KEY"] = os.environ.get("LANGSMITH_API_KEY")
os.environ["LANGSMITH_TRACING"] = os.environ.get("LANGSMITH_TRACING")
os.environ["LANGSMITH_ENDPOINT"] = os.environ.get("LANGSMITH_ENDPOINT")
os.environ["LANGSMITH_PROJECT"] = os.environ.get("LANGSMITH_PROJECT")

user_id = st.session_state.get("user_id")[0]

def get_stage_members(user_id):
s =st.session_state['stage_informations'][st.session_state['stage_index']]['flows']['flow_steps']
members = []
for i in range(len(s)):
agent_id = s[i]['agent_id']
member = get_agents(agent_id,user_id)
# st.write(member)
s[i]['agent_name']= member[0]['name']
s[i]['agent_skills']=member[0]['skills']

    members.append(member[0]['name'].replace(" ", "_"))


return members

llm = ChatOpenAI(model_name="gpt-4", temperature=0.7,max_tokens=1000)

class State(MessagesState):
next: str

def create_agent(llm, tools, prompt):
return create_react_agent(llm, tools=tools, prompt=prompt)

def create_agent_nodes(flow_steps):
agents = {}
for step in flow_steps:
# print(step)
requirements = step['requirements']
description = step['description']
agent_skills = step['agent_skills']#"skill:" + step['agent_skills_name'] +"Description:"+ step['agent_description']
agent_name = step['agent_name'].replace(" ", "_")
prompt = assign_agent(agent_name, requirements, agent_skills, description)

    # Assume `llm` object is predefined or passed as parameter
    agent = create_agent(llm, [], prompt)
    agents[agent_name] = agent
return agents

def process_flow_steps(stage_info):
flow_steps = stage_info#['flows']['flow_steps']
return create_agent_nodes(flow_steps)

def process_node(agent,node_name, state: State) -> Command[Literal["human"]]:
result = agent.invoke(state)
return Command(
update={
"messages": [
HumanMessage(content=result["messages"][-1].content, name=node_name)
]
},
goto="human",
)

def get_human_feedback() -> str:
feedback = st.chat_input("How was the conversation?")
return feedback.strip()

def human_node(state: State):# -> Command[Literal["supervisor"]]:
#input = interrupt("Waiting for user feedback...")

human_feedback = st.session_state.get("latest_feedback", "Kindly go to the next step").strip() #Kindly hold for 60 seconds as i read through the response and provide feedback

# if human_feedback == "Kindly hold for 60 seconds as i read through the response and provide feedback" :
#     return Command(goto="supervisor")

return Command(
    update={
        "messages": [
            HumanMessage(content=human_feedback, name="Human")
        ]
    },
    goto="supervisor",
)

def setup_graph(stage_info):

members = get_stage_members(st.session_state.get("user_id", [None])[0])

options = members + ["FINISH"]
supervisor_options = members + ["__end__"]

system_prompt = (
    "You are a supervisor tasked with managing a conversation between the"
    f" following workers: {members}. Given the following user request,"
    " respond with the worker to act next. Each worker will perform a"
    " task and respond with their results and status. When finished,"
    " respond with FINISH."
)

class Router(TypedDict):
    """Worker to route to next. If no workers needed, route to FINISH."""

    next: Literal[tuple(options)]

def supervisor_node(state: State) -> Command[Literal[tuple(supervisor_options)]]:
    messages = [
        {"role": "system", "content": system_prompt},
    ] + state["messages"]
    response = llm.with_structured_output(Router).invoke(messages)
    goto = response["next"]
    if goto == "FINISH":
        goto = END
    return Command(goto=goto, update={"next": goto})


# st.write(stage_info)
agents = process_flow_steps(stage_info)

builder = StateGraph(State)
builder.add_edge(START, "supervisor")
builder.add_node("supervisor", supervisor_node)

for agent_name, agent in agents.items():
    node_name = agent_name
    builder.add_node(f'{node_name}', lambda state, agent=agent,node_name=node_name: process_node(agent,node_name, state))

builder.add_node("human", human_node)  # Added human feedback node


return builder.compile()

`

my streamlit implementation :

`if not brief_output:
st.write("Processing information for this stage....")

        if "conversation_history" not in st.session_state:
            st.session_state["conversation_history"] = []

        
        stages_info = st.session_state['stage_informations'][st.session_state['stage_index']]['flows']['flow_steps']
        

        graph = setup_graph(stages_info)

      
        
        requirements = stages_info[0]['requirements']
        agent_name = stages_info[0]['agent_name']
        agent_skills = stages_info[0]['agent_skills']
        brief  = get_project_details(st.session_state["selected_project_id"])


        brief_prompt = generate_response(agent_name, agent_skills, brief[0]['title'], brief[0]['brand'],  brief[0]['description'], brief[0]['objectives'],
                  brief[0]['target_audience'], brief[0]['budget'], brief[0]['timeline'], brief[0]['website'],
                  previous_output=None, requirements=requirements)
        

        for s in graph.stream({"messages": [("user", f"{brief_prompt}")]}, subgraphs=True):

                first_part, second_part = s

                if not first_part:
                    print(s)
                    if list(second_part.keys())[0] in ['supervisor', '__interrupt__']:
                        print("First key is 'supervisor'")
                        if list(second_part.keys())[0] == 'supervisor' and second_part['supervisor'].get('next') == '__end__':
                            break
                        
                    else:

                        agent_name = list(second_part.keys())[0]
                        
                        content = second_part.get(f"{agent_name}", {}).get("messages", [])

                        for msg in content:
                            content = getattr(msg, "content", "")

                        # with st.chat_message(f"{agent_name}"):
                        #     st.markdown(content)

                        if agent_name == 'human':
                            st.session_state["conversation_history"].append(HumanMessage(content=content, role="Human"))
                        else:
                            st.session_state["conversation_history"].append(AIMessage(content=content, role=f"{agent_name}"))
                # Display conversation history
                for message in st.session_state["conversation_history"]:
                    role = "AI" if isinstance(message, AIMessage) else "Human"
                    with st.chat_message(role):
                        st.markdown(message.content)
                if user_input:

                    st.session_state["latest_feedback"] = user_input.strip()
                    # st.write(f"feedback:{st.session_state['latest_feedback']}")

`