AWS Bedrockで回答の続きを取得する方法

方法

• invoke_modelを使います。
• レスポンスが "stop_reason": "max_tokens" のとき、messages に {"role": "assistant", "content": [{"type": "text", "text": {生成されたテキスト}}]を追加して、再度invoke_modelを実行します。
• レスポンスが "stop_reason": "end_turn" になるまで続けます。

環境

• Python 3
• boto3
• モデル: Claude 3 Haiku

サンプルコード

max_tokens で打ち切られるたびに応答を再リクエストし、生成されたテキストを連結しています。stop_reason が "end_turn" になるまでこれを繰り返す仕組みです。

InvokeModel API

import json

import boto3

client = boto3.client("bedrock-runtime")
model_id = "anthropic.claude-3-haiku-20240307-v1:0"


def main() -> str:
    full_text = ""
    body_params = {
        "anthropic_version": "bedrock-2023-05-31",
        "max_tokens": 10,   # 意図的に小さい値にしています
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": "カエルについての俳句を作ってください"},
                ],
            },
        ],
    }

    while True:
        params = {
            "modelId": model_id,
            "body": json.dumps(body_params),
            "accept": "application/json",
            "contentType": "application/json",
        }

        print("Invoking model...")
        response = client.invoke_model(**params)
        response_body = json.loads(response.get("body").read())

        for content in response_body.get("content", []):
            full_text += content["text"].strip()

        stop_reason = response_body["stop_reason"]
        print("Stop reason: ", stop_reason)
        if stop_reason != "max_tokens":
            break

        body_params["messages"].append(
            {
                "role": "assistant",
                "content": [{"type": "text", "text": full_text}],
            },
        )

    return full_text


if __name__ == "__main__":
    generated_text = main()

    print("--------------------------")
    print(generated_text)

実行結果例

Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  max_tokens
Invoking model...
Stop reason:  end_turn
--------------------------
はい、カエルについての俳句を作ってみました。

池のほとりで
静かに鳴く小さな声夕暮れの影自然の中でのカエル。静かに響く声、そして夕暮れの静けさを感じさせてくれます。

Converse API

import boto3

client = boto3.client("bedrock-runtime")
model_id = "anthropic.claude-3-haiku-20240307-v1:0"


def main(
    user_prompts: list[str],
    assistant_propmts: list[str] | None = None,
    system: str | None = None,
) -> str:
    if assistant_propmts is None:
        # 初期化
        assistant_propmts = []

    # ユーザープロンプト
    user_message = {
        "role": "user",
        "content": [{"text": user_prompt} for user_prompt in user_prompts],
    }
    messages = [user_message]

    if assistant_propmts:
        # アシスタントプロンプト
        assistant_message = {
            "role": "assistant",
            "content": [{"text": assistant_propmt} for assistant_propmt in assistant_propmts],
        }
        messages.append(assistant_message)

    params = {
        "messages": messages,
        "modelId": model_id,
        "inferenceConfig": {
            "maxTokens": 10,  # 意図的に小さい値にしています
            "temperature": 0,
        },
    }

    if system:
        # システムプロンプト
        params["system"] = [{"text": system}]

    # 実行
    print("conversing...")
    response = client.converse(**params)

    # 生成テキスト
    response_text = (
        response.get("output", {}).get("message", {}).get("content", [{}])[0].get("text")
    )
    if response_text is None:
        msg = "No response from Bedrock"
        raise Exception(msg)  # noqa: TRY002

    # 生成テキストを assistant_propmts に追加する
    assistant_propmts.append(response_text)

    stop_reason = response["stopReason"]
    print("Stop reason:", stop_reason)
    if stop_reason == "max_tokens":
        # 不完全な生成の場合は、続きを生成する
        return main(user_prompts, assistant_propmts)

    # テキスト生成が完了したので、生成テキスト群を文字列にして戻す
    return "".join(assistant_propmts)


if __name__ == "__main__":
    generated_text = main(user_prompts="カエルについての俳句を作ってください")
    print(generated_text)

実行結果例

conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: max_tokens
conversing...
Stop reason: end_turn
はい、カエルについての俳句を作ってみました。

池の水面に
跳ねる小さな影
夏の夕暮れ

緑の葉陰で
鳴き声響き渡る
月夜のカ エル

蛙の声に
心和む静寂
初夏の夜

いかがで しょうか。カエルの姿や鳴き声、季節感を表現してみました。俳句の短い形式の中に、カエルの魅力を感じ取っていただければ幸いです。