LangMemの長期記憶をPostgreSQL(pgvector)に永続化する

pip install langmem langgraph-checkpoint-postgres psycopg[binary]

class UserTechInterest(BaseModel):
    """ユーザーの技術に関する関心情報"""
    topic: str = Field(..., description="技術トピック（例：プログラミング言語、クラウドサービスなど）")
    category: str | None = Field(
        None, description="技術の分野（例：フロントエンド、バックエンド、ネットワークなど）"
    )
    interest_level: conint(ge=0, le=100) | None = Field(
        None, description="関心度（0～100のスコアで表現）"
    )
    description: str | None = Field(
        None, description="その他の補足説明（具体例や詳細な知識など）"
    )

def create_conn_url():
    secret_name = os.environ['SECRET_NAME']
    session = boto3.session.Session()
    sm_client = session.client(service_name='secretsmanager')
    get_secret_value_response = sm_client.get_secret_value(SecretId=secret_name)
    secret = json.loads(get_secret_value_response['SecretString'])
    return f"postgresql://{secret["username"]}:{secret["password"]}@{secret["host"]}:{secret["port"]}/{secret["dbname"]}"

postgres_conn_url = create_conn_url()

# LangMemのStore Manager
manager = create_memory_store_manager(
    "openai:gpt-4o-2024-11-20",
    namespace=("chat", "{user_id}"),
    schemas=[UserTechInterest],
    instructions="ユーザーの興味や技術的背景を詳細に抽出してください",
    enable_inserts=True,
    enable_deletes=False,
)

# OpenAIクライアント
client = OpenAI()

def lambda_handler(event, context):
    # 1. LangGraphのPostgreStoreを作成
    with PostgresStore.from_conn_string(postgres_conn_url, index={
        "dims": 1536,
        "embed": "openai:text-embedding-3-small",
    }) as store:
        # 2. PostgerSQLのスキーママイグレーション
        store.setup()

        # 3. LangGraph Functional APIのワークフロー
        @entrypoint(store=store)
        def app(params: dict):
            message = params["message"]
            user_id = params["user_id"]
            # ストアから関連する長期記憶を検索
            memories = store.search(("chat", user_id))
            developer_msg = ("You are a helpful assistant.\n"
                             f"## Memories\n <memories>\n{memories}\n</memories> ")
            # OpenAIのResponses API実行
            response = client.responses.create(
                model="gpt-4o-2024-11-20",
                input=[{"role": "developer", "content": developer_msg}, message],
                tools=[{"type": "web_search_preview"}],
            )
            # 該当ユーザーの嗜好を抽出して長期記憶を更新
            manager.invoke(
                {"messages": [message]},
                config={"configurable": {"user_id": user_id}}
            )
            return response.output_text

        body = json.loads(event["body"])
        user_id = body["user_id"]
        # LangGraphワークフロー起動
        output = app.invoke({
            "message": {
                "role": "user",
                "content": body["prompt"]
            },
            "user_id": user_id
        })

        # レスポンス返却
        return {
            "statusCode": 200,
            "headers": {"Content-Type": "plain/text"},
            "body": output
        }

# 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
    ),
    # Secrets Managerのリソース作成(RDS統合)
    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(
        "path/to/deployment_package.zip"),
    handler="lambda_function.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)

LAMBDA_URL=$(aws lambda get-function-url-config --function-name LongTermMemHandler --query FunctionUrl --output text)
> https://xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx.lambda-url.ap-northeast-1.on.aws/
curl -H 'Content-Type: application/json' "$LAMBDA_URL" -d @- <<'EOF'
{
  "user_id": "MZ0001",
  "prompt": "クラウドサービス全般に興味があり、特にAWS、Azureの基本的な知識を徐々に深めていきたいと思っています"
}
EOF

curl -H 'Content-Type: application/json' "$LAMBDA_URL" -d @- <<'EOF'
{
  "user_id": "MZ0001",
  "prompt": "新しく気になる技術として、AIエージェントを活用した取り組みにも強い興味を持つようになりました"
}
EOF

curl -H 'Content-Type: application/json' "$LAMBDA_URL" -d @- <<'EOF'
{
  "user_id": "MZ0001",
  "prompt": "おすすめの技術トピックを3つ挙げて簡潔にまとめて！"
}
EOF