Creates an API Gateway WebSocketand Lambda functions that provides a streaming response from the LLMs in Amazon Bedrock.
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
Visit the GitHub repo for this pattern.
git clone https://github.com/aws-samples/serverless-patterns/ cd serverless-patterns/apigw-websocket-api-bedrock-streaming
sam deploy
sam delete --stack-name STACK_NAME
.