Call multiple Amazon Bedrock model offerings using Python asyncio and the Bedrock REST API

asyncio_bedrock_RESTapi.py

## Note: This is for titan-express, but for any model you use, you will bring your own payloads and simply run this code.

import json
import time
import asyncio
import requests as req
import botocore.session
from itertools import groupby
from operator import itemgetter
from botocore.auth import SigV4Auth
from typing import Dict, List, Tuple
from botocore.awsrequest import AWSRequest

## ------ Initialize constants to invoke the general async function to call REST APIs for Bedrock ------------
SERVICE_NAME: str = 'bedrock'

## ------ Utility function, see: https://stackoverflow.com/questions/31887447/how-do-i-merge-two-lists-of-tuples-based-on-a-key

def inner_join(a, b):
    L = a + b
    L.sort(key=itemgetter(0)) # sort by the first column
    for _, group in groupby(L, itemgetter(0)):
        row_a, row_b = next(group), next(group, None)
        if row_b is not None: # join
            yield row_a + row_b[1:] # cut 1st column from 2nd row

## ------------------------------------------------------------------------------------------------------------

def get_inference(model_id: str, region: str, payload: List) -> Tuple:
    print(f"making an inference request to {model_id}, payload={payload}")
    try:
        ## Initialize the runtime rest API to be called for the endpoint
        endpoint: str = f"https://{SERVICE_NAME}-runtime.{region}.amazonaws.com/model/{model_id}/invoke"

        # Converting the payload dictionary into a JSON-formatted string to be sent in the HTTP request
        request_body = json.dumps(payload[1])

        # Creating an AWSRequest object for a POST request with the service specified endpoint, JSON request body, and HTTP headers
        request = AWSRequest(method='POST',
                             url=endpoint,
                             data=request_body,
                             headers={'content-type': 'application/json'})

        # Initializing a botocore session
        session = botocore.session.Session()

        # Adding a SigV4 authentication information to the AWSRequest object, signing the request
        sigv4 = SigV4Auth(session.get_credentials(), SERVICE_NAME, region)
        sigv4.add_auth(request)

        # Prepare the request by formatting it correctly
        prepped = request.prepare()

        # Send the HTTP POST request to the prepared URL with the specified headers & JSON-formatted request body, storing the response
        response = req.post(prepped.url, headers=prepped.headers, data=request_body)

        if response.status_code == 200:
            return (payload[0], response.json())
        else:
            print(f"Error: Received status code {response.status_code}, Response: {response.text}")
            return None
    except Exception as e:
        print(f"Exception occurred: {e}")
        return None


async def async_calls_on_model(model_id, region, payload):
    return await asyncio.to_thread(get_inference, model_id, region, payload)

# Asynchronously calling all of the prompts based on user input on the specific model offering with the given payload 
async def parallel_calls(model_id, region, payloads):
    responses = await asyncio.gather(*[async_calls_on_model(model_id, region, payload) for payload in payloads])
    # the responses would usually not be in the same order as the payloads, so join the the payload and response  
    responses = list(inner_join(payloads, responses))
    return responses

# Function to create the payload
def create_payload(prompt: str) -> Dict:
    return {"inputText": prompt}

# Region and model id to use
region: str = 'us-east-1'
model_id: str = 'amazon.titan-text-express-v1'

## Initialize the number of prompts you want to invoke on the bedrock specific model
prompts = ["what is Amazon?", "what is a cat?", "what is an phone?", "what is a strawberry?", "what is a pencil?"]
payloads: List = [(i, create_payload(p)) for i, p in enumerate(prompts)]

# Start timing before sending the request
print(f"going to make {len(prompts)} requests concurrently")
start_time = time.perf_counter()
responses = await parallel_calls(model_id, region, payloads)
# Calculate the elapsed time
elapsed_time = time.perf_counter() - start_time
print(f"Total time taken for {len(prompts)} calls made concurrently: {elapsed_time:.2f} seconds")

for r in responses:
    print(f"prompt={r[1]}, response={json.dumps(r[2], indent=2)}")

## -------------------------- SAMPLE RESPONSE: Using 'titan-text-express-v1' on Bedrock --------------------------------

""" 
going to make 5 requests concurrently
making an inference request to amazon.titan-text-express-v1, payload=(0, {'inputText': 'what is Amazon?'})
making an inference request to amazon.titan-text-express-v1, payload=(1, {'inputText': 'what is a cat?'})
making an inference request to amazon.titan-text-express-v1, payload=(2, {'inputText': 'what is an phone?'})
making an inference request to amazon.titan-text-express-v1, payload=(3, {'inputText': 'what is a strawberry?'})
making an inference request to amazon.titan-text-express-v1, payload=(4, {'inputText': 'what is a pencil?'})
Total time taken for 5 calls made concurrently: 5.28 seconds
prompt={'inputText': 'what is Amazon?'}, response={
  "inputTextTokenCount": 4,
  "results": [
    {
      "tokenCount": 98,
      "outputText": "\nAmazon is an international technology business that specializes in cloud computing, online advertising, digital streaming, and artificial intelligence. It has grown to become one of the world's largest and most influential companies, with a market capitalization of over $3 trillion. Amazon was founded by Jeff Bezos in 1994, initially as an online bookstore, and has since expanded to offer a wide range of products and services, including Amazon Web Services, Amazon Prime, Amazon Music, and Amazon Echo.",
      "completionReason": "FINISH"
    }
  ]
}
prompt={'inputText': 'what is a cat?'}, response={
  "inputTextTokenCount": 5,
  "results": [
    {
      "tokenCount": 128,
      "outputText": "\nA cat is a small, furry domestic mammal that is commonly kept as a pet. Cats are known for their agility, grace, and playful nature and are valued for their companionship and ability to control pests such as rodents. There are over 100 different breeds of cats, each with its own unique characteristics and personality. Cats are carnivores and require a diet that is high in protein and fat to maintain their health. They are typically independent and require minimal attention, but they can also be very affectionate and loyal pets. Cats are known for their ability to climb and jump and can often be found perched on high surfaces,",
      "completionReason": "LENGTH"
    }
  ]
}
prompt={'inputText': 'what is an phone?'}, response={
  "inputTextTokenCount": 5,
  "results": [
    {
      "tokenCount": 87,
      "outputText": "\nA phone is a device that allows people to make and receive calls over a network, typically using a telephone number. It has various features such as a screen, keypad, and microphone for communication. Phones have evolved over time to include advanced technologies like internet connectivity, mobile apps, and cameras. They are widely used for personal and business purposes, enabling voice calls, text messaging, video conferencing, and accessing information from the internet.",
      "completionReason": "FINISH"
    }
  ]
}
prompt={'inputText': 'what is a strawberry?'}, response={
  "inputTextTokenCount": 5,
  "results": [
    {
      "tokenCount": 128,
      "outputText": "\n\"A strawberry (or Fragaria \u00d7 ananassa) is a widely grown hybrid species of the genus Fragaria. The fruit is known for its characteristic aroma, bright red color, juicy texture, and sweetness. It is widely grown and consumed in large quantities worldwide, and is often used in desserts, jams, juices, and other culinary products. Strawberries are a hybrid species of Fragaria chiloensis and Fragaria virginiana, and are members of the rose family, Rosaceae. The first strawberry was bred in Brittany, France, in the 1750s by Jean-Fran\u00e7",
      "completionReason": "LENGTH"
    }
  ]
}
prompt={'inputText': 'what is a pencil?'}, response={
  "inputTextTokenCount": 5,
  "results": [
    {
      "tokenCount": 128,
      "outputText": "\n\"A pencil is a writing implement that is made of a narrow graphite core enclosed in a wooden shaft that is used for writing and drawing. Graphite is used because it is easily formed into a solid stick that can be applied to paper with a light pressure, leaving a dark mark that can easily be erased. Pencils were invented in the 16th century by the Spanish and have been used extensively for writing, drawing, and marking ever since.\n\nThe most common type of pencil is the #2 pencil, which has a solid graphite core and a cylindrical wooden shaft. Other types of pencils include the #1 pencil,",
      "completionReason": "LENGTH"
    }
  ]
}
"""
## -----------------------------------------------------------------------------------------------------------------------