Langchainのメモリを基本的に状態を持たないstreamlitで実装する方法

streamlitは簡単に実装できますが、その反面状態管理は難しいです。一方でLLMアプリを作ろうと思うと会話の履歴やインデックスなど保持していてほしい情報がまぁまあまあります。

今回はstreamlitで特にlangchainのメモリ(会話の情報)を実装する方法を紹介します。

実装例

streamlitの状態管理　st.session_state

streamlitでは基本的に状態を持たないので変数に何かを入れても何かアクションした瞬間(ボタン押すなど)すぐ消えてしまいます。

そこでst.session_stateが使えます。これはセッション中情報を保存してくれる特別な変数を設定するものです。
使い方としてはまず初期化して

# Initialization
if 'key' not in st.session_state:
    st.session_state['key'] = 'value'

その後この変数に情報を格納

st.session_state['key'] = 'value2'

呼び出すときは

st.write(st.session_state['key'])

とすればvalue2と表示されます。今回はkeyの部分をmemoryとしてその中身に会話の情報を入れていきます。

langchainのメモリ

langchainは色々なメモリの保持の仕方がありますが、今回は一番基本的なConversationBufferMemoryを使います。

まずはConversationChainと一緒にインポートし

from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory

初期化　

def get_state(): 
     if "state" not in st.session_state: 
         st.session_state.state = {"memory": ConversationBufferMemory(memory_key="chat_history")} 
     return st.session_state.state 
state = get_state()

そしてChainの中にstate['memory']を入れます。

chat = ChatOpenAI(streaming=True, temperature=0.9)
        conversation = ConversationChain(
            llm=chat, 
            prompt=prompt,
            memory=state['memory']            
        )

これでchatの会話を保存してくれます。

過去の会話を要約したり、エンティティモデルやグラフとして格納する方法もありますので慣れたら是非memoryを入れ替えて楽しんでください。

chatの内容が保存されているのを確認できるデモアプリも作ったので遊んでみてください。

同じようにstreamlitでlangchainを実装する際にハマるストリーミング(一文字(トークン)ずつアウトプットする)の方法もまとめたので困ったときは読んでみてください。
https://qiita.com/yazoo/items/b318adbab0d5f3f58dc8