WIP: Add initial Bedrock Agent and DDB for Calendar events

backend/bedrock_agent/fetch_calendar_events.py

@@ -0,0 +1,62 @@
+# Built-in imports
+import os
+import boto3
+from boto3.dynamodb.conditions import Key
+from botocore.exceptions import ClientError
+
+# TODO: Enhance code to be production grade. This is just a POC
+# (Add logger, add error handling, add optimizations, etc...)
+
+TABLE_NAME = os.environ.get("TABLE_NAME")
+dynamodb_resource = boto3.resource("dynamodb")
+table = dynamodb_resource.Table(TABLE_NAME)
+
+
+def get_all_calendar_events(partition_key: str, sort_key_portion: str) -> list[dict]:
+    """
+    Function to run a query against DynamoDB with partition key and the sort
+    key with <begins-with> functionality on it.
+    :param partition_key (str): partition key value.
+    :param sort_key_portion (str): sort key portion to use in query.
+    """
+    print(
+        f"Starting get_all_calendar_events with "
+        f"pk: ({partition_key}) and sk: ({sort_key_portion})"
+    )
+
+    all_items = []
+    try:
+        # The structure key for a single-table-design "PK" and "SK" naming
+        key_condition = Key("PK").eq(partition_key) & Key("SK").begins_with(
+            sort_key_portion
+        )
+        limit = 50
+
+        # Initial query before pagination
+        response = table.query(
+            KeyConditionExpression=key_condition,
+            Limit=limit,
+        )
+        if "Items" in response:
+            all_items.extend(response["Items"])
+
+        # Pagination loop for possible following queries
+        while "LastEvaluatedKey" in response:
+            response = table.query(
+                KeyConditionExpression=key_condition,
+                Limit=limit,
+                ExclusiveStartKey=response["LastEvaluatedKey"],
+            )
+            if "Items" in response:
+                all_items.extend(response["Items"])
+
+        return all_items
+    except ClientError as error:
+        print(
+            f"query operation failed for: "
+            f"table_name: {TABLE_NAME}."
+            f"pk: {partition_key}."
+            f"sort_key_portion: {sort_key_portion}."
+            f"error: {error}."
+        )
+        raise error

backend/bedrock_agent/lambda_function.py

@@ -0,0 +1,50 @@
+# NOTE: This is a super-MVP code for testing. Still has a lot of gaps to solve/fix. Do not use in prod.
+
+from bedrock_agent.fetch_calendar_events import get_all_calendar_events
+
+
+def lambda_handler(event, context):
+    action_group = event["actionGroup"]
+    _function = event["function"]
+    parameters = event.get("parameters", [])
+
+    print("PARAMETERS ARE: ", parameters)
+
+    # Extract date from parameters
+    date = None
+    event_name = None
+    for param in parameters:
+        if param["name"] == "date":
+            date = param["value"]
+
+    all_events_for_user = get_all_calendar_events(
+        partition_key="USER#[email protected]",  # TODO: Extract user from input when multi-user support
+        sort_key_portion=f"DATE#{date}",
+    )
+    print("DEBUG, ", all_events_for_user)
+
+    # TODO: Add a more robust search engine/algorithm for matching/finding events
+    result_events = []
+    for calendar_event in all_events_for_user:
+        if calendar_event.get("events"):
+            result_events.extend(calendar_event["events"])
+
+    print(f"Events found: {result_events}")
+
+    # Convert the list of events to a string to be able to return it in the response as a string
+    result_events_string = "\n-".join(result_events)
+    response_body = {"TEXT": {"body": result_events_string}}
+
+    action_response = {
+        "actionGroup": action_group,
+        "function": _function,
+        "functionResponse": {"responseBody": str(response_body)},
+    }
+
+    function_response = {
+        "response": action_response,
+        "messageVersion": event["messageVersion"],
+    }
+    print("Response: {}".format(function_response))
+
+    return function_response

cdk.json

@@ -4,23 +4,25 @@
     "main_resources_name": "aws-wpp",
     "tags": {
       "Owner": "Santiago Garcia Arango",
-      "Source": "https://github.com/san99tiago/aws-whatsapp-poc",
+      "Source": "https://github.com/san99tiago/aws-whatsapp-chatbot",
       "Usage": "WhatsApp chatbot deployed on AWS for POC purposes"
     },
     "app_config": {
       "dev": {
         "deployment_environment": "dev",
         "log_level": "DEBUG",
-        "table_name": "aws-whatsapp-poc-dev",
-        "api_gw_name": "wpp-poc",
+        "table_name": "aws-whatsapp-dev",
+        "calendar_events_table_name": "aws-whatsapp-calendar-dev",
+        "api_gw_name": "wpp-dev",
         "secret_name": "/dev/aws-whatsapp-chatbot",
         "meta_endpoint": "https://graph.facebook.com/"
       },
       "prod": {
         "deployment_environment": "prod",
         "log_level": "DEBUG",
         "table_name": "aws-whatsapp-prod",
-        "api_gw_name": "wpp-webhook",
+        "calendar_events_table_name": "aws-whatsapp-calendar-prod",
+        "api_gw_name": "wpp-prod",
         "secret_name": "/prod/aws-whatsapp-chatbot",
         "meta_endpoint": "https://graph.facebook.com/"
       }

cdk/stacks/cdk_chatbot_api_stack.py

@@ -4,10 +4,13 @@
 # External imports
 from aws_cdk import (
     Duration,
+    aws_bedrock,
     aws_dynamodb,
+    aws_iam,
     aws_lambda,
     aws_lambda_event_sources,
     aws_logs,
+    aws_ssm,
     aws_secretsmanager,
     aws_stepfunctions as aws_sfn,
     aws_stepfunctions_tasks as aws_sfn_tasks,
@@ -59,6 +62,7 @@ def __init__(
         self.create_state_machine_tasks()
         self.create_state_machine_definition()
         self.create_state_machine()
+        self.create_bedrock_components()
 
         # Generate CloudFormation outputs
         self.generate_cloudformation_outputs()
@@ -79,7 +83,7 @@ def create_dynamodb_table(self):
         """
         self.dynamodb_table = aws_dynamodb.Table(
             self,
-            "DynamoDB-Table",
+            "DynamoDB-Table-Chatbot",
             table_name=self.app_config["table_name"],
             partition_key=aws_dynamodb.Attribute(
                 name="PK", type=aws_dynamodb.AttributeType.STRING
@@ -199,6 +203,43 @@ def create_lambda_functions(self) -> None:
         )
         self.secret_chatbot.grant_read(self.lambda_state_machine_process_message)
 
+        # Lambda Function for the Bedrock Agent Group (fetch recipes)
+        bedrock_agent_lambda_role = aws_iam.Role(
+            self,
+            "BedrockAgentLambdaRole",
+            assumed_by=aws_iam.ServicePrincipal("lambda.amazonaws.com"),
+            description="Role for Bedrock Agent Lambda",
+            managed_policies=[
+                aws_iam.ManagedPolicy.from_aws_managed_policy_name(
+                    "service-role/AWSLambdaBasicExecutionRole",
+                ),
+                aws_iam.ManagedPolicy.from_aws_managed_policy_name(
+                    "AmazonBedrockFullAccess",
+                ),
+                aws_iam.ManagedPolicy.from_aws_managed_policy_name(
+                    "AmazonDynamoDBFullAccess",
+                ),
+            ],
+        )
+
+        # Lambda for the Action Group (used for Bedrock Agents)
+        self.lambda_fetch_calendar_events = aws_lambda.Function(
+            self,
+            "Lambda-AG-FetchCalendarEvents",
+            runtime=aws_lambda.Runtime.PYTHON_3_11,
+            handler="bedrock_agent/lambda_function.lambda_handler",
+            function_name=f"{self.main_resources_name}-bedrock-action-group-calendar-events",
+            code=aws_lambda.Code.from_asset(PATH_TO_LAMBDA_FUNCTION_FOLDER),
+            timeout=Duration.seconds(60),
+            memory_size=512,
+            environment={
+                "ENVIRONMENT": self.app_config["deployment_environment"],
+                "LOG_LEVEL": self.app_config["log_level"],
+                "TABLE_NAME": self.app_config["calendar_events_table_name"],
+            },
+            role=bedrock_agent_lambda_role,
+        )
+
     def create_dynamodb_streams(self) -> None:
         """
         Method to create the DynamoDB Streams for the Lambda Function that will
@@ -500,6 +541,131 @@ def create_state_machine(self) -> None:
             self.state_machine.state_machine_arn,
         )
 
+    def create_bedrock_components(self) -> None:
+        """
+        Method to create the Bedrock Agent for the chatbot.
+        """
+
+        # Generic "PK" and "SK", to leverage Single-Table-Design
+        self.calendar_events_dynamodb_table = aws_dynamodb.Table(
+            self,
+            "DynamoDB-Table-CalendarEvents",
+            table_name=self.app_config["calendar_events_table_name"],
+            partition_key=aws_dynamodb.Attribute(
+                name="PK", type=aws_dynamodb.AttributeType.STRING
+            ),
+            sort_key=aws_dynamodb.Attribute(
+                name="SK", type=aws_dynamodb.AttributeType.STRING
+            ),
+            billing_mode=aws_dynamodb.BillingMode.PAY_PER_REQUEST,
+            removal_policy=RemovalPolicy.DESTROY,
+        )
+        Tags.of(self.calendar_events_dynamodb_table).add(
+            "Name", self.app_config["table_name"]
+        )
+
+        # Add permissions to the Lambda function resource policy. You use a resource-based policy to allow an AWS service to invoke your function.
+        self.lambda_fetch_calendar_events.add_permission(
+            "AllowBedrock",
+            principal=aws_iam.ServicePrincipal("bedrock.amazonaws.com"),
+            action="lambda:InvokeFunction",
+            source_arn=f"arn:aws:bedrock:{self.region}:{self.account}:agent/*",
+        )
+
+        bedrock_agent_role = aws_iam.Role(
+            self,
+            "BedrockAgentRole",
+            assumed_by=aws_iam.ServicePrincipal("bedrock.amazonaws.com"),
+            description="Role for Bedrock Agent",
+            managed_policies=[
+                aws_iam.ManagedPolicy.from_aws_managed_policy_name(
+                    "AmazonBedrockFullAccess",
+                ),
+                aws_iam.ManagedPolicy.from_aws_managed_policy_name(
+                    "AWSLambda_FullAccess",
+                ),
+                aws_iam.ManagedPolicy.from_aws_managed_policy_name(
+                    "CloudWatchLogsFullAccess",
+                ),
+            ],
+        )
+        # Add additional IAM actions for the bedrock agent
+        bedrock_agent_role.add_to_policy(
+            aws_iam.PolicyStatement(
+                effect=aws_iam.Effect.ALLOW,
+                actions=[
+                    "bedrock:InvokeModel",
+                    "bedrock:InvokeModelEndpoint",
+                    "bedrock:InvokeModelEndpointAsync",
+                ],
+                resources=["*"],
+            )
+        )
+
+        self.bedrock_agent = aws_bedrock.CfnAgent(
+            self,
+            "BedrockAgent",
+            agent_name=f"{self.main_resources_name}-agent",
+            agent_resource_role_arn=bedrock_agent_role.role_arn,
+            description="Agent for chatbot",
+            foundation_model="anthropic.claude-3-haiku-20240307-v1:0",
+            instruction="You are a specialized agent in giving back calendar events based on the user's input <date>. User must provide the date, and you will make sure it has the format 'YYYY-MM-DD' for the parameter <date>. You will use it to get the list of events for that day and return them in a structured format.",
+            auto_prepare=True,
+            action_groups=[
+                aws_bedrock.CfnAgent.AgentActionGroupProperty(
+                    action_group_name="FetchCalendarEvents",
+                    description="A function that is able to fetch the calendar events from the database from an input date.",
+                    action_group_executor=aws_bedrock.CfnAgent.ActionGroupExecutorProperty(
+                        lambda_=self.lambda_fetch_calendar_events.function_arn,
+                    ),
+                    function_schema=aws_bedrock.CfnAgent.FunctionSchemaProperty(
+                        functions=[
+                            aws_bedrock.CfnAgent.FunctionProperty(
+                                name="FetchCalendarEvents",
+                                # the properties below are optional
+                                description="Function to fetch the calendar events based on the input input date",
+                                parameters={
+                                    "date": aws_bedrock.CfnAgent.ParameterDetailProperty(
+                                        type="string",
+                                        description="Date to fetch the calendar events",
+                                        required=True,
+                                    ),
+                                },
+                            )
+                        ]
+                    ),
+                ),
+            ],
+        )
+
+        # Create an alias for the bedrock agent
+        cfn_agent_alias = aws_bedrock.CfnAgentAlias(
+            self,
+            "MyCfnAgentAlias",
+            agent_alias_name="bedrock-agent-alias",
+            agent_id=self.bedrock_agent.ref,
+            description="bedrock agent alias to simplify agent invocation",
+        )
+        cfn_agent_alias.add_dependency(self.bedrock_agent)
+
+        # This string will be as <AGENT_ID>|<AGENT_ALIAS_ID>
+        agent_alias_string = cfn_agent_alias.ref
+
+        # Create SSM Parameters for the agent alias to use in the Lambda functions
+        # Note: can not be added as Env-Vars due to circular dependency. Thus, SSM Params (decouple)
+        aws_ssm.StringParameter(
+            self,
+            "SSMAgentAlias",
+            parameter_name=f"/{self.deployment_environment}/aws-wpp/bedrock-agent-alias-id-full-string",
+            string_value=agent_alias_string,
+        )
+        aws_ssm.StringParameter(
+            self,
+            "SSMAgentId",
+            parameter_name=f"/{self.deployment_environment}/aws-wpp/bedrock-agent-id",
+            string_value=self.bedrock_agent.ref,
+        )
+
     def generate_cloudformation_outputs(self) -> None:
         """
         Method to add the relevant CloudFormation outputs.