LangMem example with Aurora Serverless v2 CDK

aurora_stack.py

from aws_cdk import (
    Stack,
    aws_rds as rds,
    aws_ec2 as ec2,
    aws_lambda as _lambda,
    aws_iam as iam,
    CfnOutput,
    Duration
)
from constructs import Construct
import os


class AuroraStack(Stack):
    def __init__(self, scope: Construct, construct_id: str, **kwargs) -> None:
        super().__init__(scope, construct_id, **kwargs)

        vpc = ec2.Vpc(
            self, "AuroraVpc",
            max_azs=2,
            ip_addresses=ec2.IpAddresses.cidr("10.1.0.0/16"),
            subnet_configuration=[
                ec2.SubnetConfiguration(
                    name="PrivateSubnet",
                    subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS,
                    cidr_mask=24
                ),
                ec2.SubnetConfiguration(
                    name="PublicSubnet",
                    subnet_type=ec2.SubnetType.PUBLIC,
                    cidr_mask=24
                )
            ],
            nat_gateways=1
        )

        lambda_sg = ec2.SecurityGroup(
            self, "LambdaSG",
            vpc=vpc,
            description="Allow Lambda to access Aurora",
            allow_all_outbound=True
        )

        aurora_sg = ec2.SecurityGroup(
            self, "AuroraSG",
            vpc=vpc,
            description="Allow Aurora access from Lambda",
            allow_all_outbound=True
        )
        # Allow Lambda to connect to Aurora
        aurora_sg.add_ingress_rule(
            peer=lambda_sg,
            connection=ec2.Port.tcp(5432),  # Default PostgreSQL port
            description="Allow Lambda to access Aurora"
        )
        # Create Aurora Serverless v2 Cluster
        dbname = "memory"
        cluster = rds.DatabaseCluster(
            self, "AuroraCluster",
            engine=rds.DatabaseClusterEngine.aurora_postgres(
                version=rds.AuroraPostgresEngineVersion.VER_16_6
            ),
            writer=rds.ClusterInstance.serverless_v2("writer"),
            readers=[
                rds.ClusterInstance.serverless_v2(
                    "reader1", scale_with_writer=True),
            ],
            vpc=vpc,
            serverless_v2_min_capacity=0,
            serverless_v2_max_capacity=1,
            vpc_subnets=ec2.SubnetSelection(
                subnet_type=ec2.SubnetType.PRIVATE_WITH_EGRESS
            ),
            credentials=rds.Credentials.from_generated_secret(
                username="langmem"),
            security_groups=[aurora_sg],
            default_database_name=dbname,
            enable_data_api=True
        )

        lambda_role = iam.Role(
            self, "LambdaExecutionRole",
            assumed_by=iam.ServicePrincipal("lambda.amazonaws.com"),
            inline_policies={
                "SecretsManagerAccessPolicy": iam.PolicyDocument(
                    statements=[
                        iam.PolicyStatement(
                            actions=["secretsmanager:GetSecretValue"],
                            resources=[cluster.secret.secret_arn]
                        )
                    ]
                )
            },
            managed_policies=[
                iam.ManagedPolicy.from_aws_managed_policy_name(
                    "service-role/AWSLambdaBasicExecutionRole"),
                # for VPC attached Lambda
                iam.ManagedPolicy.from_aws_managed_policy_name(
                    "service-role/AWSLambdaVPCAccessExecutionRole"),
            ]
        )
        lambda_function = _lambda.Function(
            self, "LongTermMemHandler",
            function_name="LongTermMemHandler",
            runtime=_lambda.Runtime.PYTHON_3_12,
            code=_lambda.Code.from_asset(
                "../langmem-example/deployment_package.zip"),
            handler="longterm_mem.lambda_handler",
            security_groups=[lambda_sg],
            vpc=vpc,
            timeout=Duration.seconds(180),
            memory_size=1024,
            environment={
                "SECRET_NAME": cluster.secret.secret_name,
                "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
            },
            role=lambda_role
        )
        lambda_url = lambda_function.add_function_url(
            auth_type=_lambda.FunctionUrlAuthType.NONE
        )
        CfnOutput(self, "LongTermMemFunctionUrl", value=lambda_url.url)