Add VectorKnowledgeBase construct for knowledge base and guardrail

Author: kris-szlapa

packages/cdk/constructs/BedrockGuardrail.ts

@@ -0,0 +1,88 @@
+import {Construct} from "constructs"
+import {bedrock} from "@cdklabs/generative-ai-cdk-constructs"
+import {Role} from "aws-cdk-lib/aws-iam"
+import {Bucket} from "aws-cdk-lib/aws-s3"
+import {
+  ContentFilterType,
+  ContentFilterStrength,
+  ManagedWordFilterType,
+  PIIType,
+  GuardrailAction
+} from "@cdklabs/generative-ai-cdk-constructs/lib/cdk-lib/bedrock/guardrails/guardrail-filters"
+
+export interface VectorKnowledgeBaseProps {
+  kbName: string
+  embeddingsModel: bedrock.BedrockFoundationModel
+  docsBucket: Bucket
+  bedrockExecutionRole: Role
+}
+
+export class VectorKnowledgeBase extends Construct {
+  public readonly knowledgeBase: bedrock.VectorKnowledgeBase
+  public readonly guardrail: bedrock.Guardrail
+
+  constructor(scope: Construct, id: string, props: VectorKnowledgeBaseProps) {
+    super(scope, id)
+
+    this.guardrail = new bedrock.Guardrail(this, "Guardrail", {
+      name: "eps-assist-guardrail",
+      description: "Guardrail for EPS Assist Me Slackbot",
+      blockedInputMessaging: "Your input was blocked.",
+      blockedOutputsMessaging: "Your output was blocked.",
+      contentFilters: [
+        {
+          type: ContentFilterType.SEXUAL,
+          inputStrength: ContentFilterStrength.HIGH,
+          outputStrength: ContentFilterStrength.HIGH
+        },
+        {
+          type: ContentFilterType.VIOLENCE,
+          inputStrength: ContentFilterStrength.HIGH,
+          outputStrength: ContentFilterStrength.HIGH
+        },
+        {
+          type: ContentFilterType.HATE,
+          inputStrength: ContentFilterStrength.HIGH,
+          outputStrength: ContentFilterStrength.HIGH
+        },
+        {
+          type: ContentFilterType.INSULTS,
+          inputStrength: ContentFilterStrength.HIGH,
+          outputStrength: ContentFilterStrength.HIGH
+        },
+        {
+          type: ContentFilterType.MISCONDUCT,
+          inputStrength: ContentFilterStrength.HIGH,
+          outputStrength: ContentFilterStrength.HIGH
+        },
+        {
+          type: ContentFilterType.PROMPT_ATTACK,
+          inputStrength: ContentFilterStrength.HIGH,
+          outputStrength: ContentFilterStrength.NONE
+        }
+      ],
+      piiFilters: [
+        {type: PIIType.General.EMAIL, action: GuardrailAction.ANONYMIZE},
+        {type: PIIType.General.PHONE, action: GuardrailAction.ANONYMIZE},
+        {type: PIIType.General.NAME, action: GuardrailAction.ANONYMIZE},
+        {type: PIIType.Finance.CREDIT_DEBIT_CARD_NUMBER, action: GuardrailAction.BLOCK}
+      ],
+      managedWordListFilters: [
+        {type: ManagedWordFilterType.PROFANITY}
+      ]
+    })
+
+    // Main construct - let it create its own default OpenSearch collection
+    this.knowledgeBase = new bedrock.VectorKnowledgeBase(this, "VectorKB", {
+      name: props.kbName,
+      description: "Knowledge base for EPS Assist Me Slackbot",
+      embeddingsModel: props.embeddingsModel,
+      existingRole: props.bedrockExecutionRole
+    })
+
+    // Add S3 data source to knowledge base
+    this.knowledgeBase.addS3DataSource({
+      bucket: props.docsBucket
+    })
+  }
+}

packages/cdk/resources/BedrockResources.ts

@@ -4,19 +4,19 @@ import {
   StackProps,
   CfnOutput
 } from "aws-cdk-lib"
-
 import {PolicyStatement} from "aws-cdk-lib/aws-iam"
 import * as cdk from "aws-cdk-lib"
 import * as iam from "aws-cdk-lib/aws-iam"
 import * as ops from "aws-cdk-lib/aws-opensearchserverless"
 import * as cr from "aws-cdk-lib/custom-resources"
+import {bedrock} from "@cdklabs/generative-ai-cdk-constructs"
 import {nagSuppressions} from "../nagSuppressions"
 import {Apis} from "../resources/Apis"
 import {Functions} from "../resources/Functions"
 import {Storage} from "../resources/Storage"
 import {Secrets} from "../resources/Secrets"
 import {OpenSearchResources} from "../resources/OpenSearchResources"
-import {BedrockResources} from "../resources/BedrockResources"
+import {VectorKnowledgeBase} from "../resources/VectorKnowledgeBase"
 
 const EMBEDDING_MODEL = "amazon.titan-embed-text-v2:0"
 const VECTOR_INDEX_NAME = "eps-assist-os-index"
@@ -83,9 +83,7 @@ export class EpsAssistMeStack extends Stack {
     bedrockExecutionRole.addToPolicy(bedrockKBDeleteRolePolicy)
 
     // Create Storage construct
-    const storage = new Storage(this, "Storage", {
-      bedrockExecutionRole
-    })
+    const storage = new Storage(this, "Storage", {bedrockExecutionRole})
 
     // Create an IAM policy for S3 access
     const s3AccessListPolicy = new PolicyStatement({
@@ -143,15 +141,15 @@ export class EpsAssistMeStack extends Stack {
 
     const endpoint = openSearchResources.collection.endpoint
 
-    // Create Bedrock Resources
-    const bedrockResources = new BedrockResources(this, "BedrockResources", {
-      bedrockExecutionRole,
-      osCollection: openSearchResources.collection.collection,
-      kbDocsBucket: storage.kbDocsBucket.bucket,
-      region
+    // Create VectorKnowledgeBase construct
+    const vectorKB = new VectorKnowledgeBase(this, "VectorKB", {
+      kbName: "eps-assist-kb",
+      embeddingsModel: bedrock.BedrockFoundationModel.TITAN_EMBED_TEXT_V2_1024,
+      docsBucket: storage.kbDocsBucket.bucket,
+      bedrockExecutionRole
     })
 
-    // Create Functions construct
+    // Functions construct: use guardrail and kb from the new construct
     const functions = new Functions(this, "Functions", {
       stackName: props.stackName,
       version: props.version,
@@ -161,10 +159,10 @@ export class EpsAssistMeStack extends Stack {
       createIndexFunctionRole,
       slackBotTokenParameter: secrets.slackBotTokenParameter,
       slackSigningSecretParameter: secrets.slackSigningSecretParameter,
-      guardrailId: bedrockResources.guardrail.guardrailId,
-      guardrailVersion: bedrockResources.guardrail.guardrailVersionId,
+      guardrailId: vectorKB.guardrail.guardrailId,
+      guardrailVersion: vectorKB.guardrail.guardrailVersion,
       collectionId: openSearchResources.collection.collection.attrId,
-      knowledgeBaseId: bedrockResources.knowledgeBase.attrKnowledgeBaseId,
+      knowledgeBaseId: vectorKB.knowledgeBase.knowledgeBaseId,
       region,
       account,
       slackBotTokenSecret: secrets.slackBotTokenSecret,
@@ -191,6 +189,7 @@ export class EpsAssistMeStack extends Stack {
     })
     openSearchResources.collection.collection.addDependency(aossAccessPolicy)
 
+    // Create a custom resource to create the OpenSearch index
     const vectorIndex = new cr.AwsCustomResource(this, "VectorIndex", {
       installLatestAwsSdk: true,
       onCreate: {
@@ -235,10 +234,10 @@ export class EpsAssistMeStack extends Stack {
     vectorIndex.node.addDependency(openSearchResources.collection.collection)
 
     // add a dependency for bedrock kb on the custom resource. Enables vector index to be created before KB
-    bedrockResources.knowledgeBase.node.addDependency(vectorIndex)
-    bedrockResources.knowledgeBase.node.addDependency(functions.functions.createIndex)
-    bedrockResources.knowledgeBase.node.addDependency(openSearchResources.collection.collection)
-    bedrockResources.knowledgeBase.node.addDependency(bedrockExecutionRole)
+    vectorKB.knowledgeBase.node.addDependency(vectorIndex)
+    vectorKB.knowledgeBase.node.addDependency(functions.functions.createIndex)
+    vectorKB.knowledgeBase.node.addDependency(openSearchResources.collection.collection)
+    vectorKB.knowledgeBase.node.addDependency(bedrockExecutionRole)
 
     // Create Apis and pass the Lambda function
     const apis = new Apis(this, "Apis", {