Python約30行で作る Bedrock x Claude3 のStreamingチャットアプリ

これまで見て見ぬふりをしてきた「Streamlit上でStreaming出力させる」プログラムを作ってみます。

ライブラリのインストール

いつのまにか「langchain-aws」なるものが生まれているので今回は最終的にはそれを使います。

pip install -U boto3 streamlit langchain langchain-aws

まずboto3からStreamingしてみる

以下を参考に、リクエストとレスポンス取得をClaude3に合わせて微調整します。

import boto3
import json

bedrock = boto3.client(service_name='bedrock-runtime')

body = json.dumps({
    "anthropic_version": "bedrock-2023-05-31",
    "max_tokens": 4000,
    "messages": [{"role": "user", "content": "カレーの作り方を説明してください"}]
})

response = bedrock.invoke_model_with_response_stream(modelId="anthropic.claude-3-haiku-20240307-v1:0", body=body)

stream = response.get("body")
if stream:
    for event in stream:
        chunk = event.get("chunk")
        if chunk:
            chunk_json = json.loads(chunk.get("bytes").decode())
            if chunk_json["type"]=="content_block_delta":
                print(chunk_json["delta"]["text"], end="", flush=True)

レスポンスの取得が少しごちゃごちゃしましたが、一応それっぽく動きました。

LangChainからStreamingしてみる

以下を参考にします。

invokeの代わりにstreamを呼び出せばchunkが取れるっぽいのでそういうプログラムを書いてみます。

from langchain_aws import ChatBedrock
import streamlit as st

LLM = ChatBedrock(model_id="anthropic.claude-3-haiku-20240307-v1:0", model_kwargs={"max_tokens": 4000})
for chunk in LLM.stream("カレーの作り方を説明してください"):
    print(chunk.content, end="", flush=True)

こちらもそれっぽく動きました。boto3を直接叩くより圧倒的に短いですね。

Streamlitでチャットアプリを作ってみる

以下のようにCallbackHandlerでon_llm_new_tokenイベントを拾って出力する方法もありそうですが、

もう少し楽に作れそうだったので今回は以下の方法を参考にします。

上記を見てみると、こちらもchainに対してinvokeではなくstreamで呼び出して、それをwrite_streamに突っ込めば良きに計らってくれるようです。

from langchain_aws import ChatBedrock
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage
import streamlit as st

# chainの定義
prompt = ChatPromptTemplate.from_messages([
        ("system","あなたはAIチャットボットです"),
        MessagesPlaceholder(variable_name="messages"),
        MessagesPlaceholder(variable_name="human_message")])
LLM = ChatBedrock(model_id="anthropic.claude-3-haiku-20240307-v1:0", model_kwargs={"max_tokens": 4000})
chain = prompt | LLM

# 初回はsession領域を作成
if "messages" not in st.session_state:
    st.session_state["messages"] = []

# 2回目以降はsessionを元に全量再描画を行う        
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# 入力された場合
if user_prompt := st.chat_input():
    # 入力内容の描画
    with st.chat_message("user"):
        st.write(user_prompt)
    # chainをstreamで実行し生成内容をstreamで出力
    with st.chat_message("assistant"):
        response = st.write_stream(chain.stream({"messages": st.session_state.messages, "human_message": [HumanMessage(content=user_prompt)]}))

    # 入力内容と生成内容をsession領域に格納
    st.session_state.messages.append({"role": "user", "content": user_prompt})
    st.session_state.messages.append({"role": "assistant", "content": response})

結構シンプルなプログラムですが、期待通りに動作しました。
上記プログラムの本体はほぼ以下の1行のみです。（ストリーミングで生成してストリーミングで出力させている）

st.write_stream(chain.stream({"messages": st.session_state.messages, "human_message": [HumanMessage(content=user_prompt)]}))

静止画だと分かりませんが、Haikuの爆速でストリーミング出力される様はなかなか壮観です。

永続化はしていませんが、チャット履歴をセッション上に持っているのでチャットアプリとして機能します。