Merge pull request #13 from aws-samples/aurora

Aurora Support

README.md

@@ -94,7 +94,7 @@ new LargeLanguageModel(this, 'ModelId', {
   model: {
     kind: ModelKind.CustomScript,
     modelId: 'modelId', // i.e. sentence-transformers/all-MiniLM-L6-v2 - this must match HuggingFace Model ID
-    codeFolder: 'localFolder', // see for example ./lib/semantic-search/embeddings-model
+    codeFolder: 'localFolder', // see for example ./lib/aurora-semantic-search/embeddings-model
     container: 'container-arn', // One from https://github.com/aws/deep-learning-containers/blob/master/available_images.md
     instanceType: 'instanceType', // i.e. g5.12xlarge
     codeBuildComputeType: codebuild.ComputeType.LARGE, // Size of CodeBuild instance. Must have enough storage to download the whole model repository from HuggingFace

bin/aws-genai-llm-chatbot.ts

@@ -6,7 +6,7 @@ import * as lambda from 'aws-cdk-lib/aws-lambda';
 import { ChatBotStack } from '../lib/chatbot-stack';
 import { ChatBotUIStack } from '../lib/chatbot-ui-stack';
 import { ChatBotVpcStack } from '../lib/chatbot-vpc-stack';
-import { SemanticSearchStack } from '../lib/semantic-search/semantic-search-stack';
+import { AuroraSemanticSearchStack } from '../lib/aurora-semantic-search/aurora-semantic-search-stack';
 
 const config = {
   deployUI: true,
@@ -16,24 +16,38 @@ const config = {
 };
 
 const app = new cdk.App();
-const chatBotVpcStack = new ChatBotVpcStack(app, `${config.prefix}-ChatBotVpcStack`);
+const chatBotVpcStack = new ChatBotVpcStack(
+  app,
+  `${config.prefix}-ChatBotVpcStack`
+);
 
 let semanticSearchApi: lambda.Function | null = null;
 if (config.deploySemanticSearch) {
-  const semanticSearch = new SemanticSearchStack(app, `${config.prefix}-SemanticSearchStack`, {
-    vpc: chatBotVpcStack.vpc,
-  });
+  const semanticSearch = new AuroraSemanticSearchStack(
+    app,
+    `${config.prefix}-AuroraSemanticSearchStack`,
+    {
+      vpc: chatBotVpcStack.vpc,
+    }
+  );
 
   semanticSearchApi = semanticSearch.semanticSearchApi;
 }
 
 const chatBotStack = new ChatBotStack(app, `${config.prefix}-ChatBotStack`, {
+  prefix: config.prefix,
   vpc: chatBotVpcStack.vpc,
   semanticSearchApi,
   maxParallelLLMQueries: config.maxParallelLLMQueries,
 });
 
 if (config.deployUI) {
-  const chatBotUIStack = new ChatBotUIStack(app, `${config.prefix}-ChatBotUIStack`);
+  const chatBotUIStack = new ChatBotUIStack(
+    app,
+    `${config.prefix}-ChatBotUIStack`,
+    {
+      prefix: config.prefix,
+    }
+  );
   chatBotUIStack.addDependency(chatBotStack);
 }

lib/semantic-search/semantic-search-stack.ts → lib/aurora-semantic-search/aurora-semantic-search-stack.ts

@@ -14,13 +14,13 @@ import { Construct } from 'constructs';
 import { DocumentIndexingStack } from './document-indexing-stack';
 import { HuggingFaceCustomScriptModel } from '../large-language-model/hf-custom-script-model';
 
-export enum SemanticSearchIndexType {
+export enum PGVectorIndexType {
   COSINE = 'cosine',
   L2 = 'l2',
   INNER = 'inner',
 }
 
-export interface SemanticSearchStackProps extends cdk.StackProps {
+export interface AuroraSemanticSearchStackProps extends cdk.StackProps {
   vpc: ec2.Vpc;
   /**
   https://github.com/pgvector/pgvector
@@ -31,51 +31,47 @@ export interface SemanticSearchStackProps extends cdk.StackProps {
   You can add an index to use approximate nearest neighbor search, which trades some recall for performance. 
   Unlike typical indexes, you will see different results for queries after adding an approximate index.
   **/
-  indexTypes?: SemanticSearchIndexType[];
+  indexTypes?: PGVectorIndexType[];
 }
 
-export class SemanticSearchStack extends cdk.Stack {
+export class AuroraSemanticSearchStack extends cdk.Stack {
   public semanticSearchApi: lambda.DockerImageFunction;
 
-  constructor(scope: Construct, id: string, props: SemanticSearchStackProps) {
+  constructor(
+    scope: Construct,
+    id: string,
+    props: AuroraSemanticSearchStackProps
+  ) {
     super(scope, id, props);
 
     const { vpc, indexTypes = [] } = props;
 
-    const { dbInstance } = this.createVectorDB({ vpc, indexTypes });
+    const { dbCluster } = this.createVectorDB({ vpc, indexTypes });
     const { embeddingsEndpoint } = this.createEmbeddingsEndpoint({ vpc });
     new DocumentIndexingStack(this, 'DocumentIndexingStack', {
       vpc,
-      dbInstance,
+      dbCluster,
       embeddingsEndpoint,
     });
 
-    this.createAPI({ vpc, dbInstance, embeddingsEndpoint: embeddingsEndpoint });
+    this.createAPI({ vpc, dbCluster, embeddingsEndpoint: embeddingsEndpoint });
   }
 
   private createVectorDB({
     vpc,
     indexTypes,
   }: {
     vpc: ec2.Vpc;
-    indexTypes: SemanticSearchIndexType[];
+    indexTypes: PGVectorIndexType[];
   }) {
-    const dbInstance = new rds.DatabaseInstance(this, 'DatabaseInstance', {
-      vpc: vpc,
-      vpcSubnets: {
-        subnetType: ec2.SubnetType.PRIVATE_ISOLATED,
-      },
-      storageEncrypted: true,
-      securityGroups: [
-        new ec2.SecurityGroup(this, 'DatabaseSecurityGroup', {
-          vpc: vpc,
-          allowAllOutbound: true,
-        }),
-      ],
-      engine: rds.DatabaseInstanceEngine.postgres({
-        version: rds.PostgresEngineVersion.VER_15_3,
+    const dbCluster = new rds.DatabaseCluster(this, 'AuroraDatabase', {
+      engine: rds.DatabaseClusterEngine.auroraPostgres({
+        version: rds.AuroraPostgresEngineVersion.VER_15_3,
       }),
       removalPolicy: cdk.RemovalPolicy.DESTROY,
+      writer: rds.ClusterInstance.serverlessV2('ServerlessInstance'),
+      vpc,
+      vpcSubnets: { subnetType: ec2.SubnetType.PRIVATE_ISOLATED },
     });
 
     const databaseSetupFunction = new lambda.DockerImageFunction(
@@ -93,8 +89,8 @@ export class SemanticSearchStack extends cdk.Stack {
       }
     );
 
-    dbInstance.secret?.grantRead(databaseSetupFunction);
-    dbInstance.connections.allowDefaultPortFrom(databaseSetupFunction);
+    dbCluster.secret?.grantRead(databaseSetupFunction);
+    dbCluster.connections.allowDefaultPortFrom(databaseSetupFunction);
 
     const databaseSetupProvider = new cr.Provider(
       this,
@@ -105,30 +101,36 @@ export class SemanticSearchStack extends cdk.Stack {
       }
     );
 
-    new cdk.CustomResource(this, 'DatabaseSetupResource', {
-      removalPolicy: cdk.RemovalPolicy.DESTROY,
-      serviceToken: databaseSetupProvider.serviceToken,
-      properties: {
-        DB_SECRET_ID: dbInstance.secret?.secretArn as string,
-        INDEX_TYPES: indexTypes.join(','),
-      },
-    });
+    const dbSetupResource = new cdk.CustomResource(
+      this,
+      'DatabaseSetupResource',
+      {
+        removalPolicy: cdk.RemovalPolicy.DESTROY,
+        serviceToken: databaseSetupProvider.serviceToken,
+        properties: {
+          DB_SECRET_ID: dbCluster.secret?.secretArn as string,
+          INDEX_TYPES: indexTypes.join(','),
+        },
+      }
+    );
+
+    dbSetupResource.node.addDependency(dbCluster);
 
-    return { dbInstance };
+    return { dbCluster };
   }
 
   private createEmbeddingsEndpoint({ vpc }: { vpc: ec2.Vpc }) {
     const embeddingsModel = new HuggingFaceCustomScriptModel(
       this,
-      'EmbeddingsCustomScriptModel',
+      'EmbeddingsModel',
       {
         vpc,
         region: this.region,
         modelId: [
           'sentence-transformers/all-MiniLM-L6-v2',
           'cross-encoder/ms-marco-MiniLM-L-12-v2',
         ],
-        codeFolder: './lib/semantic-search/embeddings-model',
+        codeFolder: './lib/aurora-semantic-search/embeddings-model',
         instanceType: 'ml.g4dn.xlarge',
         codeBuildComputeType: codebuild.ComputeType.LARGE,
       }
@@ -139,11 +141,11 @@ export class SemanticSearchStack extends cdk.Stack {
 
   private createAPI({
     vpc,
-    dbInstance,
+    dbCluster,
     embeddingsEndpoint,
   }: {
     vpc: ec2.Vpc;
-    dbInstance: rds.DatabaseInstance;
+    dbCluster: rds.DatabaseInstance | rds.DatabaseCluster;
     embeddingsEndpoint: sagemaker.CfnEndpoint;
   }) {
     const semanticSearchApi = new lambda.DockerImageFunction(
@@ -161,15 +163,15 @@ export class SemanticSearchStack extends cdk.Stack {
         environment: {
           REGION_NAME: this.region,
           LOG_LEVEL: 'DEBUG',
-          DB_SECRET_ID: dbInstance.secret?.secretArn as string,
+          DB_SECRET_ID: dbCluster.secret?.secretArn as string,
           EMBEDDINGS_ENDPOINT_NAME: embeddingsEndpoint.attrEndpointName,
           CROSS_ENCODER_ENDPOINT_NAME: embeddingsEndpoint.attrEndpointName,
         },
       }
     );
 
-    dbInstance.secret?.grantRead(semanticSearchApi);
-    dbInstance.connections.allowDefaultPortFrom(semanticSearchApi);
+    dbCluster.secret?.grantRead(semanticSearchApi);
+    dbCluster.connections.allowDefaultPortFrom(semanticSearchApi);
 
     semanticSearchApi.addToRolePolicy(
       new iam.PolicyStatement({

lib/semantic-search/document-indexing-stack.ts → lib/aurora-semantic-search/document-indexing-stack.ts

@@ -15,60 +15,70 @@ import { Construct } from 'constructs';
 
 export interface DocumentIndexingStackProps extends cdk.NestedStackProps {
   vpc: ec2.Vpc;
-  dbInstance: rds.DatabaseInstance;
+  dbCluster: rds.DatabaseInstance | rds.DatabaseCluster;
   embeddingsEndpoint: sagemaker.CfnEndpoint;
 }
 
 export class DocumentIndexingStack extends cdk.NestedStack {
   constructor(scope: Construct, id: string, props: DocumentIndexingStackProps) {
     super(scope, id, props);
 
-    const { vpc, dbInstance, embeddingsEndpoint } = props;
+    const { vpc, dbCluster, embeddingsEndpoint } = props;
 
     const dataBucket = new s3.Bucket(this, 'DataBucket', {
       blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
       removalPolicy: cdk.RemovalPolicy.DESTROY,
       autoDeleteObjects: true,
-      enforceSSL: true,
     });
 
     const dataQueue = new sqs.Queue(this, 'DataQueue', {
       visibilityTimeout: cdk.Duration.seconds(600),
     });
 
-    dataBucket.addEventNotification(s3.EventType.OBJECT_CREATED, new s3Notifications.SqsDestination(dataQueue));
+    dataBucket.addEventNotification(
+      s3.EventType.OBJECT_CREATED,
+      new s3Notifications.SqsDestination(dataQueue)
+    );
 
-    const documentIndexing = new lambda.DockerImageFunction(this, 'DocumentIndexing', {
-      code: lambda.DockerImageCode.fromImageAsset(path.join(__dirname, './functions/document-indexing')),
-      architecture: lambda.Architecture.X86_64,
-      vpc,
-      vpcSubnets: vpc.selectSubnets({
-        subnetType: ec2.SubnetType.PRIVATE_ISOLATED,
-      }),
-      timeout: cdk.Duration.minutes(10),
-      memorySize: 3072,
-      logRetention: logs.RetentionDays.ONE_DAY,
-      environment: {
-        REGION_NAME: this.region,
-        LOG_LEVEL: 'DEBUG',
-        DB_SECRET_ID: dbInstance.secret?.secretArn as string,
-        EMBEDDINGS_ENDPOINT_NAME: embeddingsEndpoint.attrEndpointName,
-      },
-    });
+    const documentIndexing = new lambda.DockerImageFunction(
+      this,
+      'DocumentIndexing',
+      {
+        code: lambda.DockerImageCode.fromImageAsset(
+          path.join(__dirname, './functions/document-indexing')
+        ),
+        architecture: lambda.Architecture.X86_64,
+        vpc,
+        vpcSubnets: vpc.selectSubnets({
+          subnetType: ec2.SubnetType.PRIVATE_ISOLATED,
+        }),
+        timeout: cdk.Duration.minutes(10),
+        memorySize: 3072,
+        logRetention: logs.RetentionDays.ONE_DAY,
+        environment: {
+          REGION_NAME: this.region,
+          LOG_LEVEL: 'DEBUG',
+          DB_SECRET_ID: dbCluster.secret?.secretArn as string,
+          EMBEDDINGS_ENDPOINT_NAME: embeddingsEndpoint.attrEndpointName,
+        },
+      }
+    );
 
-    dbInstance.secret?.grantRead(documentIndexing);
-    dbInstance.connections.allowDefaultPortFrom(documentIndexing);
+    dbCluster.secret?.grantRead(documentIndexing);
+    dbCluster.connections.allowDefaultPortFrom(documentIndexing);
     dataBucket.grantReadWrite(documentIndexing);
 
     documentIndexing.addToRolePolicy(
       new iam.PolicyStatement({
         actions: ['sagemaker:InvokeEndpoint'],
         resources: [embeddingsEndpoint.ref],
-      }),
+      })
     );
 
     dataQueue.grantConsumeMessages(documentIndexing);
-    documentIndexing.addEventSource(new lambdaEventSources.SqsEventSource(dataQueue));
+    documentIndexing.addEventSource(
+      new lambdaEventSources.SqsEventSource(dataQueue)
+    );
 
     new cdk.CfnOutput(this, 'DataBucketName', {
       value: dataBucket.bucketName,