lambda_function_bedrock.py

import json
import boto3
import os

region = 'us-east-1'
bedrock_runtime_client = boto3.client('bedrock-runtime', region_name=region)

def lambda_handler(event, context):
    # TODO implement
    # print(event['prompt'])
    # return {
    #     'statusCode': 200,
    #     'body': event['queryStringParameters']['prompt']
    # }
    # return event
    # print(boto3.__version__)
    
    if (event['queryStringParameters']['type'] == 'image'):
        modelId = 'stability.stable-diffusion-xl-v0'
        
        body = json.dumps({
            "text_prompts": [{"text": event['queryStringParameters']['prompt']}],
            "cfg_scale": 10,
            "seed": 0,
            "steps": 50
        })

    else:
        modelId = 'ai21.j2-mid-v1'
        
        body = json.dumps({
            "prompt": event['queryStringParameters']['prompt'],
            "maxTokens": 200,
            "temperature": 0.7,
            "topP": 1,
            "stopSequences": [],
            "countPenalty": {"scale":0},
            "presencePenalty": {"scale":0},
            "frequencyPenalty": {"scale":0}
        })
        
    contentType = 'application/json'
    accept = 'application/json'
    
    response = bedrock_runtime_client.invoke_model(
        modelId=modelId,
        contentType=contentType,
        accept=accept, 
        body=body
    )
    
    response_body = json.loads(response.get('body').read())
    
    print(response_body)
    
    return {
        'statusCode': 200,
        'body': response_body
    }