StreamlitとLangGraphで簡単にエージェントチャットを作る

LangGraphとは

LangGraphはエージェントシステム用のオーケストレーションフレームワークです。
LLMを利用したグラフ構造のワークフローを簡単に作ることができます。
名前の通りLangChainファミリーなので、LangChainとも簡単に組み合わせることができます。

www.langchain.com

Streamlit

Streamlitは簡単なpythonコードでWebアプリを作ることができるライブラリです。
たった数行のコードでコンポーネントを配置しておしゃれなWebアプリを作ることができるので、処理のコードに集中することができます。
LLM用のチャットのコンポーネントも用意されており、今回はこれを使って手軽にUIを作成したいと思います。

streamlit.io

まずチャットアプリ

ということで以前はTypeScriptでLangChainを触っていたのですが、今回はStreamlitとの連携を考えてPythonです。
LangChain自体もJS版よりPython版のほうが機能が豊富ですね。
Pythonはあまりに触ったことがないので、コードのお作法とかで変なところがあったらごめんなさい。
こんな感じのアプリになります。

LLMはGoogle Geminiを使っています。
GoogleのアカウントさえあればAPIキーの発行も簡単ですし、おためしでちょっと触る分には毎日の無料枠で十分なのでよい感じです。
もちろん他のLLMも利用できるので、お好みで試してみてください。

# 環境変数読み込み
load_dotenv()

# タイトルの設定
st.title("Simple Chat App")

# モデルの読み込み
client = ChatGoogleGenerativeAI(
    model="gemini-1.5-flash"
)

まずはこの辺り。APIキーはどこから持ってきても良いですが、個人的には.envで持ちたいのでpython-dotenvを使っています。
ChatGoogleGenerativeAIの場合は「GOOGLE_API_KEY」という環境変数に読み込まれていれば勝手に利用してくれます。

# メッセージの初期化
if "messages" not in st.session_state:
    st.session_state.messages = []

# メッセージの表示
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

ここがStreamlitのUIのところです。
こんなにシンプルなコードでユーザーとアシスタントの会話部分のコンポーネントが作れてしまいます。
st.chat_messageで枠を作成し、st.markdownでコンテンツを書き込むだけです。便利！

# チャットボットとの対話
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"):
        message = client.invoke(st.session_state.messages)
        response = message.content
        st.markdown(response)
    st.session_state.messages.append({"role": "assistant", "content": response})

チャットのやり取りのあたりはこれだけです。
st.chat_inputを使って入力欄を作成し、入力があるとifの中身がレンダリングされます。
st.session_state.messagesはステートなので現在のこの画面の情報を一時的に保存しておく領域です。

何かしらのアクションが起こると画面全体が再レンダリングされます。
ちなみにこの辺りはWebSocketで実現しているようです。
ユーザとアシスタントのメッセージをst.session_state.messagesに追加する処理をコメントアウトしてチャットをしてみると意味が分かりやすいかと思います。
LLMとのやり取りはclient.invokeの部分ですが、これはまだLangChainの機能ですね。
コードはimport部分を除いてこれで全部です。簡単ですね。

ストリーミングで表示

チャットアプリといえば、しばらく待ってから全部のコンテンツが出力されるよりもストリーミングで表示されるほうがUXとして良いですよね。
というわけでストリーミング対応してみましょう。

# モデルの読み込み
client = ChatGoogleGenerativeAI(
    model="gemini-1.5-flash",
    stream=True
)

まずはモデルの作成時にstream=Trueを指定します。

# チャットボットの応答
with st.chat_message("assistant"):
    stream = client.stream(st.session_state.messages)
    response = st.write_stream(stream)
st.session_state.messages.append({"role": "assistant", "content": response})

次にLLMの呼び出しをinvokeからstreamに変えて、Streamlitへの書き込みをst.markdownからst.write_streamに変えます。
これだけでチャンクごとにストリーミングでレンダリングされるようになります。すごい。

エージェントにしてみる

さて、ではここからLangGraphを使ってエージェントアプリにします。
エージェントの定義の仕方はいろいろあって、自分でワークフローを作るのも良いのですが、今回はcreate_react_agentを使うことにします。
ここでいうreactはReactではなくReActのほうです。ややこしい。
簡単に言えば、自分で考えて自分で行動するみたいな振る舞いをします。

@tool
def get_now():
    """Useful for when you need to cuttent datetime."""
    from datetime import datetime
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

まずは、エージェントが利用するツールを定義します。
今回は簡単に現在日時が取れる関数を定義しましょう。LLMはこういうの苦手ですからね。

# ツールの設定
tools = [get_now, DuckDuckGoSearchRun(), WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())]

一つだけだと寂しいので、ありもののツールも2つくらい追加しておきます。
これはあなたはこういうことができますよというのをLLMに教えるためのものです。
なんだか見覚えがありますか？そうです。RAGが出てきたころにすっかり影が薄くなってしまったfunction_callingそのものですね。
最近はfunction_callingだったりtool_callingだったりまちまちで呼ばれているような気がします。

# graphの構築
graph = create_react_agent(client, tools=tools)

そしてcreate_react_agentを使います。カスタマイズしたい場合は自分でグラフ（ワークフロー）を作ると良いです。
システムプロンプトにはhwchase17/reactを使うことが多いですが、今回は後ろの処理の関係でシステムプロンプトはなしです。
なので、ツール自体は認識されていますが、詳しくこう使ってねみたいな指示がないぶん精度がちょっと下がります。

一番変わるのがLLMへのリクエストの部分です。
graph.streamの戻り値をst.write_streamに入れてみるとわかるのですが、いまいちうまく処理してくれません。
なのでLangGraphのstreamingに関するページを参考にコードを変更します。

langchain-ai.github.io

# チャットボットの応答
with st.chat_message("assistant"):

    expander = st.expander("tool use proccess")
    message_placeholder = st.empty()
    contents = ""
    tool_output = ""

    stream = graph.stream({"messages": st.session_state.messages}, stream_mode=["updates", "messages"])
    for token in stream:
        # メッセージでかつツールからの回答ではない場合に表示
        if token[0] == "messages" and isinstance(token[1][0], AIMessageChunk) and token[1][0].content != "":
            contents += token[1][0].content
            message_placeholder.markdown(contents)
        # ツールからの回答の場合に表示
        elif token[0] == "updates":
            if "tools" in token[1]:
                tmp = f"#### Using the tool ： {token[1]['tools']['messages'][0].name}  \nOutput ： {token[1]['tools']['messages'][0].content}"
                tool_output += tmp
                expander.markdown(tmp)

    st.session_state.messages.append({"role": "assistant", "content": contents, "tool": tool_output})

ちょっと長いですが、streamをループして必要な情報を取り出しています。
stream_mode=["updates", "messages"]とすると、処理の流れの情報とメッセージチャンクの両方が取れるようになります。
messagesトークンが取れた場合には画面に表示を、ツールを使ったよという情報が取れた場合にはst.expanderで折りたたんで出力します。

もう少し詳しく

さて、この場合はツールを使い終わると画面に表示されますが、その間がしばらく無反応のように見えてしまいます。
また、ツール呼び出しの結果は見えますが、どうやって呼び出したかもよくわかりません。
streamingのページに記載されているもう一つの方法であるastream_eventsを使ってみます。

# チャットボットの応答
with st.chat_message("assistant"):

    expander = st.expander("tool use proccess")
    message_placeholder = st.empty()
    contents = ""
    tool_outputs = []

    async for event in graph.astream_events({"messages": st.session_state.messages}, version="v1"):
        # メッセージの表示
        if event["event"] == "on_chat_model_stream":
            content = event["data"]["chunk"].content
            if content:
                contents += content
                message_placeholder.markdown(contents)
        # ツール利用の開始
        elif event["event"] == "on_tool_start":
            tmp = f"#### Start using the tool ： {event['name']}  \nInputs: {event['data'].get('input')}"
            tool_outputs.append(tmp)
            expander.markdown(tmp)
        # ツール利用の終了
        elif event["event"] == "on_tool_end":
            tmp = f"#### Finish using the tool ： {event['name']}  \nOutput ： {event['data'].get('output')}"
            tool_outputs.append(tmp)
            expander.markdown(tmp)
    
    st.session_state.messages.append({"role": "assistant", "content": contents, "tools": tool_outputs})

こうするとよりスムーズに詳しい情報まで出るようになります。
astream_eventsは非同期の関数なのでasync forで回します。
なので、全体をasync def main():で関数化して、asyncio.run(main())のような形にするとよいかと思います。

プロンプトに会話履歴をただ突っ込んでいるだけなので、あまりReActっぽい動きはしませんでした。
この辺りはシステムプロンプトをもう少しちゃんとすればいい感じになるとは思います。

UI作らなくていいってすごく楽

いくら便利なエージェントを作っても、UIが使いづらいとなかなか使ってもらえません。
ただ、チャットのUIを自分で作ろうとするとなかなかに大変です。
Streamlitは簡単なアプリを作る際にはすごく良い選択肢になりそうです。

そして今回は簡単な例しか出しませんでしたが、LangGraphはグラフ構造を採用することによってすごく複雑なワークフローも表現することができます。
エージェントアプリを1から作るのは大変ですが、便利なライブラリを使って気軽に作っていきたいですね。
それでは。