Claude 3 on Bedrockの超簡単チャットを作る。

LangChainの対応を待ちきれず、Boto3とStreamlitで作りました。

Claude 3のMessages APIに対応しました。
どうしてもストリーミングレスポンスに対応させたかったのでそこだけ一工夫しました。

70行弱です。

import json

import boto3
import streamlit as st

st.header("Claude 3 Sonnet on Bedrock")

if "messages" not in st.session_state:
    st.session_state.messages = []


def create_body(
    user_prompt: str, system_prompt: str = None, assistant_prompt: str = None
):
    st.session_state.messages.append(
        {"role": "user", "content": [{"type": "text", "text": user_prompt}]}
    )
    if assistant_prompt:
        st.session_state.messages.append(
            {
                "role": "assistant",
                "content": [{"type": "text", "text": assistant_prompt}],
            }
        )

    body = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 1000,
        "messages": st.session_state.messages,
    }

    if system_prompt:
        body["system"] = system_prompt

    return body


def stream(response: dict):
    for event in response.get("body"):
        chunk = json.loads(event["chunk"]["bytes"])
        if "delta" in chunk and "text" in chunk["delta"]:
            yield chunk["delta"]["text"]


for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"][0]["text"])


if prompt := st.chat_input("何でも聞いてください。"):

    st.chat_message("user").markdown(prompt)

    client = boto3.client("bedrock-runtime")
    response = client.invoke_model_with_response_stream(
        modelId="anthropic.claude-3-sonnet-20240229-v1:0",
        body=json.dumps(
            create_body(
                user_prompt=prompt, system_prompt="フレンドリーに回答してください。"
            )
        ),
    )

    with st.chat_message("assistant"):
        response = st.write_stream(stream(response))
    st.session_state.messages.append(
        {"role": "assistant", "content": [{"type": "text", "text": response}]}
    )