Amazon Bedrock Knowledge Bases が RetrieveAndGenerateStream API をサポート

内容

AWS re:Invent 2024 で発表があった Amazon Bedrock Knowledge Bases が RetrieveAndGenerateStream API をサポートしたことについてシンプルに纏めます。

• Knowledge Bases を利⽤する RAG で、完全な応答を待たずに、順次応答を受け取ることができる（ストリーム）
• InvokeModelWithResponseStream の Knowledge Bases みたいな感じ？

対応リージョン

• Knowledge Bases を利用可能な全てのリージョン

設定

前提

ストリームができるかの確認ですので、Knowledge Bases は構築済とします。

retrieve_and_generate_stream API の実行と結果

import boto3
import json

bedrock_agent = boto3.client("bedrock-agent-runtime", region_name="us-west-2")
knowledgeBase_Id = "<ナレッジベースID>"
model_Arn = "arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-3-5-sonnet-20241022-v2:0"

prompt = ("2023年の山手線の利用者数ベスト10を教えて下さい")

response = bedrock_agent.retrieve_and_generate_stream(
    input={"text": prompt},
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            "knowledgeBaseId": knowledgeBase_Id,
            "modelArn": model_Arn,
        },
    }
)
stream = response.get("stream")

if stream:
    for event in stream:
        if "output" in event:
            print(event['output']['text'], end="", flush=True)