class StreamHandler(BaseCallbackHandler):
    def __init__(self, stream_key: str, initial_text=""):
        # self.text = initial_text
        self.key = stream_key

    def on_llm_start(
        self,
        *args,
        **kwargs: Any,
    ) -> Any:
        st.session_state.current_text = ''
        st.session_state.partial = True
        message(st.session_state.current_text, key=key, partial=st.session_state.partial)


    def on_llm_new_token(self, token: str, **kwargs) -> None:
        st.session_state.current_text += token
        st.experimental_rerun()

    def on_llm_end(
        self,
        response: LLMResult,
        **kwargs
    ):
        st.session_state.partial = False

key = f"message_{len(st.session_state.langchain_messages) + 1}"
stream_handler = StreamHandler(key)
response = agent.run(prompt, callbacks=[stream_handler])

class StreamHandler(BaseCallbackHandler):
    def __init__(self, container, initial_text=""):
        self.container = container
        self.text = initial_text

    def on_llm_new_token(self, token: str, **kwargs) -> None:
        self.text += token
        self.container.markdown(self.text)

with st.chat_message("assistant"):
            stream_handler = StreamHandler(st.empty())
            response = st.session_state.agent.run(prompt, callbacks=[stream_handler])