Gestionar el acceso a AWS bedrock utilizando AWS Lambda y API Gateway

Desde el lanzamiento y disponibilidad de AWS Bedrock en distintas regiones [1][2], las pruebas de conceptos son necesarias para validar y probar las funcionalidades y caractersticas que ofrece el servicio.

Para agregar una capa de seguridad y lmites a las invocaciones construimos un proxy utilizando AWS lambda y APIGateway en los ambientes de pruebas y testing que realizamos, sobre todo si tenemos que habilitar el servicio para otros sectores dentro de la empresa o a nuestros clientes.

El API key siempre debera ir acompaado de una cuota(Lmite de invocaciones), y en la mayora de los casos para ambientes efmeros o de pruebas; para ms seguridad siempre deberamos agregar Autorizers a nuestras API, como AWS cognito/JWT Customs , etc.

A continuacin un ejemplo de cmo agregar un proxy a AWS bedrock y exponerlo a travs de una API con cuota.

Conectarnos a AWS bedrock desde Lambda

Para conectarnos usaremos Boto3, y nos conectaremos a modo de prueba a dos de los LLM disponibles tras BedRock. Titan(amazon.titan-text-express-v1) y Llama(meta.llama2-70b-chat-v1) .

# src/client_bedrock.pyimport boto3import botocoreimport jsonbedrock_runtime = boto3.client('bedrock-runtime', region_name='us-east-1')

Y luego, un par de mtodos para cada LLM al que queremos conectarnos, para el caso de titan(amazon.titan-text-express-v1):

# src/client_bedrock.pydef titan_model(prompt): body = json.dumps({ "inputText": prompt, "textGenerationConfig":{ "maxTokenCount":4096, "stopSequences":[], "temperature":0, "topP":0.9 } }) modelId = 'amazon.titan-text-express-v1' # change this to use a different version from the model provider accept = 'application/json' contentType = 'application/json' outputText = "\n" text = bedrock_runtime.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType) response_body = json.loads(text.get('body').read()) outputText = response_body.get('results')[0].get('outputText') response = {"statusCode": 200, "body": json.dumps(outputText)} return response

Para llama(meta.llama2-70b-chat-v1) usamos un mtodo similar:

# src/client_bedrock.pydef llama_model(prompt): body = json.dumps({ 'prompt': prompt, 'max_gen_len': 512, 'top_p': 0.9, 'temperature': 0.2 }) modelId = 'meta.llama2-70b-chat-v1' # change this to use a different version from the model provider accept = 'application/json' contentType = 'application/json' outputText = "\n" response = bedrock_runtime.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType) response_body = json.loads(response.get('body').read().decode('utf-8')) outputText = response_body['generation'].strip() response = {"statusCode": 200, "body": json.dumps(outputText)} return response

Y luego, un handler para decidir a que LLM queremos llamar:

# src/handler.pyimport jsonfrom src.client_bedrock import titan_model, llama_modelfrom enum import Enumclass Model(Enum): TITAN = 'titan' LLAMA = 'llama' NONE = 'none' @classmethod def _missing_(cls, value): return cls.NONEdef hello(event, context): q = json.loads(event['body']) model = Model(q['model']) match model: case Model.TITAN: text = titan_model(json.dumps(q['question'])) case Model.LLAMA: text = llama_model(json.dumps(q['question'])) case Model.NONE: text = "Model does not match" response = {"statusCode": 200, "body": json.dumps(text)} return response

El despliegue y gestin de las cuotas

Para el despliegue utilic Serverless Framework , de la configuracin del Serverless Framework importante destacar:

• Configuracin del API key:

  apiGateway: apiKeys: - name: tollmkey description: tollmkey api key # Optional usagePlan: quota: limit: 100 # Limint by moth offset: 2 period: MONTH throttle: burstLimit: 200 rateLimit: 100

En el proceso de deploy Serverless Framework nos muestra como salidas el nombre de la API y la API Key:

Deploying lambda-bedrock-proxy to stage dev (us-east-1) Service deployed to stack lambda-bedrock-proxy-dev (80s)# Value of APIKEYapi keys: tollmkey: YOURAPIKEYHERE - tollmkey api key# Optionalendpoint: GET - https://xxxxxxxxxx.execute-api.us-east-1.amazonaws.com/dev/bedrockfunctions: hello: lambda-bedrock-proxy-dev-hello (3.1 kB)layers: bedrock: arn:aws:lambda:us-east-1:XXXXXXXXX:layer:bedrock-dependencies:4

• Configuracin de la Lambda layer: El runtime para Python de lambda toma una versin de Boto3(1.26) por default, pero para conectarse con bedrock es necesario una versin especfica o superior (boto3>=1.28.57), para esto se configur Boto3 dentro de una layer al igual que el resto de dependencias.

# layers/tools/aws_requirements.txtboto3>=1.28.57awscli>=1.29.57botocore>=1.31.57langchain>=0.0.350

Invocar la API con el API key

Para invocar los modelos el body tiene dos campos: question y model. question es la consulta que le pasaremos al modelo y model el tipo de modelo, para este ejemplo estn configurados dos modelos Titan(amazon.titan-text-express-v1) y Llama(meta.llama2-70b-chat-v1) .

curl --location 'https://YOURID.execute-api.us-east-1.amazonaws.com/dev/bedrock' \--header 'Content-Type: application/json' \--header 'x-api-key: YOURAPIKEY' \--data '{ "question": "What is a dystopia", "model": "llama"}'

Probando el lmite de la cuota

Si superamos el lmite de las cuotas que definimos en el apartado Apigateway/usagePlan del serverless.yml nos indicar la respuesta de la API con un mensaje de que el lmite se ha superado y un cdigo HTTP 429

HTTP/1.1 429 Content-Type: application/jsonContent-Length: 28Connection: closex-amzn-RequestId: 893834e6-5f4c-4bb6-95e0-a5b7f255bf51x-amzn-ErrorType: LimitExceededExceptionx-amz-apigw-id: TvNwHEKSIAMEQiQ=X-Cache: Error from cloudfront{ "message": "Limit Exceeded"}

Codigo fuente:

El ejemplo completo sobre el que escrib este resumen est en el siguiente repositorio: https://github.com/olcortesb/lambda-bedrock-proxy/tree/main

Conclusiones:

• Configuramos un solo endpoint para llamar a distintos modelos que son soportados por AWS Bedrock, se pueden extender agregando la funcin para cada LLM y un campo ms en el Enum

• Se configur una cuota para una API key y se pueden configurar distintas cuotas para distintos equipos, una cuota para cada API key

• Se configur Boto3 para usar con lambda las versiones soportadas dentro de una Lambda Layer

• Para ms seguridad se podra habilitar Authorizer a la lambda con AWS Cognito o un JWT Custom

Referencias

[1]. https://aws.amazon.com/es/about-aws/whats-new/2023/09/amazon-bedrock-generally-available/

[2]. https://aws.amazon.com/es/about-aws/whats-new/2023/10/amazon-bedrock-europe-frankfurt-aws-region/