Streamlitのtab内でchat_inputを正しく取り扱う

遭遇した課題

Streamlitを用いて、LLMを用いたchat interfaceを含むアプリケーションを実装しようとしていたところ
st.tabs などのレイアウトオブジェクトの中で chat_input のレイアウトが想定しているケースにならないことを発見しました
その際の対処法について記載します

前提

この記事は以下のバージョン構成での実装を前提としています

python==3.11.5
streamlit==1.34.0

通常のChat Interface

Streamlitが公式に公開してくれている、ChatGPTライクなチャットアプリの実装方法については以下に記載されています。
https://docs.streamlit.io/develop/tutorials/llms/build-conversational-apps

このコードを実行すると、以下のように

• ページ上部に st.session_state.messages に格納された会話履歴を表示するコンポーネント
• ページ最下部にユーザーからの入力を受け付けるChat Interfaceのコンポーネント
  というレイアウトで表示されます。

st.tabs内でwrapした場合

複数tabを備えるアプリケーションを実装し、その中の一画面をチャット画面にしようとする場合、以下のような実装になると思います。

tab1, tab2 = st.tabs(["chat", "dummy"])

with tab1: # Chatアプリケーションの実態を tab1 内に記述
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    if prompt := st.chat_input("What is up?"):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            stream = client.chat.completions.create(
                model=st.session_state["openai_model"],
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                stream=True,
            )
            response = st.write_stream(stream)
        st.session_state.messages.append({"role": "assistant", "content": response})

with tab2:
    st.write("dummy tab")

しかし、これをそのまま実行すると以下のように想定していないレイアウトになってしまいます。

※ページ最下部にChat Interfaceが固定されていない

※履歴と現在の会話の間にChat入力用のInterfaceが入る

st.tabs などのレイアウトコンテナ内で chat_input などが表示できるようになったのが 1.31.0 からなので、まだ細かいレイアウトの実装が追いついていないのかなと思います。
https://docs.streamlit.io/develop/quick-reference/changelog#version-1310

回避方法

同じことで困っている方のIssueが既に起票されていました。
https://github.com/streamlit/streamlit/issues/8564

ここに記載されているように、 st.container で st.chat_message を描画する場所を括ってあげると、ある程度想定されるレイアウトに近くなります。

tab1, tab2 = st.tabs(["chat", "dummy"])

with tab1:
    chat_container = st.container(height=600) # st.containerでブロックを定義
    prompt = st.chat_input("What is up?")

    # 以下、`st.chat_message` を `chat_container.chat_message` に置換し、
    # チャットメッセージをcontainer内に表示する
    for message in st.session_state.messages:
        with chat_container.chat_message(message["role"]):
            st.markdown(message["content"])
    

    if prompt:
        st.session_state.messages.append({"role": "user", "content": prompt})
        with chat_container.chat_message("user"):
            st.markdown(prompt)

        with chat_container.chat_message("assistant"):
            stream = client.chat.completions.create(
                model=st.session_state["openai_model"],
                messages=[
                    {"role": m["role"], "content": m["content"]}
                    for m in st.session_state.messages
                ],
                stream=True,
            )
            response = st.write_stream(stream)
        st.session_state.messages.append({"role": "assistant", "content": response})

with tab2:
    st.write("dummy tab")

これを実行すると以下のように描画されます。

展望

ただ、 st.tabs などのレイアウトコンポーネント内でレイアウトルールが変わるのは直感的に実装できないので、上記Issueから修正を後押しするか、余力があれば自分でcontributionしていこうかなと思います。