Bedrock Lambda API Gateway Streamer (Rust)

introduction-to-bedrock/bedrock-lambda-apigw-streamer/.gitignore

@@ -0,0 +1,10 @@
+*.js
+!jest.config.js
+*.d.ts
+node_modules
+.DS_Store
+
+# CDK asset staging directory
+.cdk.staging
+cdk.out
+

introduction-to-bedrock/bedrock-lambda-apigw-streamer/README.md

@@ -0,0 +1,82 @@
+# Bedrock Lambda API Gateway Streamer
+
+Bedrock Lambda API Gateway Streamer demonstrates a simple, serverless approach to integrating Amazon Bedrock with AWS Lambda and Amazon API Gateway. Written in Rust 🦀, it showcases how to efficiently stream responses from Amazon Bedrock to a client via WebSocket connections. The example serves as a practical illustration of implementing real-time, serverless communication between Bedrock's GenAI capabilities and end-users through AWS infrastructure.
+
+Once you've successfully built and deployed the application, you'll be able to:
+
+1. Establish a WebSocket connection with Amazon API Gateway.
+2. Use the connection to interact with Amazon Bedrock for simple story generation.
+3. Receive real-time streaming of the generated story content directly to your client via the established WebSocket.
+
+This setup allows for immediate, continuous delivery of the story as it's being created, providing an interactive experience for users.
+
+## Architecture
+
+<p align="center">
+  <img src="images/streaming-response.png" alt="Architecture Diagram"/>
+</p>
+
+## Prerequisites (Local)
+
+* [npm](https://docs.npmjs.com/downloading-and-installing-node-js-and-npm)
+* [Node.js](https://nodejs.org/en/download/package-manager) v18.20.4 or higher
+* [CDK](https://docs.aws.amazon.com/cdk/v2/guide/getting_started.html) latest version
+* [Docker](https://docs.docker.com/engine/install/) is running locally (needed for Rust cross-complilation Lambda build)
+* [Rust 🦀](https://www.rust-lang.org/) v1.79.0 or higher
+* [Cargo Lambda](https://www.cargo-lambda.info/)
+* [cross](https://crates.io/crates/cross) for Rust cross-compilation
+* [wscat](https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-how-to-call-websocket-api-wscat.html) for CLI WebSocket capabilities
+
+## Prerequisites (AWS Accounts)
+
+* Access to Amazon Bedrock
+* Access to the Amazon Bedrock model used in this example - currently `Claude` as listed in the Bedrock console. In you AWS account, make sure you have `Access granted` to this model. Make sure you have access in the same region where you're deploying the application.
+
+<p align="center">
+  <img src="images/amazon-bedrock-models.png" alt="Architecture Diagram"/>
+</p>
+
+
+## Deployment
+
+You can use the following commands at the root of this repository to build and deploy the application:
+
+```bash
+# Everything lives here
+cd cdk
+
+# Install Node.js dependencies
+npm install
+
+# Synthesize the CDK app to produce a cloud assembly for deployment
+cdk synth
+
+# Deploy the CDK stack into your AWS environment
+cdk deploy
+```
+After you run the `cdk deploy` command, under the `Outputs` section at the end, you'll see the `BedrockStreamerStack.WebSocketURL` followed by the WebSocket URL you'll use to connect to API Gateway. It should look something like `wss://{YOUR_API_ID_HERE}.execute-api.{YOUR_REGION_HERE}.amazonaws.com/prod`
+
+## Run It
+
+Now that you've sucessfully built and deployed the application, it's time to run it! You can use `wscat` to connect via WebSocket to API Gateway and send it commands. The commands you send will allow a Bedrock LLM to generate a short story based on the type of story you'd like to produce.
+
+```bash
+# Connect to the API Gateway URL via WebSocket
+wscat -c wss://{YOUR_API_ID_HERE}.execute-api.{YOUR_REGION_HERE}.amazonaws.com/prod
+
+Connected (press CTRL+C to quit)
+> {"storyType": "FICTION"}
+< {"type":"completion","completion":" Here","stop_reason":null,"stop":null}
+< {"type":"completion","completion":" is","stop_reason":null,"stop":null}
+< {"type":"completion","completion":" a","stop_reason":null,"stop":null}
+< {"type":"completion","completion":" very","stop_reason":null,"stop":null}
+< {"type":"completion","completion":" short","stop_reason":null,"stop":null}
+< {"type":"completion","completion":" fictional","stop_reason":null,"stop":null}
+< {"type":"completion","completion":" story","stop_reason":null,"stop":null}
+< {"type":"completion","completion":":","stop_reason":null,"stop":null}
+.
+.
+.
+< {"type":"completion","completion":"","stop_reason":"stop_sequence","stop":"\n\nHuman:"}
+```
+As the `wscat` CLI says, press `CTRL+C` to disconnect 

introduction-to-bedrock/bedrock-lambda-apigw-streamer/cdk/.npmignore

@@ -0,0 +1,6 @@
+*.ts
+!*.d.ts
+
+# CDK asset staging directory
+.cdk.staging
+cdk.out

introduction-to-bedrock/bedrock-lambda-apigw-streamer/cdk/bin/bedrock-streamer.ts

@@ -0,0 +1,21 @@
+#!/usr/bin/env node
+import 'source-map-support/register';
+import * as cdk from 'aws-cdk-lib';
+import { BedrockStreamerStack } from '../lib/bedrock-streamer-stack';
+
+const app = new cdk.App();
+new BedrockStreamerStack(app, 'BedrockStreamerStack', {
+  /* If you don't specify 'env', this stack will be environment-agnostic.
+   * Account/Region-dependent features and context lookups will not work,
+   * but a single synthesized template can be deployed anywhere. */
+
+  /* Uncomment the next line to specialize this stack for the AWS Account
+   * and Region that are implied by the current CLI configuration. */
+  // env: { account: process.env.CDK_DEFAULT_ACCOUNT, region: process.env.CDK_DEFAULT_REGION },
+
+  /* Uncomment the next line if you know exactly what Account and Region you
+   * want to deploy the stack to. */
+  // env: { account: '123456789012', region: 'us-east-1' },
+
+  /* For more information, see https://docs.aws.amazon.com/cdk/latest/guide/environments.html */
+});

introduction-to-bedrock/bedrock-lambda-apigw-streamer/cdk/cdk.json

@@ -0,0 +1,71 @@
+{
+  "app": "npx ts-node --prefer-ts-exts bin/bedrock-streamer.ts",
+  "watch": {
+    "include": [
+      "**"
+    ],
+    "exclude": [
+      "README.md",
+      "cdk*.json",
+      "**/*.d.ts",
+      "**/*.js",
+      "tsconfig.json",
+      "package*.json",
+      "yarn.lock",
+      "node_modules",
+      "test"
+    ]
+  },
+  "context": {
+    "@aws-cdk/aws-lambda:recognizeLayerVersion": true,
+    "@aws-cdk/core:checkSecretUsage": true,
+    "@aws-cdk/core:target-partitions": [
+      "aws",
+      "aws-cn"
+    ],
+    "@aws-cdk-containers/ecs-service-extensions:enableDefaultLogDriver": true,
+    "@aws-cdk/aws-ec2:uniqueImdsv2TemplateName": true,
+    "@aws-cdk/aws-ecs:arnFormatIncludesClusterName": true,
+    "@aws-cdk/aws-iam:minimizePolicies": true,
+    "@aws-cdk/core:validateSnapshotRemovalPolicy": true,
+    "@aws-cdk/aws-codepipeline:crossAccountKeyAliasStackSafeResourceName": true,
+    "@aws-cdk/aws-s3:createDefaultLoggingPolicy": true,
+    "@aws-cdk/aws-sns-subscriptions:restrictSqsDescryption": true,
+    "@aws-cdk/aws-apigateway:disableCloudWatchRole": true,
+    "@aws-cdk/core:enablePartitionLiterals": true,
+    "@aws-cdk/aws-events:eventsTargetQueueSameAccount": true,
+    "@aws-cdk/aws-ecs:disableExplicitDeploymentControllerForCircuitBreaker": true,
+    "@aws-cdk/aws-iam:importedRoleStackSafeDefaultPolicyName": true,
+    "@aws-cdk/aws-s3:serverAccessLogsUseBucketPolicy": true,
+    "@aws-cdk/aws-route53-patters:useCertificate": true,
+    "@aws-cdk/customresources:installLatestAwsSdkDefault": false,
+    "@aws-cdk/aws-rds:databaseProxyUniqueResourceName": true,
+    "@aws-cdk/aws-codedeploy:removeAlarmsFromDeploymentGroup": true,
+    "@aws-cdk/aws-apigateway:authorizerChangeDeploymentLogicalId": true,
+    "@aws-cdk/aws-ec2:launchTemplateDefaultUserData": true,
+    "@aws-cdk/aws-secretsmanager:useAttachedSecretResourcePolicyForSecretTargetAttachments": true,
+    "@aws-cdk/aws-redshift:columnId": true,
+    "@aws-cdk/aws-stepfunctions-tasks:enableEmrServicePolicyV2": true,
+    "@aws-cdk/aws-ec2:restrictDefaultSecurityGroup": true,
+    "@aws-cdk/aws-apigateway:requestValidatorUniqueId": true,
+    "@aws-cdk/aws-kms:aliasNameRef": true,
+    "@aws-cdk/aws-autoscaling:generateLaunchTemplateInsteadOfLaunchConfig": true,
+    "@aws-cdk/core:includePrefixInUniqueNameGeneration": true,
+    "@aws-cdk/aws-efs:denyAnonymousAccess": true,
+    "@aws-cdk/aws-opensearchservice:enableOpensearchMultiAzWithStandby": true,
+    "@aws-cdk/aws-lambda-nodejs:useLatestRuntimeVersion": true,
+    "@aws-cdk/aws-efs:mountTargetOrderInsensitiveLogicalId": true,
+    "@aws-cdk/aws-rds:auroraClusterChangeScopeOfInstanceParameterGroupWithEachParameters": true,
+    "@aws-cdk/aws-appsync:useArnForSourceApiAssociationIdentifier": true,
+    "@aws-cdk/aws-rds:preventRenderingDeprecatedCredentials": true,
+    "@aws-cdk/aws-codepipeline-actions:useNewDefaultBranchForCodeCommitSource": true,
+    "@aws-cdk/aws-cloudwatch-actions:changeLambdaPermissionLogicalIdForLambdaAction": true,
+    "@aws-cdk/aws-codepipeline:crossAccountKeysDefaultValueToFalse": true,
+    "@aws-cdk/aws-codepipeline:defaultPipelineTypeToV2": true,
+    "@aws-cdk/aws-kms:reduceCrossAccountRegionPolicyScope": true,
+    "@aws-cdk/aws-eks:nodegroupNameAttribute": true,
+    "@aws-cdk/aws-ec2:ebsDefaultGp3Volume": true,
+    "@aws-cdk/aws-ecs:removeDefaultDeploymentAlarm": true,
+    "@aws-cdk/custom-resources:logApiResponseDataPropertyTrueDefault": false
+  }
+}

introduction-to-bedrock/bedrock-lambda-apigw-streamer/cdk/jest.config.js

@@ -0,0 +1,8 @@
+module.exports = {
+  testEnvironment: 'node',
+  roots: ['<rootDir>/test'],
+  testMatch: ['**/*.test.ts'],
+  transform: {
+    '^.+\\.tsx?$': 'ts-jest'
+  }
+};

introduction-to-bedrock/bedrock-lambda-apigw-streamer/cdk/lib/bedrock-streamer-stack.ts

@@ -0,0 +1,69 @@
+import * as cdk from 'aws-cdk-lib';
+import * as apigatewayv2 from 'aws-cdk-lib/aws-apigatewayv2';
+import * as lambda from 'aws-cdk-lib/aws-lambda';
+import * as iam from 'aws-cdk-lib/aws-iam';
+import { WebSocketLambdaIntegration } from 'aws-cdk-lib/aws-apigatewayv2-integrations';
+import { Construct } from 'constructs';
+import { RustFunction } from 'cargo-lambda-cdk';
+
+export class BedrockStreamerStack extends cdk.Stack {
+  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
+    super(scope, id, props);
+
+    // Rust Lambda Functions
+    const streamingLambda = new RustFunction(this, 'BedrockStreamer', {
+      manifestPath: 'lib/resources/lambdas/bedrock-streamer/Cargo.toml',
+      bundling: {
+        architecture: lambda.Architecture.ARM_64,
+        cargoLambdaFlags: [
+          '--compiler',
+          'cross',
+          '--release',
+        ],
+      },
+      timeout: cdk.Duration.seconds(60),
+    });
+
+    // Grant Lambda permission to invoke Bedrock
+    streamingLambda.addToRolePolicy(new iam.PolicyStatement({
+      actions: ['bedrock:InvokeModelWithResponseStream'],
+      resources: ['*'],
+    }));
+
+    // Create the WebSocket API
+    const webSocketApi = new apigatewayv2.WebSocketApi(this, 'WebSocketApi', {
+      connectRouteOptions: {
+        integration: new WebSocketLambdaIntegration('ConnectIntegration', streamingLambda),
+      },
+      disconnectRouteOptions: {
+        integration: new WebSocketLambdaIntegration('DisconnectIntegration', streamingLambda),
+      },
+      defaultRouteOptions: {
+        integration: new WebSocketLambdaIntegration('DefaultIntegration', streamingLambda),
+      },
+    });
+
+    // Create a stage for the WebSocket API
+    const stage = new apigatewayv2.WebSocketStage(this, 'Stage', {
+      webSocketApi,
+      stageName: 'prod',
+      autoDeploy: true,
+    });
+
+    // Grant the Lambda function permission to manage WebSocket connections
+    streamingLambda.addToRolePolicy(new iam.PolicyStatement({
+      actions: [
+        'execute-api:ManageConnections',
+      ],
+      resources: [
+        `arn:aws:execute-api:${this.region}:${this.account}:${webSocketApi.apiId}/${stage.stageName}/*`,
+      ],
+    }));
+
+    // Output the WebSocket URL
+    new cdk.CfnOutput(this, 'WebSocketURL', {
+      value: stage.url,
+      description: 'WebSocket URL',
+    });
+  }
+}

introduction-to-bedrock/bedrock-lambda-apigw-streamer/cdk/lib/resources/lambdas/bedrock-streamer/.gitignore

@@ -0,0 +1 @@
+/target