streamlitのマルチページアプリでsession_stateをページごとに管理する方法

概要

　StreamlitでマルチページのWebアプリを作る際、session_stateはページをまたいで共有されます。そのため、たとえばpage1.pyでtextというキーを使用している場合、page2.pyのtextも連動してしまいます。そこで、ページごとにsession_stateを管理できるpgssというパッケージを作成しました。

背景

　あなたは上司からある日「StreamlitでChatGPTを組み込んだ複数ページのアプリ作って」と言われたことはあるだろうか？わたしはある。そして、「Streamlitってなんだっけ？」と心の中で思いつつも、「仕事ができる人間だと思われたい」という悲しき性（さが）から、二つ返事で「了解しました！」と答えてしまった。
　しかし、さすが流行りのフレームワークだけあり、チュートリアルを試しているうちに、「なんでも作れそう」という気分になる。実際2～3個のアプリページを作成し、「なんだおれは天才だったのか」と思った矢先、確認作業中に気づいた。「あれ、ページAのメッセージ履歴がページBにも表示されてる...session_stateがページ間で共有されてないか？」

問題の確認

以下のような構成のマルチページWebアプリを作成したとする。

app
├── main.py
├── page1.py
└── page2.py

main.pyの内容

import streamlit as st
pg = st.navigation([st.Page("page1.py"), st.Page("page2.py")])
pg.run()

page1.pyの内容

import streamlit as st
if 'count' not in st.session_state: 
	st.session_state.count = 0
if st.button("Increment"):
    st.session_state.count += 1
st.write("This is page1")
st.write(f"count: ", st.session_state.count)

page2.pyの内容(page1.pyと同じ)

import streamlit as st
if 'count' not in st.session_state: 
	st.session_state.count = 0
if st.button("Increment"):
    st.session_state.count += 1
st.write("This is page2")
st.write(f"count: ", st.session_state.count)

各ページでcountをsession_stateに保持し、ボタンが押されるたびにインクリメントして表示するシンプルな例である。実際にこれを動かすと、まずpage1が開く。ボタンを数回押した後の画面は以下の通り。

次にpage2に移動すると、本来countが0であるべきだが、以下のようにcountが3のままになっている。

このように、session_stateはページ間で共有されているため、複数ページのチャットアプリなどを作成すると、会話履歴がすべてのページで共有されてしまう問題が発生する。

解決方法①：ページ名をsession_stateに含める

　最初に思い付く解決方法は、session_stateのキーにページ固有の名前を追加することだ。たとえば、以下のようにpage1とpage2を修正すると、countの値がページごとに独立するようになる。

page1.pyの内容（ファイル部分以外page2.pyも同様）

import streamlit as st
page_name = "page1.py"
if f"{page_name}count" not in st.session_state: 
	st.session_state[f"{page_name}count"] = 0
if st.button("Increment"):
    st.session_state[f"{page_name}count"] += 1
st.write("This is page1")
st.write(f"count: ", st.session_state[f"{page_name}count"])

page1で数回ボタンを押した状態：

page2に移動した状態。countは独立しており、0になっている。

これで問題は解決したかに見えるが、すべてのコードにページ名を追加する作業は非常に手間がかかる。そこで、趣味でより簡単な解決策を考えて、以下のようなパッケージを作成した。

解決方法②：session_stateのラッパーを用意する

　発想はシンプルで、解決方法①のsession_stateにページ名を追加する部分をラッパーで隠蔽しただけである。このアイデアを基に、pgssというパッケージを作成した。これを使用すると、以下のように記述できる。

page1.py（ファイル名部分以外page2.pyも同様）：

import streamlit as st
from pgss import PageSessionState
ps = PageSessionState("page1.py")

if "count" not in ps: 
   ps.count = 0
if st.button("Increment"):
   ps.count += 1
st.write("This is page1")
st.write(f"count: ", ps.count)

最初にps = PageSessionState("page1.py")でページ固有の管理インスタンスを作成して、あとはst.sessoin_stateの代わりにpsを使うだけである。"page1.py"の部分はページ名以外でもよい。

keyの設定も簡単にできるようにした。

# 普通の書き方
st.text_input("Text input1", key="text_input1")
st.write(st.session_state.text_input1)
# pgssを使った書き方
st.text_input("Text input2", key=ps("text_input2"))
st.write(ps.text_input2)

ページ名を追加する部分を隠しただけのシンプルなパッケージであるが、自分にとってはとても便利である。pgssは以下のコマンドでインストールできるので、興味があればぜひ試してみてほしい。

pip install pgss