API Gateway が REST API でストリーミング対応！Strands TypeScript で試してみた。

はじめに
レスポンスストリーミングとは？
試してみます。
まとめ

はじめに

こんにちは！アプリケーションサービス本部ディベロップメントサービス１課の森山です。

今回は以下のアップデートを確認していきます。元々、Amazon Bedrock AgentCore Gateway 等のアップデートも含めて検証をしていたのですが、少し複雑になるので、複数回に分けて記事にしてみます。

Lambda 自体は以前より Lambda 関数 URL でレスポンスストリーミングに対応していましたが、今回 API Gateway が REST API でもレスポンスストリーミングに対応しました。

これにより、API Gateway + Lambda 構成においても、TTFB（最初の 1 バイトまでの時間）の改善、タイムアウトの最大 15 分への延長、10 MB を超えるペイロードのサポートが実現されます。

今回はStrandsも利用しながら動作確認をしてみました。

レスポンスストリーミングとは？

従来の API Gateway では、バックエンドからの完全なレスポンスをバッファリングしてからクライアントに返していましたが、レスポンスストリーミングでは、バックエンドからデータが利用可能になった時点で段階的にクライアントへ送信します。

特に LLM からのレスポンスをリアルタイムで表示できるため、ユーザー体験が大幅に向上します。

（私は最初、LLM からの応答が文字を打っているかのように表示されるのは演出かと思っていたことは秘密です）

試してみます。

なお、今回は環境を CDK で作成しております。

ソースは以下に格納しておきます。

バックエンドの Lambda の作成

今回はバックエンドを Lambda で作成します。

レスポンスストリーム対応 Lambda を作成する場合、awslambda.streamifyResponse()で関数をラップしてあげる必要があります。

また、簡単に Agent を使いたいので、Strands を使って作成しています。

レスポンスストリーミング対応で少しコードが長いですが、Agent 部分で言うと変数agentだけなので、2 行だけです！

import { Agent } from "@strands-agents/sdk";
  
const agent = new Agent();
  
export const handler = awslambda.streamifyResponse(
  async (event: any, responseStream: any) => {
    const { message } = JSON.parse(event.body || "{}");
    const stream = awslambda.HttpResponseStream.from(responseStream, {
      statusCode: 200,
      headers: {
        "Content-Type": "text/plain",
      },
    });
  
    if (!message) return stream.end("message is required");

    for await (const chunk of agent.stream(message)) {
      if (
        chunk.type === "modelContentBlockDeltaEvent" &&
        chunk.delta.type === "textDelta"
      ) {
        stream.write(chunk.delta.text);
      }
    }
    stream.end();
  }
);

なお、Strands の TypeScript 対応は現時点でプレビュー版なのでご注意ください。詳細は以下の記事で紹介してくれているので、ご確認いただければと思います。

API Gateway の作成

CDK では v2.230.0 から L2 Construct がレスポンスストリーミングに対応していて、Lambda 統合のプロパティresponseTransferModeをapigateway.ResponseTransferMode.STREAMに指定するのみです。

// /chat endpoint with streaming
const chat = this.api.root.addResource("chat");
const integration = new apigateway.LambdaIntegration(this.strandsFunction, {
  proxy: true,
  responseTransferMode: apigateway.ResponseTransferMode.STREAM,
});
chat.addMethod("POST", integration, {
  methodResponses: [{ statusCode: "200" }],
});

マネージドコンソール上だと、統合リクエストの設定にて確認することができます。

なお v2.230.0 で対応する前までは以下のような対応が必要でした。

const method = restApi.root
  .addResource("chat")
  .addMethod("POST", new apigateway.LambdaIntegration(lambda, { proxy: true }));
 
// Lambda Response Streaming設定
const cfnMethod = method.node.defaultChild as apigateway.CfnMethod;
cfnMethod.addOverride(
  "Properties.Integration.Uri",
  `arn:aws:apigateway:${region}:lambda:path/2021-11-15/functions/${lambda.functionArn}/response-streaming-invocations`
);
cfnMethod.addOverride("Properties.Integration.ResponseTransferMode", "STREAM");

Properties.Integration.Uriの上書きをしているところからわかるのですが、レスポンスストリーミングでは、Lambda 統合の ARN フォーマットが異なりますね。

	ARN フォーマット
通常	`arn:aws:apigateway:{region}:lambda:path/2015-03-31/functions/{functionArn}/invocations`
ストリーミング	`arn:aws:apigateway:{region}:lambda:path/2021-11-15/functions/{functionArn}/response-streaming-invocations`

動作確認

curl で確認していきます。

なお、curl でストリーミングレスポンスを確認する場合、--no-buffer（または -N）オプションをつける必要があります。

export API_ENDPOINT="https://xxxxxxx.execute-api.us-east-1.amazonaws.com/prod/chat"
curl -N -X POST $API_ENDPOINT \
  -H "Content-Type: application/json" \
  -d '{"message": "AWSのサーバーレスアーキテクチャについて、初心者にもわかりやすく詳しく説明してください"}'

結果は以下の通りです。

youtu.be

まとめ

API Gateway のレスポンスストリーミング対応により、生成 AI アプリケーションにおいても、API Gateway を利用する選択肢が出てきました！

また、CDK v2.230.0 以降では responseTransferMode プロパティを指定するだけで簡単に実装できます。ロジック部分でも Strands SDK の TypeScript 対応と組み合わせることで、型安全性と開発体験の向上の恩恵を受けつつ、シンプルなコードで AI エージェントのストリーミング API を構築できました。

誰かのお役に立てば幸いです！