Call Amazon Bedrock with AWS SigV4 (requests/botocore)

bedrock_boto3_sigv4.py

# pylint: disable=invalid-name
"""
Call Amazon Bedrock via Boto3 with AWS SigV4
"""

import os
import logging

import http.client as http_client

import requests

from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials

#############
# Variables #
#############

# The service name
service = "bedrock"

# The endpoint prefix (=/= service name)
endpoint_prefix = "bedrock-runtime"

# The target AWS region
# https://docs.aws.amazon.com/bedrock/latest/userguide/bedrock-regions.html
region = os.environ.get('AWS_DEFAULT_REGION', "us-east-1")

# The model
# https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
model_id = "amazon.titan-embed-text-v1"

# The endpoint
host = f"{endpoint_prefix}.{region}.amazonaws.com"
endpoint = f"https://{host}/model/{model_id}/invoke"

# The data (payload)
payload = '{"inputText": "love"}'

###########
# Logging #
###########

# Enable debugging at httplib level
# requests -> urllib3 -> http.client
http_client.HTTPConnection.debuglevel = 1

# Initialize logging
# https://requests.readthedocs.io/en/latest/api/#api-changes
logging.basicConfig()
logging.getLogger().setLevel(logging.DEBUG)
requests_log = logging.getLogger("requests.packages.urllib3")
requests_log.setLevel(logging.DEBUG)
requests_log.propagate = True

########
# Main #
########

# Create request
req = AWSRequest(
    method="POST",
    url=endpoint,
    data=payload,
)

# Define credentials
credentials = Credentials(
    os.environ['AWS_ACCESS_KEY_ID'],
    os.environ['AWS_SECRET_ACCESS_KEY'],
    os.environ['AWS_SESSION_TOKEN']
)

# Sign the request with AWS SigV4
SigV4Auth(credentials, service, region).add_auth(req)

# Prepare the request
req = req.prepare()

# and send it
response = requests.request(
    method=req.method,
    url=req.url,
    headers=req.headers,
    data=req.body,
    timeout=30
)