BedrockでWebSocketとLambda関数URLを使ったストリーミング実装

に公開

はじめに

LLMを使ったアプリケーションを開発していると、レスポンスの待ち時間が長くてユーザー体験が今ひとつ...という状況になってしまうこともあると思います。
特に全てのテキスト生成が完了するまでユーザーが待ち続けなければならない構成では、体感的な遅延が大きくなりがちです。

今回は、Amazon BedrockのレスポンスをストリーミングAPIで扱う実装について紹介します。
WebSocketを使った構成と、Lambda関数URLを使った構成の2パターンを実装してみました。

課題点

下記の構成では、Bedrockが全てのテキストを生成し終わるまでクライアントは待ち続けることになります。

Client → API GW(REST API) → Lambda → Bedrock
シーケンス図

200文字程度のレスポンスでも、体感で数秒(5、6秒)ほど待たされる感覚です。
ChatGPTやClaudeのように、生成されたテキストが逐次表示される体験に慣れていると、この待ち時間は少し長く感じられます。

使用したBedrockモデル

今回はClaude 4.5 Sonnetの日本向け推論プロファイルを使用しました。

jp.anthropic.claude-sonnet-4-5-20250929-v1:0
推論プロファイル

このプロファイルを使うと、日本国内(東京 or 大阪)でクロスリージョン推論が行われます。
地理的な面でも国外にルーティングされるよりかは、国内に閉じることでレイテンシが少しでも早くならないのかとの思いです(実際のところはわからないです)。

改善策:ストリーム処理の活用

Bedrockのストリーム可能なAPI(invoke_model_with_response_stream)を活用し、生成されたテキストを逐次クライアントに返す仕組みを実装しました。

検証した構成は以下の2パターンです。

パターン1: WebSocketを使った構成

User → フロント画面(React) → API Gateway (WebSocket) → Lambda → Bedrock
構成

シーケンス図

事前にHTTPからプロトコルをアップグレードしてWebsocketでの接続を確立していれば、
Bedrockから返ってきたチャンクをそのままWebSocketのフレームで送り返す実装です。

実装

1. Bedrockのストリーミング処理

import json
import boto3

bedrock = boto3.client("bedrock-runtime", region_name="ap-northeast-1")
model_id = "jp.anthropic.claude-sonnet-4-5-20250929-v1:0"

payload = {
    "anthropic_version": "bedrock-2023-05-31",
    "max_tokens": 2048,
    "messages": [
        {"role": "user", "content": [{"type": "text", "text": "日本の首都は?"}]}
    ]
}

# ストリーム形式のレスポンスを受け取り
response = bedrock.invoke_model_with_response_stream(
    modelId=model_id,
    body=json.dumps(payload)
)

# チャンクを逐次受信
for event in response["body"]:
    if "chunk" in event:
        chunk = json.loads(event["chunk"]["bytes"])
        if chunk["type"] == "content_block_delta":
            print(chunk["delta"]["text"], end="", flush=True)

invoke_model_with_response_streamAPIを使用しています。
通常のinvoke_modelだと生成完了後に一括でレスポンスが返ってきますが、ストリーミングAPIなら生成中に逐次チャンクでレスポンスを取得できます。

DevTools > WebSocket から確認したら、こま切れのメッセージが返ってきてるのを確認できました。

2. WebSocketでのフレーム送信

Lambda内で受け取ったチャンクを、API Gatewayのpost_to_connectionで送り返します。

api_gateway = boto3.client('apigatewaymanagementapi', 
                            endpoint_url=f'https://{domain_name}/{stage}')

for event in response["body"]:
    if "chunk" in event:
        chunk = json.loads(event["chunk"]["bytes"])
        if chunk["type"] == "content_block_delta":
            text = chunk["delta"]["text"]
            # WebSocketでクライアントに送信
            api_gateway.post_to_connection(
                ConnectionId=connection_id,
                Data=text
            )

パターン2: 関数URL化したLambdaのレスポンスストリーミング

ユーザ → フロント画面(React) → Lambda Function URL (Response Streaming) → Bedrock
構成

コンソール画面からLambdaを関数URL化します。

関数URLを作成すると、以下の形式のエンドポイントが発行されます。

https://<function-id>.lambda-url.<region>.on.aws/

レスポンスストリームの設定

関数URL化したLambdaには設定オプションで「RESPONSE_STREAM」という項目があります。

これを使えば、WebSocketを使わなくてもHTTPレスポンスペイロードを元のクライアントにストリーミングできます。
このオプション設定を有効化した場合のLambdaは、現時点(2025/10/22)ではマネージドランタイムだとNode.jsのみサポートされています。

参考:Lambda 関数のレスポンスストリーミング

実装

import { streamifyResponse } from 'aws-lambda-stream';

export const handler = streamifyResponse(
  async (event, responseStream) => {
    const { prompt } = JSON.parse(event.body);
    
    responseStream.setContentType('text/plain');
    
    // Bedrockのストリーミングレスポンス
    const response = await bedrock.send(
      new InvokeModelWithResponseStreamCommand({
        modelId: 'jp.anthropic.claude-sonnet-4-5-20250929-v1:0',
        body: JSON.stringify({
          anthropic_version: 'bedrock-2023-05-31',
          max_tokens: 4096,
          messages: [
            {
              role: 'user',
              content: [{ type: 'text', text: prompt }]
            }
          ]
        })
      })
    );
    
    // チャンクを逐次送信
    for await (const chunk of response.body) {
      if (chunk.chunk?.bytes) {
        const text = JSON.parse(
          new TextDecoder().decode(chunk.chunk.bytes)
        ).delta?.text;
        if (text) responseStream.write(text);
      }
    }
    
    responseStream.end();
  }
);

Bedrockのストリーム処理をstreamifyResponse()でラップするイメージです。
クライアントに逐次レスポンスを返せます。API GWを立てないでいい分WebSocketよりもシンプルな構成だと思います。

レスポンス時間

下記の3パターンで実際に10回ほどコールして、Webブラウザの画面上にBedrockで生成された最初の1文字が画面に表示されるまでの時間を計ってみました。

  1. WebSocket(レスポンスを逐次返却)
  2. 関数URL化したLambda(レスポンスを逐次返却)
  3. REST API(レスポンスを全て生成してから返却)
計測条件
  • プロンプト:旅行で東京に行った際のおすすめ観光地を200文字くらいで教えてください
  • Lambdaの状態:ウォームスタート状態
  • Lambdaの設定値:
    • メモリ:128MB
    • タイムアウト:60秒
    • ランタイム:
      • 関数URL化ver:Node.js 20.x
      • それ以外:python3.12
    • アーキテクチャ:arm64

結果、12に関しては、約2秒程度で、3に関しては6秒程度最初の1文字が表示されるのに時間がかかっていました。
WebSocketとLambda関数URLの差は、体感的な違いはほとんど感じられなかったです。
WebSocket、関数URL化したLambda どちらの構成も従来のREST API構成と比べてかなり早くレスポンスを表示できるようになりました。

もっと早くするには

今回の検証では実施していませんが、さらに速度を追求するなら以下の方法も考えられます。

  • 軽量で高速な Claude 4.5 Haiku モデルを使用
  • Bedrockのプロビジョンドスループットを購入してモデル実行枠を確保
  • Lambdaのメモリを増やす(今回は128MBで検証)
  • Bedrockのプロンプトキャッシュを活用

参考資料

Discussion