Pankaj Katariya
116 lines
4.1 KiB
Python
116 lines
4.1 KiB
Python
# Import necessary libraries
|
|
import json
|
|
import logging
|
|
import boto3
|
|
from botocore.exceptions import ClientError
|
|
|
|
# Configure logging
|
|
logger = logging.getLogger()
|
|
logger.setLevel(logging.INFO)
|
|
|
|
# Lambda function handler
|
|
def lambda_handler(event, context):
|
|
logger.info("Lambda function invoked.")
|
|
logger.info(f"Received event: {json.dumps(event)}")
|
|
|
|
# Call the bedrock_llm function and return its result
|
|
try:
|
|
return call_bedrock_llm(event)
|
|
except ThrottlingException:
|
|
logger.error("Too many requests. Please wait before trying again.")
|
|
raise
|
|
|
|
# Function to invoke the bedrock_llm with necessary parameters
|
|
def invoke(wrapper, event):
|
|
logger.info(f"Event details: {event}")
|
|
|
|
# Extract parameters from the event
|
|
prompt = event['prompt']
|
|
model_id = event['model_id']
|
|
max_gen_len = event['max_gen_len']
|
|
temperature = event['temperature']
|
|
top_p = event['top_p']
|
|
|
|
try:
|
|
# Invoke the bedrock_llm using the provided wrapper
|
|
completion = wrapper.invoke_llm(prompt, model_id, max_gen_len, temperature, top_p)
|
|
logger.info("Bedrock LLM invoked successfully.")
|
|
return completion
|
|
|
|
except ClientError as e:
|
|
# Log an exception if there is an issue invoking the model
|
|
logger.exception(f"Couldn't invoke model {model_id}. Error: {str(e)}")
|
|
raise
|
|
|
|
# Function to call the Bedrock LLM
|
|
def call_bedrock_llm(event):
|
|
logger.info("Calling Bedrock LLM.")
|
|
|
|
# Initialize Bedrock Runtime client
|
|
client = boto3.client(service_name="bedrock-runtime", region_name="us-east-1")
|
|
|
|
# Create an instance of BedrockRuntimeWrapper
|
|
wrapper = BedrockRuntimeWrapper(client)
|
|
|
|
# Invoke the wrapper and return the result
|
|
try:
|
|
answer = invoke(wrapper, event)
|
|
logger.info("Bedrock LLM call completed.")
|
|
return answer
|
|
except ClientError as e:
|
|
if e.response['Error']['Code'] == 'ThrottlingException':
|
|
raise ThrottlingException
|
|
else:
|
|
logger.error(f"Error calling Bedrock LLM: {str(e)}")
|
|
raise
|
|
|
|
# Class to wrap Bedrock Runtime functionality
|
|
class BedrockRuntimeWrapper:
|
|
# Constructor to initialize the wrapper with a Bedrock Runtime client
|
|
def __init__(self, bedrock_runtime_client):
|
|
self.bedrock_runtime_client = bedrock_runtime_client
|
|
|
|
# Method to invoke the Bedrock LLM model
|
|
def invoke_llm(self, prompt, model_id, max_gen_len, temperature, top_p):
|
|
logger.info("Invoking Bedrock LLM model.")
|
|
|
|
try:
|
|
# Prepare the request body
|
|
if 'claude' in model_id.lower():
|
|
# If yes, change the parameter name to 'max_tokens_to_sample' for claude
|
|
prompt='Human:'+prompt+'\n\nAssistant:You read and understand the USA healthcrae contract and able to answer questions based on given contract.'
|
|
body = {
|
|
"prompt": prompt,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_tokens_to_sample": max_gen_len, # Change the parameter name
|
|
}
|
|
else:
|
|
# If no, use the original parameter name 'max_gen_len'
|
|
body = {
|
|
"prompt": prompt,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"max_gen_len": max_gen_len,
|
|
}
|
|
|
|
# Invoke the Bedrock Runtime model with the specified parameters
|
|
response = self.bedrock_runtime_client.invoke_model(
|
|
modelId=model_id, body=json.dumps(body)
|
|
)
|
|
|
|
# Parse the response body from JSON
|
|
response_body = json.loads(response["body"].read())
|
|
logger.info(f"Response Body from Bedrock LLM: {response_body}")
|
|
logger.info("Bedrock LLM model invoked successfully.")
|
|
# Return the response body
|
|
return response_body
|
|
|
|
except ClientError as e:
|
|
# Log an error if there is an issue invoking the model
|
|
logger.error(f"Problem in invoking model. Error: {str(e)}")
|
|
raise
|
|
|
|
class ThrottlingException(Exception):
|
|
pass
|