Error rendering st.status() container in st.chat_message() component

[rizkhalifa99]

Running into an issue when trying to render st.status component within st.chat_input. I am trying to reproduce the look and feel of the langchain + streamlit callback demo but through a conversation: Streamlit | 🦜️🔗 Langchain. It works if you type in one question, but if you go to type in a second question, it doesn’t get rendered. However if you proceed to type in a third question it does work and the third question becomes your second interaction. The jist is you have to type everything twice for it to work. No errors get thrown in the process.

Another wrinkle is the code works if I remove the cookie manager or I comment out rendering the status container. I use cookies to manage login, so would be ideal to keep.

Reproducable code that DOES NOT WORK:

import os
import streamlit as st
from langchain.schema import ChatMessage
from pydantic import BaseModel
from langchain.schema import ChatMessage
from pydantic import BaseModel
from streamlit_cookies_manager import EncryptedCookieManager
from langchain.callbacks import StreamlitCallbackHandler

cookies = EncryptedCookieManager(
    # This prefix will get added to all your cookie names.
    # This way you can run your app on Streamlit Cloud without cookie name clashes with other apps.
    prefix="ktosiek/streamlit-cookies-manager/",
    # You should really setup a long COOKIES_PASSWORD secret if you're running on Streamlit Cloud.
    password=os.environ.get("COOKIES_PASSWORD", "My secret password"),
)

class AssistantIntermediateStep(BaseModel):
    tool_name: str
    input: str
    output: str

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

messages = st.session_state.messages
for i, msg in enumerate(messages):
    message = msg["message"]
    logs = msg["logs"]
    with st.chat_message(message.role):
        if logs:
            for log in logs:
                with st.status(f"{log.tool_name}: {log.input}", state='complete'):
                    st.write(log.output)

            with st.status("Complete!", state='complete'):
                st.write(message.content)
        
        st.write(message.content)

if prompt:=st.chat_input(key="chat_input"):

    with st.chat_message("user"):
        st.write(prompt)

    with st.chat_message("assistant"):

        # commented out to not incur cost; mocked response is able to reproduce bug; can uncomment to test with any langchain agent
        # st_callback = StreamlitCallbackHandler(st.container())
        # response = agent({"input": prompt}, callbacks=[st_callback]) 
        response = {
            "output": "test output",
        }

        st.write(response["output"])

        # uncomment if using real agent
        # assistant_intermediate_steps = [AssistantIntermediateStep(**{"tool_name": step[0].tool, "input": json.dumps(step[0].tool_input), "output": step[1], "log": step[0].log}) for step in response["intermediate_steps"]]
        st.session_state.messages.append({"message": ChatMessage(role="user", content=prompt), "logs": None})

        st.session_state.messages.append(
            {
                "message": ChatMessage(role="assistant", content=response["output"]),
                "logs": [
                    AssistantIntermediateStep(**{
                        "tool_name": "test",
                        "input": "test input",
                        "output": "test output"
                    })
                ]
            }
        )

WORKS w/o cookie manager:

import os
import streamlit as st
from langchain.schema import ChatMessage
from pydantic import BaseModel
from langchain.schema import ChatMessage
from pydantic import BaseModel
from streamlit_cookies_manager import EncryptedCookieManager
from langchain.callbacks import StreamlitCallbackHandler

# cookies = EncryptedCookieManager(
#     # This prefix will get added to all your cookie names.
#     # This way you can run your app on Streamlit Cloud without cookie name clashes with other apps.
#     prefix="ktosiek/streamlit-cookies-manager/",
#     # You should really setup a long COOKIES_PASSWORD secret if you're running on Streamlit Cloud.
#     password=os.environ.get("COOKIES_PASSWORD", "My secret password"),
# )

... <rest of code>

WORKS w/o logs:

... <rest of code>
messages = st.session_state.messages
for i, msg in enumerate(messages):
    message = msg["message"]
    logs = msg["logs"]
    with st.chat_message(message.role):
        # if logs:
        #     for log in logs:
        #         with st.status(f"{log.tool_name}: {log.input}", state='complete'):
        #             st.write(log.output)

        #     with st.status("Complete!", state='complete'):
        #         st.write(message.content)
        
        st.write(message.content)

... <rest of code>

streamlit version: 1.27.2
python version: 3.11.4
os: mac