aws-genai-llm-chatbot

1
import json
2
import boto3
3
import psycopg2
4
import cfnresponse
5
from aws_lambda_powertools import Logger
6
from aws_lambda_powertools.utilities.typing import LambdaContext
7
from pgvector.psycopg2 import register_vector
8

9
logger = Logger()
10
secretsmanager_client = boto3.client("secretsmanager")
11

12

13
@logger.inject_lambda_context(log_event=True)
14
def lambda_handler(event, context: LambdaContext):
15
    request_type = event["RequestType"]
16
    resource_properties = event["ResourceProperties"]
17
    AURORA_DB_SECRET_ID = resource_properties["AURORA_DB_SECRET_ID"]
18

19
    secret_response = secretsmanager_client.get_secret_value(
20
        SecretId=AURORA_DB_SECRET_ID
21
    )
22
    database_secrets = json.loads(secret_response["SecretString"])
23
    dbhost = database_secrets["host"]
24
    dbport = database_secrets["port"]
25
    dbuser = database_secrets["username"]
26
    dbpass = database_secrets["password"]
27

28
    if request_type == "Create" or request_type == "Update":
29
        dbconn = psycopg2.connect(
30
            host=dbhost, user=dbuser, password=dbpass, port=dbport, connect_timeout=10
31
        )
32

33
        dbconn.set_session(autocommit=True)
34

35
        cur = dbconn.cursor()
36

37
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
38
        register_vector(dbconn)
39

40
        cur.execute("SELECT typname FROM pg_type WHERE typname = 'vector';")
41
        rows = cur.fetchall()
42

43
        for row in rows:
44
            logger.info(f"pg_type.typname: {row}")
45

46
        cur.close()
47
        dbconn.close()
48

49
        logger.info("Created vector extension")
50

51
    cfnresponse.send(event, context, cfnresponse.SUCCESS, {"ok": True})
52

53
    return {"ok": True}
54