feat(construct): Bug fix for aws-contentgen-appsync-lambda and aws-su…

…mmarization-appsync-stepfn construct (#722) * feat(content-gen): update modelid in content gen construct --------- Co-authored-by: Dinesh Sajwan <[email protected]> Co-authored-by: Alain Krok <[email protected]>
awslabs · Oct 7, 2024 · 9b1b838 · 9b1b838
1 parent 82f48c1
commit 9b1b838
Show file tree

Hide file tree

Showing 8 changed files with 75 additions and 54 deletions.
diff --git a/README.md b/README.md
@@ -136,7 +136,7 @@ The following constructs are available in the library:
 | [SageMaker model deployment (JumpStart)](./src/patterns/gen-ai/aws-model-deployment-sagemaker/README_jumpstart.md) | Deploy a foundation model from Amazon SageMaker JumpStart to an Amazon SageMaker endpoint. | Amazon SageMaker |
 | [SageMaker model deployment (Hugging Face)](./src/patterns/gen-ai/aws-model-deployment-sagemaker/README_hugging_face.md) | Deploy a foundation model from Hugging Face to an Amazon SageMaker endpoint. | Amazon SageMaker |
 | [SageMaker model deployment (Custom)](./src/patterns/gen-ai/aws-model-deployment-sagemaker/README_custom_sagemaker_endpoint.md) | Deploy a foundation model from an S3 location to an Amazon SageMaker endpoint. | Amazon SageMaker |
-| [Content Generation](./src/patterns/gen-ai/aws-contentgen-appsync-lambda/README.md) | Generate images from text using Amazon titan-image-generator-v1 or stability.stable-diffusion-xl model.  | AWS Lambda, Amazon Bedrock, AWS AppSync |
+| [Content Generation](./src/patterns/gen-ai/aws-contentgen-appsync-lambda/README.md) | Generate images from text using Amazon titan-image-generator-v1 or stability.stable-diffusion-xl-v1 model.  | AWS Lambda, Amazon Bedrock, AWS AppSync |
 | [Web crawler](./src/patterns/gen-ai/aws-web-crawler/README.md) | Crawl websites and RSS feeds on a schedule and store changeset data in an Amazon Simple Storage Service bucket.  | AWS Lambda, AWS Batch, AWS Fargate, Amazon DynamoDB  |
 | [Amazon Bedrock Monitoring (Amazon CloudWatch Dashboard)](./src/patterns/gen-ai/aws-bedrock-cw-dashboard/README.md) | Amazon CloudWatch dashboard to monitor model usage from Amazon Bedrock. | Amazon CloudWatch |
 

diff --git a/lambda/aws-contentgen-appsync-lambda/src/image_generator.py b/lambda/aws-contentgen-appsync-lambda/src/image_generator.py
@@ -15,6 +15,8 @@
 from datetime import datetime
 from requests_aws4auth import AWS4Auth
 from aws_lambda_powertools import Logger, Tracer, Metrics
+from util  import MODEL_NAME
+
 
 logger = Logger(service="CONTENT_GENERATION")
 tracer = Tracer(service="CONTENT_GENERATION")
@@ -49,7 +51,6 @@ def __init__(self,input_text, rekognition_client,comprehend_client,bedrock_clien
 
 
 
-    @tracer.capture_method
     def upload_file_to_s3(self,imgbase64encoded,file_name):
 
         """Upload generated file to S3 bucket"""
@@ -68,7 +69,6 @@ def upload_file_to_s3(self,imgbase64encoded,file_name):
             "bucket_name":self.bucket,
         }
 
-    @tracer.capture_method
     def text_moderation(self):
 
         """Check input text has any toxicity or not. The comprehend is trained 
@@ -96,7 +96,6 @@ def text_moderation(self):
 
         return response
 
-    @tracer.capture_method
     def image_moderation(self,file_name):
 
         """Detect image moderation on the generated image to avoid any toxicity/nudity"""
@@ -197,12 +196,12 @@ def send_job_status(self,variables):
             auth=aws_auth_appsync,
             timeout=10
         )
-        logger.info('res :: {}',responseJobstatus)
+        logger.info(f"sending response :: {responseJobstatus}")
 
 def get_model_payload(modelid,params,input_text,negative_prompts):
-      
+
      body=''
-     if modelid=='stability.stable-diffusion-xl' :
+     if modelid==MODEL_NAME.STABILITY_DIFFUSION :
         body = json.dumps({
                 "text_prompts": (
                         [{"text": input_text, "weight": 1.0}]
@@ -218,7 +217,7 @@ def get_model_payload(modelid,params,input_text,negative_prompts):
                 "height": params['height']
                 })
         return body
-     if modelid=='amazon.titan-image-generator-v1' :
+     if modelid==MODEL_NAME.TITAN_IMAGE :
 
         body = json.dumps({
                        "taskType": "TEXT_IMAGE",

diff --git a/lambda/aws-contentgen-appsync-lambda/src/lambda.py b/lambda/aws-contentgen-appsync-lambda/src/lambda.py
@@ -20,7 +20,7 @@
 from aws_lambda_powertools.utilities.typing import LambdaContext
 from aws_lambda_powertools.metrics import MetricUnit
 from aws_lambda_powertools.utilities.validation import validate, SchemaValidationError
-
+from util  import MODEL_NAME
 
 
 logger = Logger(service="CONTENT_GENERATION")
@@ -88,9 +88,9 @@ def handler(event,  context: LambdaContext) -> dict:
             num_of_images=0 #if multiple image geneated iterate through all
             for image in parsed_reponse['image_generated']:
                 logger.info(f'num_of_images {num_of_images}')
-                if model_id=='stability.stable-diffusion-xl' :
+                if model_id==MODEL_NAME.STABILITY_DIFFUSION :
                     imgbase64encoded= parsed_reponse['image_generated'][num_of_images]["base64"]
-                if model_id=='amazon.titan-image-generator-v1' :
+                if model_id==MODEL_NAME.TITAN_IMAGE :
                     imgbase64encoded= parsed_reponse['image_generated'][num_of_images]
                 imageGenerated=img.upload_file_to_s3(imgbase64encoded,file_name)
                 num_of_images=+1
@@ -127,14 +127,14 @@ def parse_response(query_response,model_id):
     else:
         response_dict = json.loads(query_response["body"].read())
 
-        if model_id=='stability.stable-diffusion-xl' :
+        if model_id==MODEL_NAME.STABILITY_DIFFUSION :
 
             if(response_dict['artifacts'] is None):
                 parsed_reponse['image_generated_status']='Failed'
             else:
                  parsed_reponse['image_generated']=response_dict['artifacts']
 
-        if model_id=='amazon.titan-image-generator-v1' :
+        if model_id==MODEL_NAME.TITAN_IMAGE :
             if(response_dict['images'] is None):
                 parsed_reponse['image_generated_status']='Failed'
             else:
@@ -143,4 +143,3 @@ def parse_response(query_response,model_id):
 
         return parsed_reponse
 
-
diff --git a/lambda/aws-contentgen-appsync-lambda/src/util.py b/lambda/aws-contentgen-appsync-lambda/src/util.py
@@ -0,0 +1,17 @@
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+from enum import StrEnum
+
+class MODEL_NAME(StrEnum):
+    STABILITY_DIFFUSION = 'stability.stable-diffusion-xl-v1',
+    TITAN_IMAGE='amazon.titan-image-generator-v1'
diff --git a/src/patterns/gen-ai/aws-contentgen-appsync-lambda/README.md b/src/patterns/gen-ai/aws-contentgen-appsync-lambda/README.md
@@ -43,7 +43,7 @@ The workflow is as follows:
 
 3. Lambda function first implement text moderation using Amazon Comprehend to check for inappropriate content.
 
-4. The functions then generate an image from the text using Amazon Bedrock with the stability.stable-diffusion-xl/amazon.titan-image-generator-v1 model.
+4. The functions then generate an image from the text using Amazon Bedrock with the stability.stable-diffusion-xl-v1/amazon.titan-image-generator-v1 model.
 
 5. Next, image moderation is performed using Amazon Rekognition to further ensure appropriateness.
 
@@ -52,7 +52,7 @@ The workflow is as follows:
 
 This construct builds a Lambda function from a Docker image, thus you need [Docker desktop](https://www.docker.com/products/docker-desktop/) running on your machine.
 
-Make sure the model (stability.stable-diffusion-xl/amazon.titan-image-generator-v1) is enabled in your account. Please follow the [Amazon Bedrock User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html) for steps related to enabling model access.
+Make sure the model (stability.stable-diffusion-xl-v1/amazon.titan-image-generator-v1) is enabled in your account. Please follow the [Amazon Bedrock User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html) for steps related to enabling model access.
 
 AWS Lambda functions provisioned in this construct use [Powertools for AWS Lambda (Python)](https://github.com/aws-powertools/powertools-lambda-python) for tracing, structured logging and custom metrics creation.
 
@@ -214,7 +214,7 @@ Expected response: It invoke an asynchronous summarization process thus the resp
 Where:
 - job_id: id which can be used to filter subscriptions on client side.
 - status: this field will be used by the subscription to update the status of the image generation process.
-- model_config: configure model id amazon.titan-image-generator-v1/stability.stable-diffusion-xl.
+- model_config: configure model id amazon.titan-image-generator-v1/stability.stable-diffusion-xl-v1.
 - model_kwargs: Image generation model driver for Stable Diffusion models and Amazon Titan generator on Amazon Bedrock.
 
 

diff --git a/src/patterns/gen-ai/aws-contentgen-appsync-lambda/index.ts b/src/patterns/gen-ai/aws-contentgen-appsync-lambda/index.ts
@@ -177,13 +177,11 @@ export class ContentGenerationAppSyncLambda extends BaseClass {
     if (props?.existingVpc) {
       this.vpc = props.existingVpc;
     } else {
-      this.vpc = vpc_helper.buildVpc(scope, {
-        defaultVpcProps: props?.vpcProps,
-        vpcName: 'cgAppSyncLambdaVpc',
-      });
+      this.vpc = new ec2.Vpc(this, 'Vpc', props.vpcProps);
       // vpc endpoints
       vpc_helper.AddAwsServiceEndpoint(scope, this.vpc, [vpc_helper.ServiceEndpointTypeEnum.S3,
-        vpc_helper.ServiceEndpointTypeEnum.BEDROCK_RUNTIME, vpc_helper.ServiceEndpointTypeEnum.REKOGNITION]);
+        vpc_helper.ServiceEndpointTypeEnum.BEDROCK_RUNTIME, vpc_helper.ServiceEndpointTypeEnum.REKOGNITION,
+        vpc_helper.ServiceEndpointTypeEnum.COMPREHEND]);
     }
 
     // Security group
@@ -285,6 +283,7 @@ export class ContentGenerationAppSyncLambda extends BaseClass {
           ],
         },
         xrayEnabled: this.enablexray,
+        visibility: appsync.Visibility.GLOBAL,
         logConfig: {
           fieldLogLevel: this.fieldLogLevel,
           retention: this.retention,
@@ -471,7 +470,7 @@ export class ContentGenerationAppSyncLambda extends BaseClass {
       description: 'Lambda function for generating image',
       vpc: this.vpc,
       tracing: this.lambdaTracing,
-      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_ISOLATED },
+      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
       securityGroups: [this.securityGroup],
       memorySize: lambdaMemorySizeLimiter(this, 1_769 * 4),
       timeout: Duration.minutes(15),

diff --git a/src/patterns/gen-ai/aws-summarization-appsync-stepfn/index.ts b/src/patterns/gen-ai/aws-summarization-appsync-stepfn/index.ts
@@ -11,7 +11,7 @@
  *  and limitations under the License.
  */
 import * as path from 'path';
-import { Duration, Aws } from 'aws-cdk-lib';
+import { Duration, Aws, RemovalPolicy } from 'aws-cdk-lib';
 import * as appsync from 'aws-cdk-lib/aws-appsync';
 import * as cognito from 'aws-cdk-lib/aws-cognito';
 import * as ec2 from 'aws-cdk-lib/aws-ec2';
@@ -258,15 +258,10 @@ export class SummarizationAppsyncStepfn extends BaseClass {
     if (props?.existingVpc) {
       this.vpc = props.existingVpc;
     } else {
-      this.vpc = vpc_helper.buildVpc(scope, {
-        defaultVpcProps: props?.vpcProps,
-        vpcName: 'sumAppSyncStepFnVpc',
-      });
-
+      this.vpc = new ec2.Vpc(this, 'Vpc', props.vpcProps);
       // vpc endpoints
       vpc_helper.AddAwsServiceEndpoint(scope, this.vpc, [vpc_helper.ServiceEndpointTypeEnum.S3,
-        vpc_helper.ServiceEndpointTypeEnum.BEDROCK_RUNTIME, vpc_helper.ServiceEndpointTypeEnum.REKOGNITION,
-        vpc_helper.ServiceEndpointTypeEnum.APP_SYNC]);
+        vpc_helper.ServiceEndpointTypeEnum.BEDROCK_RUNTIME, vpc_helper.ServiceEndpointTypeEnum.REKOGNITION]);
     }
 
     // Security group
@@ -303,6 +298,7 @@ export class SummarizationAppsyncStepfn extends BaseClass {
         encryption: s3.BucketEncryption.S3_MANAGED,
         enforceSSL: true,
         versioned: true,
+        removalPolicy: RemovalPolicy.DESTROY,
         lifecycleRules: [{
           expiration: Duration.days(90),
         }],
@@ -321,17 +317,15 @@ export class SummarizationAppsyncStepfn extends BaseClass {
       this.inputAssetBucket = new s3.Bucket(this,
         'inputAssetsSummaryBucket' + this.stage, props.bucketInputsAssetsProps);
     } else {
-      const bucketName = generatePhysicalNameV2(this,
-        'input-assets-bucket' + this.stage,
-        { maxLength: 63, lower: true });
-      this.inputAssetBucket = new s3.Bucket(this, bucketName,
+
+      this.inputAssetBucket = new s3.Bucket(this, 'inputAssetsSummaryBucket' + this.stage,
         {
           blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
           encryption: s3.BucketEncryption.S3_MANAGED,
-          bucketName: bucketName,
           serverAccessLogsBucket: serverAccessLogBucket,
           enforceSSL: true,
           versioned: true,
+          removalPolicy: RemovalPolicy.DESTROY,
           lifecycleRules: [{
             expiration: Duration.days(90),
           }],
@@ -350,18 +344,14 @@ export class SummarizationAppsyncStepfn extends BaseClass {
       this.processedAssetBucket = new s3.Bucket(this,
         'processedAssetsSummaryBucket' + this.stage, props.bucketProcessedAssetsProps);
     } else {
-      const bucketName = generatePhysicalNameV2(this,
-        'processed-assets-bucket' + this.stage,
-        { maxLength: 63, lower: true });
-
-      this.processedAssetBucket = new s3.Bucket(this, bucketName,
+      this.processedAssetBucket = new s3.Bucket(this, 'processedAssetsSummaryBucket' + this.stage,
         {
           blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
           encryption: s3.BucketEncryption.S3_MANAGED,
-          bucketName: bucketName,
           serverAccessLogsBucket: serverAccessLogBucket,
           enforceSSL: true,
           versioned: true,
+          removalPolicy: RemovalPolicy.DESTROY,
           lifecycleRules: [{
             expiration: Duration.days(90),
           }],
@@ -495,7 +485,7 @@ export class SummarizationAppsyncStepfn extends BaseClass {
       description: 'Lambda function to validate input for summary api',
       vpc: this.vpc,
       tracing: this.lambdaTracing,
-      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_ISOLATED },
+      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
       securityGroups: [this.securityGroup],
       memorySize: lambdaMemorySizeLimiter(this, 1_769 * 1),
       timeout: Duration.minutes(5),
@@ -592,7 +582,7 @@ export class SummarizationAppsyncStepfn extends BaseClass {
       functionName: 'summary_document_reader' + this.stage,
       description: 'Lambda function to read the input transformed document',
       vpc: this.vpc,
-      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_ISOLATED },
+      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
       securityGroups: [this.securityGroup],
       memorySize: lambdaMemorySizeLimiter(this, 1_769 * 1),
       tracing: this.lambdaTracing,
@@ -695,7 +685,7 @@ export class SummarizationAppsyncStepfn extends BaseClass {
       description: 'Lambda function to generate the summary',
       code: lambda.DockerImageCode.fromImageAsset(path.join(__dirname, '../../../../lambda/aws-summarization-appsync-stepfn/summary_generator')),
       vpc: this.vpc,
-      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_ISOLATED },
+      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS },
       securityGroups: [this.securityGroup],
       memorySize: lambdaMemorySizeLimiter(this, 1_769 * 4),
       timeout: Duration.minutes(10),
@@ -809,9 +799,7 @@ export class SummarizationAppsyncStepfn extends BaseClass {
 
     const logGroupName = generatePhysicalNameV2(this, logGroupPrefix,
       { maxLength: maxGeneratedNameLength, lower: true });
-    const summarizationLogGroup = new logs.LogGroup(this, 'summarizationLogGroup', {
-      logGroupName: logGroupName,
-    });
+
 
     // step function definition
     const definition = inputValidationTask.next(
@@ -824,12 +812,11 @@ export class SummarizationAppsyncStepfn extends BaseClass {
     );
 
     // step function
-
     const summarizationStepFunction = new sfn.StateMachine(this, 'summarizationStepFunction', {
       definitionBody: sfn.DefinitionBody.fromChainable(definition),
       timeout: Duration.minutes(15),
       logs: {
-        destination: summarizationLogGroup,
+        destination: getLoggroup(this, logGroupName),
         level: sfn.LogLevel.ALL,
       },
       tracingEnabled: this.enablexray,
@@ -888,3 +875,17 @@ export class SummarizationAppsyncStepfn extends BaseClass {
   }
 }
 
+function getLoggroup(stack: Construct, logGroupName: string) {
+  const existingLogGroup = logs.LogGroup.fromLogGroupName(
+    stack, 'ExistingSummarizationLogGroup', logGroupName);
+
+  if (existingLogGroup.logGroupName) {
+    return existingLogGroup;
+  } else {
+    return new logs.LogGroup(stack, 'SummarizationLogGroup', {
+      logGroupName: logGroupName,
+      retention: logs.RetentionDays.ONE_MONTH,
+      removalPolicy: RemovalPolicy.DESTROY,
+    });
+  }
+}
diff --git a/...patterns/gen-ai/aws-summarization-appsync-stepfn/aws-summarization-appsync-stepfn.test.ts b/...patterns/gen-ai/aws-summarization-appsync-stepfn/aws-summarization-appsync-stepfn.test.ts
@@ -56,11 +56,17 @@ describe('Summarization Appsync Stepfn construct', () => {
             cidrMask: 24,
           },
           {
-            name: 'private',
+            name: 'isolated',
             subnetType: ec2.SubnetType.PRIVATE_ISOLATED,
             cidrMask: 24,
           },
+          {
+            name: 'private',
+            subnetType: ec2.SubnetType.PRIVATE_WITH_EGRESS,
+            cidrMask: 24,
+          },
         ],
+        natGateways: 1,
       },
     );
     const mergedapiRole = new iam.Role(
@@ -129,9 +135,9 @@ describe('Summarization Appsync Stepfn construct', () => {
               'GraphQLUrl',
             ],
           },
-          INPUT_ASSET_BUCKET: { Ref: Match.stringLikeRegexp('testinputassetsbucket') },
+          INPUT_ASSET_BUCKET: { Ref: Match.stringLikeRegexp('testinputAssetsSummaryBucket') },
           IS_FILE_TRANSFORMED: 'false',
-          TRANSFORMED_ASSET_BUCKET: { Ref: Match.stringLikeRegexp('testprocessedassetsbucket') },
+          TRANSFORMED_ASSET_BUCKET: { Ref: Match.stringLikeRegexp('testprocessedAssetsSummaryBucket') },
         },
       },
     });
@@ -142,7 +148,7 @@ describe('Summarization Appsync Stepfn construct', () => {
         Variables: {
           ASSET_BUCKET_NAME: {
             Ref: Match.stringLikeRegexp
-            ('testprocessedassetsbucket'),
+            ('testprocessedAssetsSummaryBucket'),
           },
           GRAPHQL_URL: {
             'Fn::GetAtt': [