Streamlitで自分用ChatGPTを作る

自分用ChatGPTのようなものを作りたいと思いました。UIの実装は、Streamlitを使うと楽できそうです。

できたもの

というでわけで、こんな感じのものがサクッと作れました（新しい発言が上に表示される仕様です）。UIには前述の通りStreamlitを、ロジック部分にはLangChainを使っています。

こんなことができます。

• 研究アシスタントとして振る舞う
• チャット風のUIで会話できる
• 会話履歴に基づいて会話できる

今のところは、ChatGPTのWeb版を使うのとそんなに変わりません。ソースコードは、以下のリポジトリに置いてあります。

https://github.com/kentaro/research_assistant

開発する上でハマったこと

チャットのロジック部分に関しては、「LangChain の チャットモデル (ChatGPTの新しい抽象化) を試す｜npaka」などを参考にすれば、簡単に実装できるでしょう。

会話履歴基づいて発言できるようにするためには、LangChainのMemory機能を用いるのですが、この実装に少しハマりました。フォームをsubmitすると、会話履歴を忘れてしまうのです（初期化されてしまう）。

そこで、st.cache_resourceを使って会話を司るチェーンを丸ごとキャッシュすることで、以前の履歴を引き継げるようにしました。@st.cache_resourceとして、アトリビュートをつけてあげるだけでOKです。

@st.cache_resource
def load_conversation():
  llm = ChatOpenAI(
    streaming=True,
    callback_manager=CallbackManager([
      StreamlitCallbackHandler(),
      StreamingStdOutCallbackHandler()
    ]),
    verbose=True,
    temperature=0,
    max_tokens=1024
  )
  memory = ConversationBufferMemory(return_messages=True)
  conversation = ConversationChain(
    memory=memory,
    prompt=prompt,
    llm=llm
  )
  return conversation

Streamlitにおけるキャッシュの仕組みの詳細についてはCaching - Streamlit Docsをご覧ください。

解決したいけどできてないこと

現状の実装では、LangChainの処理が終わるまでUI上には結果が表示されません。ChatGPTのWeb UIのように、トークンが返されるたびに表示したいところです。実際には、ストリームでデータを受け取ってはいて、標準出力には随時出力されます（実行してみてください）。

LangChainにはlangchain/langchain/callbacks
/streamlit.pyというものがあって、トークンが出力されるたびに表示することができそうなのですが、動かすことができませんでした。どなたか、動かしてみてほしいと思います。

おわりに

前述の通り、ここまでの実装では「ChatGPTを使えばいいじゃん」という感じです。一方で、ここからカスタマイズしたり情報をあれこれ追加したりしできるので、自分用に作る意義は十分あるように思えます。

• プロンプトをいじることで、最初から好みのキャラクタを持つAIを作り込める
• LangChainを使っているので、様々な外部情報を組み込める
• ChatGPTの有料版よりAPIコールの方がずっと安い

UIを作るのが面倒だったので、今回はStreamlitを使ってみました。ハマりポイントや改善したい点があったので、途中経過を紹介しました。