Implement syncKnowledgeBaseFunction Lambda

Author: kris-szlapa

.github/workflows/cdk_package_code.yml

@@ -63,6 +63,7 @@ jobs:
         run: |
           pip3 install -r packages/slackBotFunction/requirements.txt -t packages/slackBotFunction
           pip3 install -r packages/createIndexFunction/requirements.txt -t packages/createIndexFunction
+          pip3 install -r packages/syncKnowledgeBaseFunction/requirements.txt -t packages/syncKnowledgeBaseFunction
 
       - name: 'Tar files'
         run: |

README.md

@@ -9,9 +9,10 @@ The solution consists of:
 
 - **Slack Bot Function**: AWS Lambda function that handles Slack slash commands and integrates with Amazon Bedrock Knowledge Base
 - **Create Index Function**: AWS Lambda function that creates and manages OpenSearch vector indices for the knowledge base
+- **Sync Knowledge Base Function**: AWS Lambda function that automatically triggers knowledge base ingestion when documents are uploaded to S3
 - **OpenSearch Serverless**: Vector database for storing and searching document embeddings
 - **Amazon Bedrock Knowledge Base**: RAG (Retrieval-Augmented Generation) service with guardrails
-- **S3 Storage**: Document storage for the knowledge base
+- **S3 Storage**: Document storage for the knowledge base with automatic sync triggers
 - **AWS CDK**: Infrastructure as Code for deployment
 
 ## Project Structure
@@ -20,13 +21,14 @@ This is a monorepo with the following structure:
 
 ```
 packages/
-├── cdk/                   # AWS CDK infrastructure code
-│   ├── bin/               # CDK app entry point
-│   ├── constructs/        # Reusable CDK constructs
-│   ├── resources/         # AWS resource definitions
-│   └── stacks/            # CDK stack definitions
-├── createIndexFunction/   # Lambda function for OpenSearch index management
-└── slackBotFunction/      # Lambda function for Slack bot integration
+├── cdk/                      # AWS CDK infrastructure code
+│   ├── bin/                  # CDK app entry point
+│   ├── constructs/           # Reusable CDK constructs
+│   ├── resources/            # AWS resource definitions
+│   └── stacks/               # CDK stack definitions
+├── createIndexFunction/      # Lambda function for OpenSearch index management
+├── slackBotFunction/         # Lambda function for Slack bot integration
+└── syncKnowledgeBaseFunction/ # Lambda function for automatic knowledge base sync
 ```
 
 ## Contributing

packages/cdk/resources/Functions.ts

@@ -18,12 +18,14 @@ export interface FunctionsProps {
   logLevel: string
   createIndexManagedPolicy: ManagedPolicy
   slackBotManagedPolicy: ManagedPolicy
+  syncKnowledgeBaseManagedPolicy: ManagedPolicy
   slackBotTokenParameter: StringParameter
   slackSigningSecretParameter: StringParameter
   guardrailId: string
   guardrailVersion: string
   collectionId: string
   knowledgeBaseId: string
+  dataSourceId: string
   region: string
   account: string
   slackBotTokenSecret: Secret
@@ -74,9 +76,25 @@ export class Functions extends Construct {
     props.slackBotTokenSecret.grantRead(slackBotLambda.function)
     props.slackBotSigningSecret.grantRead(slackBotLambda.function)
 
+    // Lambda function to sync knowledge base on S3 events
+    const syncKnowledgeBaseFunction = new LambdaFunction(this, "SyncKnowledgeBaseFunction", {
+      stackName: props.stackName,
+      functionName: `${props.stackName}-SyncKnowledgeBaseFunction`,
+      packageBasePath: "packages/syncKnowledgeBaseFunction",
+      entryPoint: "app.py",
+      logRetentionInDays: props.logRetentionInDays,
+      logLevel: props.logLevel,
+      environmentVariables: {
+        "KNOWLEDGEBASE_ID": props.knowledgeBaseId || "placeholder",
+        "DATA_SOURCE_ID": props.dataSourceId || "placeholder"
+      },
+      additionalPolicies: [props.syncKnowledgeBaseManagedPolicy]
+    })
+
     this.functions = {
       createIndex: createIndexFunction,
-      slackBot: slackBotLambda
+      slackBot: slackBotLambda,
+      syncKnowledgeBase: syncKnowledgeBaseFunction
     }
   }
 }

packages/cdk/resources/IamResources.ts

@@ -24,6 +24,7 @@ export class IamResources extends Construct {
   public readonly bedrockExecutionRole: Role
   public readonly createIndexManagedPolicy: ManagedPolicy
   public readonly slackBotManagedPolicy: ManagedPolicy
+  public readonly syncKnowledgeBaseManagedPolicy: ManagedPolicy
 
   constructor(scope: Construct, id: string, props: IamResourcesProps) {
     super(scope, id)
@@ -158,5 +159,23 @@ export class IamResources extends Construct {
         slackBotGuardrailPolicy
       ]
     })
+
+    // Create managed policy for SyncKnowledgeBase Lambda function
+    const syncKnowledgeBasePolicy = new PolicyStatement({
+      actions: [
+        "bedrock-agent:StartIngestionJob",
+        "bedrock-agent:GetIngestionJob",
+        "bedrock-agent:ListIngestionJobs"
+      ],
+      resources: [
+        `arn:aws:bedrock:${props.region}:${props.account}:knowledge-base/*`,
+        `arn:aws:bedrock:${props.region}:${props.account}:knowledge-base/*/data-source/*`
+      ]
+    })
+
+    this.syncKnowledgeBaseManagedPolicy = new ManagedPolicy(this, "SyncKnowledgeBaseManagedPolicy", {
+      description: "Policy for SyncKnowledgeBase Lambda to trigger ingestion jobs",
+      statements: [syncKnowledgeBasePolicy]
+    })
   }
 }

packages/cdk/resources/VectorKnowledgeBaseResources.ts

@@ -17,6 +17,7 @@ export interface VectorKnowledgeBaseProps {
 export class VectorKnowledgeBaseResources extends Construct {
   public readonly knowledgeBase: CfnKnowledgeBase
   public readonly guardrail: CfnGuardrail
+  public readonly dataSource: CfnDataSource
 
   constructor(scope: Construct, id: string, props: VectorKnowledgeBaseProps) {
     super(scope, id)
@@ -81,7 +82,7 @@ export class VectorKnowledgeBaseResources extends Construct {
     })
 
     // Create S3 data source for knowledge base documents
-    new CfnDataSource(this, "S3DataSource", {
+    this.dataSource = new CfnDataSource(this, "S3DataSource", {
       knowledgeBaseId: this.knowledgeBase.attrKnowledgeBaseId,
       name: `${props.stackName}-s3-datasource`,
       dataSourceConfiguration: {

packages/cdk/stacks/EpsAssistMeStack.ts

@@ -4,6 +4,8 @@ import {
   StackProps,
   CfnOutput
 } from "aws-cdk-lib"
+import {EventType} from "aws-cdk-lib/aws-s3"
+import {LambdaDestination} from "aws-cdk-lib/aws-s3-notifications"
 import {nagSuppressions} from "../nagSuppressions"
 import {Apis} from "../resources/Apis"
 import {Functions} from "../resources/Functions"
@@ -82,12 +84,14 @@ export class EpsAssistMeStack extends Stack {
       logLevel,
       createIndexManagedPolicy: iamResources.createIndexManagedPolicy,
       slackBotManagedPolicy: iamResources.slackBotManagedPolicy,
+      syncKnowledgeBaseManagedPolicy: iamResources.syncKnowledgeBaseManagedPolicy,
       slackBotTokenParameter: secrets.slackBotTokenParameter,
       slackSigningSecretParameter: secrets.slackSigningSecretParameter,
       guardrailId: "", // Will be set after vector KB is created
       guardrailVersion: "", // Will be set after vector KB is created
       collectionId: openSearchResources.collection.collection.attrId,
       knowledgeBaseId: "", // Will be set after vector KB is created
+      dataSourceId: "", // Will be set after vector KB is created
       region,
       account,
       slackBotTokenSecret: secrets.slackBotTokenSecret,
@@ -119,6 +123,22 @@ export class EpsAssistMeStack extends Stack {
     functions.functions.slackBot.function.addEnvironment("GUARD_RAIL_VERSION", vectorKB.guardrail.attrVersion)
     functions.functions.slackBot.function.addEnvironment("KNOWLEDGEBASE_ID", vectorKB.knowledgeBase.attrKnowledgeBaseId)
 
+    // Update SyncKnowledgeBase Lambda environment variables with vector KB info
+    functions.functions.syncKnowledgeBase.function.addEnvironment(
+      "KNOWLEDGEBASE_ID",
+      vectorKB.knowledgeBase.attrKnowledgeBaseId
+    )
+    functions.functions.syncKnowledgeBase.function.addEnvironment(
+      "DATA_SOURCE_ID",
+      vectorKB.dataSource.attrDataSourceId
+    )
+
+    // Add S3 event notification to trigger sync function
+    storage.kbDocsBucket.bucket.addEventNotification(
+      EventType.OBJECT_CREATED,
+      new LambdaDestination(functions.functions.syncKnowledgeBase.function)
+    )
+
     // Create Apis and pass the Lambda function
     const apis = new Apis(this, "Apis", {
       stackName: props.stackName,

packages/syncKnowledgeBaseFunction/app.py

@@ -0,0 +1,52 @@
+import os
+import boto3
+from botocore.exceptions import ClientError
+from aws_lambda_powertools import Logger
+
+# Configure logging
+logger = Logger()
+
+# Initialize Bedrock client
+bedrock_agent = boto3.client("bedrock-agent")
+
+
+@logger.inject_lambda_context
+def lambda_handler(event, context):
+    """
+    Lambda handler for S3 events that triggers knowledge base ingestion.
+    """
+    knowledge_base_id = os.environ.get("KNOWLEDGEBASE_ID")
+    data_source_id = os.environ.get("DATA_SOURCE_ID")
+
+    if not knowledge_base_id or not data_source_id:
+        logger.error("Missing required environment variables: KNOWLEDGEBASE_ID or DATA_SOURCE_ID")
+        return {"statusCode": 500, "body": "Configuration error"}
+
+    try:
+        # Process S3 event records
+        for record in event.get("Records", []):
+            if record.get("eventSource") == "aws:s3":
+                bucket = record["s3"]["bucket"]["name"]
+                key = record["s3"]["object"]["key"]
+                event_name = record["eventName"]
+
+                logger.info(f"Processing S3 event: {event_name} for {bucket}/{key}")
+
+                # Start ingestion job for the knowledge base
+                response = bedrock_agent.start_ingestion_job(
+                    knowledgeBaseId=knowledge_base_id,
+                    dataSourceId=data_source_id,
+                    description=f"Auto-sync triggered by S3 event: {event_name} on {key}",
+                )
+
+                job_id = response["ingestionJob"]["ingestionJobId"]
+                logger.info(f"Started ingestion job {job_id} for knowledge base {knowledge_base_id}")
+
+        return {"statusCode": 200, "body": "Ingestion triggered successfully"}
+
+    except ClientError as e:
+        logger.error(f"AWS service error: {e}")
+        return {"statusCode": 500, "body": f"AWS error: {str(e)}"}
+    except Exception as e:
+        logger.error(f"Unexpected error: {e}")
+        return {"statusCode": 500, "body": f"Error: {str(e)}"}

packages/syncKnowledgeBaseFunction/requirements.txt

@@ -0,0 +1,2 @@
+boto3>=1.34.0
+aws-lambda-powertools>=2.0.0