Select your cookie preferences

We use cookies and similar tools to enhance your experience, provide our services, deliver relevant advertising, and make improvements. Approved third parties also use these tools to help us deliver advertising and provide certain site features.

WebSocket API to Lambda to Bedrock with streaming response

Created with SnapAPI Gateway WebSocket APIAWS LambdaAmazon Bedrock

Creates an API Gateway WebSocketand Lambda functions that provides a streaming response from the LLMs in Amazon Bedrock.

This sample project demonstrates how to use WebSocket API as a front door to Lambda functions that perform inference on Amazon Bedrock LLMs using the [InvokeModelWithResponseStream API, where the LLM responses are returned in a stream.
With WebSocket API, developers of interactive LLM chatbot interfaces can provide a better user experience by displaying the LLM responses as they are generated (which at times can be long, rather than relying on the synchronous InvokeModel API request.
This pattern deploys one API Gateway WebSocket API, four Lambda functions, and one DynamoDB table.

AWSTemplateFormatVersion: 2010-09-09
Transform: 'AWS::Serverless-2016-10-31'
Description: An Amazon API Gateway WebSocket API and an AWS Lambda function.

# Global values that are applied to all applicable resources in this template
Globals:
  Function:
    CodeUri: ./src
    Runtime: python3.11
    MemorySize: 128
    Timeout: 300

Metadata:
  AWS::CloudFormation::Interface:
    ParameterGroups:
      - Label:
          default: "Amazon DynamoDB Configuration"
        Parameters:
          - WebSocketConnectionsTableNameSuffix
          - ReadCapacityUnits
          - WriteCapacityUnits
    ParameterLabels:
      ReadCapacityUnits:
        default: Read Capacity Units
      WriteCapacityUnits:
        default: Write Capacity Units
      WebSocketConnectionsTableNameSuffix:
        default: WebSocket Connections Table Name Suffix

Parameters:
  WebSocketConnectionsTableNameSuffix:
    Type: String
    Default: 'websocket_connections'
    Description: >
      (uksb-1tthgi812) (tag:apigw-websocket-api-bedrock-streaming)
      (Required) The suffix of the name of the Amazon DynamoDB table where connection identifiers 
      will be stored. The table name will be prefixed by the Stack Name.
    MinLength: 3
    MaxLength: 50
    AllowedPattern: ^[A-Za-z_]+$
    ConstraintDescription: 'Required. Can be characters and underscore only. No numbers or special characters allowed.'
  ReadCapacityUnits:
    Description: Provisioned read throughput
    Type: Number
    Default: 5
    MinValue: 5
    MaxValue: 10000
    ConstraintDescription: must be between 5 and 10000
  WriteCapacityUnits:
    Description: Provisioned write throughput
    Type: Number
    Default: 5
    MinValue: 5
    MaxValue: 10000
    ConstraintDescription: must be between 5 and 10000

Resources:
  # API Gateway WebSocket API
  WebSocketApi:
    Type: 'AWS::ApiGatewayV2::Api'
    Properties:
      Name: !Ref AWS::StackName
      Description: An Amazon API Gateway WebSocket API and an AWS Lambda function.
      ProtocolType: WEBSOCKET
      RouteSelectionExpression: "$request.body.action"
  # Lambda Function - uses Globals to define additional configuration values
  OnConnectLambdaFunction:
    Type: 'AWS::Serverless::Function'
    Properties:
      FunctionName: !Sub '${AWS::StackName}-onconnect-function'
      Handler: onConnect.lambda_handler
      MemorySize: 256
      Environment:
        Variables:
          WEBSOCKETS_DDB_TABLE: !Ref WebSocketConnectionsTableName
      Policies:
      - DynamoDBCrudPolicy:
          TableName: !Ref WebSocketConnectionsTableName
  # Function permissions grant an AWS service or another account permission to use a function
  OnConnectFunctionResourcePermission:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      Principal: apigateway.amazonaws.com
      FunctionName: !Ref OnConnectLambdaFunction
      SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*'
  OnConnectIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref WebSocketApi
      Description: OnConnect Integration
      IntegrationType: AWS_PROXY
      IntegrationUri: 
        Fn::Sub:
            arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnConnectLambdaFunction.Arn}/invocations
  OnConnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref WebSocketApi
      RouteKey: $connect
      AuthorizationType: NONE
      ApiKeyRequired: False
      OperationName: OnConnectRoute
      Target: !Join
        - '/'
        - - 'integrations'
          - !Ref OnConnectIntegration
  # Lambda Function - uses Globals to define additional configuration values
  InvokeModelLambdaFunction:
    Type: 'AWS::Serverless::Function'
    Properties:
      FunctionName: !Sub '${AWS::StackName}-invoke-function'
      Handler: invoke.lambda_handler
      MemorySize: 256
      Environment:
        Variables:
          WEBSOCKETS_DDB_TABLE: !Ref WebSocketConnectionsTableName
      Layers:
        - !Ref BedrockBoto3Layer
      Policies:
      - DynamoDBCrudPolicy:
          TableName: !Ref WebSocketConnectionsTableName
      - Statement:
        - Effect: Allow
          Action:
          - 'execute-api:ManageConnections'
          Resource:
          - !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*'
      - Statement:
        - Effect: Allow
          Action:
            - bedrock:InvokeModelWithResponseStream
          Resource: 'arn:aws:bedrock:*::foundation-model/*'
  # Function permissions grant an AWS service or another account permission to use a function
  InvokeModelFunctionResourcePermission:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      Principal: apigateway.amazonaws.com
      FunctionName: !Ref InvokeModelLambdaFunction
      SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*'
  InvokeModelIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref WebSocketApi
      Description: Invoke Model Integration
      IntegrationType: AWS_PROXY
      IntegrationUri: 
        Fn::Sub:
            arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${InvokeModelLambdaFunction.Arn}/invocations
  InvokeModelRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref WebSocketApi
      RouteKey: invokeModel
      AuthorizationType: NONE
      OperationName: InvokeModelRoute
      Target: !Join
        - '/'
        - - 'integrations'
          - !Ref InvokeModelIntegration
  # Lambda Function - uses Globals to define additional configuration values
  OnDisconnectLambdaFunction:
    Type: 'AWS::Serverless::Function'
    Properties:
      FunctionName: !Sub '${AWS::StackName}-ondisconnect-function'
      Handler: onDisconnect.lambda_handler
      MemorySize: 256
      Environment:
        Variables:
          TABLE_NAME: !Ref WebSocketConnectionsTableName
      Policies:
      - DynamoDBCrudPolicy:
          TableName: !Ref WebSocketConnectionsTableName
  # Function permissions grant an AWS service or another account permission to use a function
  OnDisconnectFunctionResourcePermission:
    Type: 'AWS::Lambda::Permission'
    Properties:
      Action: 'lambda:InvokeFunction'
      Principal: apigateway.amazonaws.com
      FunctionName: !Ref OnDisconnectLambdaFunction
      SourceArn: !Sub 'arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${WebSocketApi}/*'
  OnDisconnectIntegration:
    Type: AWS::ApiGatewayV2::Integration
    Properties:
      ApiId: !Ref WebSocketApi
      Description: OnDisconnect Integration
      IntegrationType: AWS_PROXY
      IntegrationUri: 
        Fn::Sub:
            arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${OnDisconnectLambdaFunction.Arn}/invocations
  OnDisconnectRoute:
    Type: AWS::ApiGatewayV2::Route
    Properties:
      ApiId: !Ref WebSocketApi
      RouteKey: $disconnect
      AuthorizationType: NONE
      OperationName: OnDisconnectRoute
      Target: !Join
        - '/'
        - - 'integrations'
          - !Ref OnDisconnectIntegration
  WebSocketConnectionsTableName:
    Type: AWS::DynamoDB::Table
    Properties:
      AttributeDefinitions:
        - AttributeName: connectionId
          AttributeType: S
      KeySchema:
        - AttributeName: connectionId
          KeyType: HASH
      ProvisionedThroughput:
        ReadCapacityUnits: !Ref ReadCapacityUnits
        WriteCapacityUnits: !Ref WriteCapacityUnits
      TableName: 
        !Sub
          - ${AWS::StackName}-${TableNameSuffix}
          - { TableNameSuffix: !Ref WebSocketConnectionsTableNameSuffix }
  Deployment:
    Type: AWS::ApiGatewayV2::Deployment
    DependsOn:
      - OnConnectRoute
      - InvokeModelRoute
      - OnDisconnectRoute
    Properties:
      ApiId: !Ref WebSocketApi
  Stage:
    Type: AWS::ApiGatewayV2::Stage
    Properties:
      StageName: prod
      Description: Prod Stage
      DeploymentId: !Ref Deployment
      ApiId: !Ref WebSocketApi
  # Boto3 version with Bedrock support
  BedrockBoto3Bucket:
    Type: AWS::S3::Bucket
    Properties:
      BucketEncryption:
        ServerSideEncryptionConfiguration: 
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
  BedrockBoto3ZipFunction:
    Type: 'AWS::Serverless::Function'
    Properties:
      FunctionName: !Sub '${AWS::StackName}-boto3-function'
      Handler: getBoto3.lambda_handler
      MemorySize: 256
      Environment:
        Variables:
          BOTO3_BUCKET: !Ref BedrockBoto3Bucket
      Policies:
      - S3CrudPolicy:
          BucketName: !Ref BedrockBoto3Bucket
  BedrockBoto3Zip:
    Type: Custom::BedrockBoto3Zip
    Properties:
      ServiceToken: !GetAtt BedrockBoto3ZipFunction.Arn
      # Rerun BedrockBoto3ZipFunction if any of the following parameters change
      BOTO3_BUCKET: !Ref BedrockBoto3Bucket
      VERSION: 1
  BedrockBoto3Layer:
    Type: "AWS::Lambda::LayerVersion"
    Properties:
      Content:
        S3Bucket: !GetAtt BedrockBoto3Zip.Bucket
        S3Key: !GetAtt BedrockBoto3Zip.Key
      CompatibleRuntimes:
        - python3.10

Outputs:
  OnConnectLambdaFunctionArn:
    Description: "OnConnect function ARN"
    Value: !GetAtt OnConnectLambdaFunction.Arn
  OnDisconnectLambdaFunctionArn:
    Description: "OnDisconnect function ARN"
    Value: !GetAtt OnDisconnectLambdaFunction.Arn
  InvokeModelLambdaFunctionArn:
    Description: "Post function ARN"
    Value: !GetAtt InvokeModelLambdaFunction.Arn
  WebSocketURL:
    Description: "The WSS Protocol URL to connect to"
    Value: !Join [ '', [ 'wss://', !Ref WebSocketApi, '.execute-api.',!Ref 'AWS::Region','.amazonaws.com/',!Ref 'Stage'] ]
  WebSocketConnectionsTableNameArn:
    Description: "WebSocket connections table ARN"
    Value: !GetAtt WebSocketConnectionsTableName.Arn

< Back to all patterns


GitHub icon Visit the GitHub repo for this pattern.

Download

git clone https://github.com/aws-samples/serverless-patterns/ cd serverless-patterns/apigw-websocket-api-bedrock-streaming

Deploy

sam deploy


Testing

See the GitHub repo for detailed testing instructions.

Cleanup

Delete the stack: sam delete --stack-name STACK_NAME.

Presented by Windrich

I am a Solutions Architect in AWS who works with customers in industries such as finance and transport, to help accelerate their cloud adoption journey. He is especially interested in Serverless technologies and how customers can leverage them to bring values to their business.

Follow on LinkedIn