Streaming Chat Component Issue with Button Execution

Using Streamlit
, , ,
1

Hello Streamlit community!
I’m having an issue with the chat component. While conversing with my LLM, I want to display a button that can also execute code, particularly to display a stream. So, I’m using Streamlit buttons with the “on_click” and “kwargs” attributes. The code and logic work fine. However, the stream result displays at the top of the conversation instead of appearing at the end. Could anyone help me solve this issue?

Here’s a complete code to replicate my use case.

import streamlit as st
import time, random

# Simple history init
title = "Streaming problems from function callback"
st.set_page_config(page_title=title)
st.title(title)

if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "Hello"}]

def build_history(messages):
    for msg in messages:
        st.chat_message(msg["role"]).write(msg["content"])

build_history(st.session_state.messages)

# Simple streaming function to illustrate the problem
def stream_single_message_for_demo(message):
    for token in message.split(" "):
        time.sleep(random.random() * 0.15)
        yield f"{token} "

# My callback function
def func(**kwargs):
    full_buffer = st.chat_message("assistant").write_stream(stream_single_message_for_demo(kwargs["content"]))
    st.session_state.messages.append({"role": "assistant", "content": full_buffer})
    
# My chat input
if prompt := st.chat_input():
    # Display user message
    st.chat_message("user").write(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display assistant message, miror of user for this example
    st.chat_message("assistant").write(prompt)
    st.session_state.messages.append({"role": "assistant", "content": prompt})
    
    # Streamlit builder function
    builder = st.button
    # Arguments to the final function
    kargs = { "content": "This is a streamed message from function click, he should appear at the bottom, but is streamed at the top of the page!" }
    # Arguments to streamlit component builder
    args = { "label": "Click to stream the demo", "on_click": func, "kwargs": kargs }
    
    # Building streamlit component
    builder(**args)
2

The reason why it is at the top is because the message is located at callback function and as you can see the callback function is on top of:

if prompt := st.chat_input():

To solve this issue, stream the message under the if prompt ... condition.

We cannot enter under the prompt unless the user write a message. What we need is to enter under prompt even if the user does not send a message.

We will create a variable to enter under prompt if the button is clicked. So we need to revise the callback function as well.

Create a boolean variable to stream or not stream.

if 'is_stream' not in st.session_state:
    st.session_state.is_stream = False

Create another variable to save the content to be streamed.

if 'content_to_stream' not in st.session_state:
    st.session_state.content_to_stream = None

This is the revised callback function. The user has pressed the button. is_stream is activated and content is saved.

def func(**kwargs):
    st.session_state.is_stream = True
    st.session_state.content_to_stream = kwargs["content"]

The revised prompt entry uses the is_stream. We do this because we would like to send the message below the page.

if prompt := st.chat_input() or st.session_state.is_stream:
    if st.session_state.is_stream:
        st.session_state.is_stream = False  # reset
        full_buffer = st.chat_message("assistant").write_stream(stream_single_message_for_demo(st.session_state.content_to_stream))
        st.session_state.messages.append({"role": "assistant", "content": full_buffer})
    else:
        # Display user message
        st.chat_message("user").write(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})

...

Complete code.

"""
https://discuss.streamlit.io/t/streaming-chat-component-issue-with-button-execution/67895
"""


import streamlit as st
import time, random


# Simple history init
title = "Streaming problems from function callback"
st.set_page_config(page_title=title)
st.title(title)


if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "Hello"}]

if 'is_stream' not in st.session_state:
    st.session_state.is_stream = False

if 'content_to_stream' not in st.session_state:
    st.session_state.content_to_stream = None


def build_history(messages):
    for msg in messages:
        st.chat_message(msg["role"]).write(msg["content"])


build_history(st.session_state.messages)


# Simple streaming function to illustrate the problem
def stream_single_message_for_demo(message):
    for token in message.split(" "):
        time.sleep(random.random() * 0.15)
        yield f"{token} "


# My callback function
def func(**kwargs):
    st.session_state.is_stream = True
    st.session_state.content_to_stream = kwargs["content"]
    
    
# My chat input
if prompt := st.chat_input() or st.session_state.is_stream:
    if st.session_state.is_stream:
        st.session_state.is_stream = False
        full_buffer = st.chat_message("assistant").write_stream(stream_single_message_for_demo(st.session_state.content_to_stream))
        st.session_state.messages.append({"role": "assistant", "content": full_buffer})
    else:
        # Display user message
        st.chat_message("user").write(prompt)
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display assistant message, miror of user for this example
        st.chat_message("assistant").write(prompt)
        st.session_state.messages.append({"role": "assistant", "content": prompt})
        
        # Streamlit builder function
        builder = st.button
        # Arguments to the final function
        kargs = { "content": "This is a streamed message from function click, he should appear at the bottom, but is streamed at the top of the page!" }
        # Arguments to streamlit component builder
        args = { "label": "Click to stream the demo", "on_click": func, "kwargs": kargs }
        
        # Building streamlit component
        builder(**args)

Reference:

1 Like
3

Hello @ferdy,

Thank you, it’s working perfectly. I’ve closed this topic.

Have a great day!

1 Like
4

This topic was automatically closed 2 days after the last reply. New replies are no longer allowed.