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

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

curl - v https://96wd9grlbd.execute-api.us-west-2.amazonaws.com/Prod/bedrock?prompt=%22how%20to%20cook%20chicken?%22

'X-Amzn-Trace-Id': 'Root=1-6660d328-424386874f29954758acde95'

fields @log, @timestamp, @message
| filter @requestId = "576f9ca4-630c-4237-a7a1-a8a44060f2d4" or @message like "620c7d45-353a-4096-9541-1e1a5b8950b9" or @message like "1-6660d328-424386874f29954758acde95" or @message like "6660d328424386874f29954758acde95"
| sort @timestamp, @message desc

fields @timestamp, @message
| filter @message like "Root=1-6660d328-424386874f29954758acde95"

# 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)