API GatewayでBedrockのストリーム応答を試してみた

AWSTemplateFormatVersion: '2010-09-09'
Description: 'API Gateway Response Streaming with Lambda and Bedrock Haiku 4.5'

Resources:
  LambdaExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: BedrockInvokePolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - bedrock:InvokeModelWithResponseStream
                Resource: '*'

  StreamingLambda:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: BedrockStreamingFunction
      Runtime: nodejs20.x
      Handler: index.handler
      Role: !GetAtt LambdaExecutionRole.Arn
      Timeout: 300
      Code:
        ZipFile: |
          const { BedrockRuntimeClient, InvokeModelWithResponseStreamCommand } = require("@aws-sdk/client-bedrock-runtime");

          const client = new BedrockRuntimeClient({ region: process.env.AWS_REGION });

          exports.handler = awslambda.streamifyResponse(async (event, responseStream, context) => {
            const httpResponseMetadata = {
              statusCode: 200,
              headers: {
                'Content-Type': 'text/plain',
                'Access-Control-Allow-Origin': '*'
              }
            };

            responseStream = awslambda.HttpResponseStream.from(responseStream, httpResponseMetadata);

            try {
              const body = JSON.parse(event.body || '{}');
              const prompt = body.prompt || "こんにちは";

              const command = new InvokeModelWithResponseStreamCommand({
                modelId: "jp.anthropic.claude-haiku-4-5-20251001-v1:0",
                contentType: "application/json",
                accept: "application/json",
                body: JSON.stringify({
                  anthropic_version: "bedrock-2023-05-31",
                  max_tokens: 100000,
                  messages: [{
                    role: "user",
                    content: prompt
                  }]
                })
              });

              const response = await client.send(command);

              for await (const chunk of response.body) {
                if (chunk.chunk?.bytes) {
                  const text = JSON.parse(new TextDecoder().decode(chunk.chunk.bytes));
                  if (text.type === 'content_block_delta' && text.delta?.text) {
                    responseStream.write(text.delta.text);
                  }
                }
              }

              responseStream.end();
            } catch (error) {
              responseStream.write(`Error: ${error.message}`);
              responseStream.end();
            }
          });

  LambdaInvokePermission:
    Type: AWS::Lambda::Permission
    Properties:
      FunctionName: !Ref StreamingLambda
      Action: lambda:InvokeFunction
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${RestApi}/*'

  RestApi:
    Type: AWS::ApiGateway::RestApi
    Properties:
      Name: BedrockStreamingAPI
      Description: API Gateway with response streaming for Bedrock

  StreamResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId: !Ref RestApi
      ParentId: !GetAtt RestApi.RootResourceId
      PathPart: stream

  StreamMethod:
    Type: AWS::ApiGateway::Method
    Properties:
      RestApiId: !Ref RestApi
      ResourceId: !Ref StreamResource
      HttpMethod: POST
      AuthorizationType: NONE
      Integration:
        Type: AWS_PROXY
        IntegrationHttpMethod: POST
        Uri: !Sub 'arn:aws:apigateway:${AWS::Region}:lambda:path/2021-11-15/functions/${StreamingLambda.Arn}/response-streaming-invocations'
        ResponseTransferMode: STREAM
        TimeoutInMillis: 300000

  OptionsMethod:
    Type: AWS::ApiGateway::Method
    Properties:
      RestApiId: !Ref RestApi
      ResourceId: !Ref StreamResource
      HttpMethod: OPTIONS
      AuthorizationType: NONE
      Integration:
        Type: MOCK
        IntegrationResponses:
          - StatusCode: 200
            ResponseParameters:
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key'"
              method.response.header.Access-Control-Allow-Methods: "'POST,OPTIONS'"
              method.response.header.Access-Control-Allow-Origin: "'*'"
            ResponseTemplates:
              application/json: ''
        RequestTemplates:
          application/json: '{"statusCode": 200}'
      MethodResponses:
        - StatusCode: 200
          ResponseParameters:
            method.response.header.Access-Control-Allow-Headers: true
            method.response.header.Access-Control-Allow-Methods: true
            method.response.header.Access-Control-Allow-Origin: true

  Deployment:
    Type: AWS::ApiGateway::Deployment
    DependsOn:
      - StreamMethod
      - OptionsMethod
    Properties:
      RestApiId: !Ref RestApi

  Stage:
    Type: AWS::ApiGateway::Stage
    Properties:
      RestApiId: !Ref RestApi
      DeploymentId: !Ref Deployment
      StageName: prod

Outputs:
  ApiEndpoint:
    Description: API Gateway endpoint URL
    Value: !Sub 'https://${RestApi}.execute-api.${AWS::Region}.amazonaws.com/prod/stream'

  TestCommand:
    Description: Test command using curl
    Value: !Sub |
      curl --no-buffer -X POST https://${RestApi}.execute-api.${AWS::Region}.amazonaws.com/prod/stream \
        -H "Content-Type: application/json" \
        -d '{"prompt":"日本のAWSリージョンについて教えてください"}'