Sample code/setup in CloudFormation for monitoring and observability for API Gateway and Lambda

Here is sample code with CloudFormation to setup API Gateway and Lambda with monitoring enabled by using CloudWatch, in short.

• Enable Lambda function and API Gateway to push logs to a log group in CloudWatch
• Then we can query log for a request ID using CloudWatch Log Insights.
• We also can analyze request via stage using X-Ray in CloudWatch
• Optionally, can setup ADOT by using a Lambda layer.

Trace in CloudWatch

Segment Timeline

1. CloudFormation for a Lambda function

AWSTemplateFormatVersion: 2010-09-09
Description: 'a lambda function'
Resources:
  # Create an IAM Role for lambda function
  LambdaExecutionRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - lambda.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /demo/logaccess/
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:PutLogEvents'
                Resource: 'arn:aws:logs:*:*:*'
              - Effect: Allow
                Action:
                  - 'bedrock:InvokeModel'
                Resource: !Sub arn:${AWS::Partition}:bedrock:${AWS::Region}::foundation-model/*
  # Create a Lambda function
  LambdaFunction:
    Type: AWS::Lambda::Function
    Properties:
      Handler: 'index.handler'
      Role: !GetAtt LambdaExecutionRole.Arn
      # Environment variables
      Environment:
        Variables:
          REGION: !Sub ${AWS::Region}
      Code:
        ZipFile: |
          import json
          import time
          def handler(event, context):
            # heavy processing job
            # time.sleep(10)
            # raise exception to test api code 500 
            if 1==2:
              raise Exception('Malformed input ...')
            # return 
            return {
              'statusCode': 200,
              'body': json.dumps('Hello from Lambda!')
            }
      Runtime: 'python3.10'
      Timeout: 30
Outputs:
  LambdaFunctionArn:
    Value: !GetAtt LambdaFunction.Arn
    Export:
      Name:
        Fn::Sub: ${AWS::StackName}-LambdaFunctionArn

2. CloudFormation for API Gateway with log enabled

AWSTemplateFormatVersion: '2010-09-09'
Description: Create an rest api
Parameters:
  BedrockLambdaStackName:
    Type: String
    Default: 'cfn-lambda-bedrock'
  ApiName:
    Type: String
    Default: 'rest-api-bedrock'
Resources:
  # Create clodwatch log role for api gateway
  ApiGatewayCloudWatchRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: 'Allow'
            Principal:
              Service:
                - 'apigateway.amazonaws.com'
            Action:
              - 'sts:AssumeRole'
      Path: '/'
      Policies:
        - PolicyName: 'root'
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: 'Allow'
                Action:
                  - 'logs:CreateLogGroup'
                  - 'logs:CreateLogStream'
                  - 'logs:DescribeLogGroups'
                  - 'logs:DescribeLogStreams'
                  - 'logs:PutLogEvents'
                  - 'logs:GetLogEvents'
                  - 'logs:FilterLogEvents'
                Resource: '*'
  # Create an IAM Role for api gateway to invoke lambda functions
  ApiGatewayLambdaRole:
    Type: 'AWS::IAM::Role'
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - apigateway.amazonaws.com
            Action:
              - 'sts:AssumeRole'
      Path: /
      Policies:
        - PolicyName: root
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action: 'lambda:*'
                Resource: '*'
  # Create a rest api
  RestApi:
    Type: AWS::ApiGateway::RestApi
    Properties:
      Name: !Ref ApiName
      Description: 'rest api bedrock'
      EndpointConfiguration:
        Types:
          - REGIONAL
  # Create a resource named bedrock
  BedrockResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId: !Ref RestApi
      ParentId: !GetAtt RestApi.RootResourceId
      PathPart: 'bedrock'
  # Create a GET method for the bedrock resource and integrate with a Lambda function
  GetBedrockMethod:
    Type: AWS::ApiGateway::Method
    Properties:
      RestApiId: !Ref RestApi
      ResourceId: !Ref BedrockResource
      HttpMethod: GET
      AuthorizationType: NONE
      MethodResponses:
        - StatusCode: 200
          ResponseParameters:
            method.response.header.Access-Control-Allow-Origin: true
            method.response.header.Access-Control-Allow-Methods: true
            method.response.header.Access-Control-Allow-Headers: true
          ResponseModels:
            application/json: Empty
        - StatusCode: 500
          ResponseParameters:
            method.response.header.Access-Control-Allow-Origin: true
            method.response.header.Access-Control-Allow-Methods: true
            method.response.header.Access-Control-Allow-Headers: true
          ResponseModels:
            application/json: Empty
      Integration:
        Type: AWS_PROXY
        # get role arn from above created role
        Credentials: !GetAtt ApiGatewayLambdaRole.Arn
        IntegrationResponses:
          - StatusCode: 200
            ResponseParameters:
              method.response.header.Access-Control-Allow-Origin: "'*'"
              method.response.header.Access-Control-Allow-Methods: "'GET,POST,PUT,DELETE,OPTIONS'"
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token'"
          - StatusCode: 500
            # http status regex
            SelectionPattern: 'Malformed.*'
            ResponseParameters:
              method.response.header.Access-Control-Allow-Origin: "'*'"
              method.response.header.Access-Control-Allow-Methods: "'GET,POST,PUT,DELETE,OPTIONS'"
              method.response.header.Access-Control-Allow-Headers: "'Content-Type,X-Amz-Date,Authorization,X-Api-Key,X-Amz-Security-Token'"
        IntegrationHttpMethod: POST
        Uri: !Sub
          - arn:${AWS::Partition}:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations
          - lambdaArn: !ImportValue
              Fn::Sub: ${BedrockLambdaStackName}-LambdaFunctionArn
  # Create deployment stage named Prod
  Deployment:
    Type: AWS::ApiGateway::Deployment
    DependsOn:
      - GetBedrockMethod
    Properties:
      RestApiId: !Ref RestApi
      StageName: Prod

3. Lambda handler example

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import json
import boto3
import os

# model id
MODEL_ID = "anthropic.claude-3-haiku-20240307-v1:0"

# bedrock runtime client
bedrock_client = boto3.client("bedrock-runtime", region_name=os.environ["REGION"])


def call_bedrock(prompt):
    """
    call bedrock claude 3 to describe image
    """
    # body to invoke bedrock claude 3
    messages = []
    # current image
    messages.append(
        {
            "role": "user",
            "content": [
                {"type": "text", "text": prompt},
            ],
        }
    )
    # build body
    body = json.dumps(
        {
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 2048,
            "temperature": 0.5,
            "top_k": 250,
            "top_p": 0.999,
            "messages": messages,
        }
    )
    # invoke model
    response = bedrock_client.invoke_model(
        body=body,
        contentType="application/json",
        accept="*/*",
        modelId=MODEL_ID,
    )
    # model response
    response_body = json.loads(response.get("body").read())
    # response
    return response_body["content"][0]["text"]


def handler(event, context) -> json:
    """
    simple lambda function
    """

    print(event)

    # parse image from event
    prompt = event["queryStringParameters"]["prompt"]

    # describe image
    bot_response = call_bedrock(prompt)

    # return
    return {
        "statusCode": 200,
        "headers": {
            "Access-Control-Allow-Origin": "*",
            "Access-Control-Allow-Headers": "Content-Type",
            "Access-Control-Allow-Methods": "OPTIONS,GET",
        },
        "body": json.dumps({"bot": bot_response}),
    }


if __name__ == "__main__":
    res = handler(
        event={"queryStringParameters": {"prompt": "how to cook chicken soup?"}},
        context=None,
    )
    print(res)