Accelerate building agents for Amazon Bedrock using Powertools for AWS Lambda and Generative AI CDK constructs

				
					import os
import boto3
import json

from time import time
from datetime import date

from pydantic import EmailStr, FutureDate
from typing_extensions import Annotated

from aws_lambda_powertools import Logger, Tracer
from aws_lambda_powertools.event_handler import BedrockAgentResolver
from aws_lambda_powertools.event_handler.openapi.params import Body, Query
from aws_lambda_powertools.utilities.typing import LambdaContext

logger = Logger()
tracer = Tracer()

S3_BUCKET = os.environ["S3_BUCKET"]

# Init an event resolver for Amazon Bedrock Agents
app = BedrockAgentResolver()

@app.get("/list_project_source_files", description="Lists available source code files")  
@tracer.capture_method
def list_project_source_files(
    project: Annotated[str, Query(title="Project", description="The project to filter on when listing source code files.")],  
) -> Annotated[dict, Body(description="List of available source code files for the given project")]:   
    logger.append_keys(
        session_id=app.current_event.session_id,
        action_group=app.current_event.action_group,
        input_text=app.current_event.input_text,
    )
    
    s3_client = boto3.client("s3")
    paginator = s3_client.get_paginator('list_objects_v2')
    
    objects = []
    try:
        for page in paginator.paginate(Bucket=S3_BUCKET, Prefix=project):
            if 'Contents' in page:
                for obj in page['Contents']:
                    objects.append(obj['Key'])
        
        logger.info("list_project_source_files response successful")
        return {"files": objects}
    except Exception as e:
        logger.error(f"Error listing files: {str(e)}")
        raise

@app.get("/get_project_source_file", description="Fetches a given project source code file.")  
@tracer.capture_method
def get_project_source_file(
    project_file: Annotated[str, Query(title="Project file", description="The project file to fetch.")],
) -> Annotated[dict, Body(description="The source code for the given project file")]:
    logger.append_keys(
        session_id=app.current_event.session_id,
        action_group=app.current_event.action_group,
        input_text=app.current_event.input_text,
    )

    s3_client = boto3.client("s3")
    try:
        response = s3_client.get_object(Bucket=S3_BUCKET, Key=project_file)
        file_content = response['Body'].read().decode('utf-8')

        logger.info("get_project_source_file response successful")
        return {"file_content": file_content}
    except Exception as e:
        logger.error(f"Error fetching file: {str(e)}")
        raise

@app.post("/set_project_source_file", description="Saves the given project source code file.")  
@tracer.capture_method
def set_project_source_file(
    project_file: Annotated[str, Body(title="Project file", description="The project file to write to.")],
    project_file_content: Annotated[str, Body(title="Project file contents", description="The source code contents to write to the project file.")],
) -> Annotated[bool, Body(description="Successfully updated project source code file")]:
    logger.append_keys(
        session_id=app.current_event.session_id,
        action_group=app.current_event.action_group,
        input_text=app.current_event.input_text,
    )

    s3_client = boto3.client("s3")
    try:
        s3_client.put_object(Bucket=S3_BUCKET, Key=project_file, Body=project_file_content)

        logger.info("set_project_source_file response successful")
        return True
    except Exception as e:
        logger.error(f"Error writing file: {str(e)}")
        raise

@app.post("/exec_sql", description="Execute the provided SQL statement against the specified database endpoint.")  
@tracer.capture_method
def exec_sql(
    database_connection: Annotated[str, Body(title="Database connection", description="The database connection to use to connect to the target database.")],
    sql_statement: Annotated[str, Body(title="SQL statement", description="The SQL statement to execute against the target database.")],
) -> Annotated[bool, Body(description="Successfully executed SQL statement")]:
    logger.append_keys(
        session_id=app.current_event.session_id,
        action_group=app.current_event.action_group,
        input_text=app.current_event.input_text,
    )

    # MOCK execution of the SQL statement
    logger.info({"database": database_connection, "sql_statement": sql_statement})

    return True

@app.post("/save_misc_file", description="Saves misc files generated as part of code conversion. E.g. conversion reports")  
@tracer.capture_method
def save_misc_file(
    project: Annotated[str, Body(title="Project", description="The project in which to save the file under.")],  
    file_name: Annotated[str, Body(title="Filename", description="The name of the file to write to.")],
    file_content: Annotated[str, Body(title="File contents", description="The contents to write to the file.")],
) -> Annotated[bool, Body(description="Successfully updated misc project file")]:
    logger.append_keys(
        session_id=app.current_event.session_id,
        action_group=app.current_event.action_group,
        input_text=app.current_event.input_text,
    )

    s3_client = boto3.client("s3")
    try:
        s3_client.put_object(Bucket=S3_BUCKET, Key=f"{project}/{file_name}", Body=file_content)

        logger.info("save_misc_file response successful")
        return True
    except Exception as e:
        logger.error(f"Error writing file: {str(e)}")
        raise


@logger.inject_lambda_context(log_event=True)
@tracer.capture_lambda_handler
def lambda_handler(event: dict, context: LambdaContext):
    return app.resolve(event, context)

if __name__ == "__main__":  
    print(app.get_openapi_json_schema())