Amazon Bedrock... let's go raw! 🥩

bedrock_raw.py

# pylint: disable=line-too-long
"""
## Amazon Bedrock... let's go raw! 🥩

Your mission (if you choose to accept it) is to create a bedrock.invokeModel request from the ground up.

In a nutshell, we have to prepare the request, SigV4-sign it and then send it to the Bedrock Runtime endpoint.

Sounds fun? Then, let's get started...

⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣠⡶⠶⣦⣄⠀⠀⢀⣴⣿⣷⡄
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣠⣴⣿⣿⡇⡖⠂⠙⠗⣠⣾⣿⣿⣿⣥⣀
⠀⠀⠀⠀⠀⠀⠀⢀⣀⣠⣤⣶⣿⣿⣿⣿⣿⣇⢣⠀⣠⣾⣿⣿⣿⣿⣿⣿⣿⠇
⠀⠀⠀⠀⠀⠀⢰⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡌⢧⠘⠿⠟⠛⣉⠉
⠀⠀⠀⠀⠀⠀⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣌⠳⣄⠀⠀⣿⡀
⠀⠀⠀⠀⠀⠀⢸⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣷⣌⣉⣁⡿
⠀⢀⣤⣤⣤⣤⡀⢻⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⣿⡿⠟⠛⠉
⠀⢻⣿⣿⣿⣿⣷⡀⠹⣿⣿⣿⣿⣿⣿⣿⣿⠟⠋⠁
⠀⠀⠉⢩⣿⣿⣿⠋⠀⠈⠻⢿⣿⣿⣿⠋⠁
⠀⠀⠀⠸⣿⡿⠁⠀⠀⠀⠀⠀⠈⠉⠁


## References 📚

AWS IAM User Guide > AWS Signature Version 4 for API requests
https://docs.aws.amazon.com/IAM/latest/UserGuide/reference_aws-signing.html

Amazon Bedrock API Reference > Amazon Bedrock Runtime > InvokeModel
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html

## Other Implementations 👾

[botocore]
https://github.com/boto/botocore/blob/develop/botocore/auth.py
https://github.com/boto/botocore/blob/develop/botocore/data/bedrock-runtime/2023-09-30/service-2.json

[awscurl]
https://github.com/okigan/awscurl/blob/master/awscurl/awscurl.py
"""

import os
import logging
import hmac

import http.client as http_client
import urllib.parse as urllib_parse

from datetime import (
    datetime,
    timezone
)

import hashlib
import requests

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

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

# Initialize logging
logging.basicConfig()
logging.getLogger().setLevel(logging.DEBUG)

requests_log = logging.getLogger("requests.packages.urllib3")
requests_log.setLevel(logging.DEBUG)
requests_log.propagate = True

####################
# Helper functions #
####################

# The algorithm used to create the hash of the canonical request
ALGORITHM = "AWS4-HMAC-SHA256"

def lowercase(text):
    """
    Convert the string to lowercase.
    """
    return text.lower()

def trim(text):
    """
    Remove any leading or trailing whitespace.
    """
    return text.strip()

def uri_encode(text, safe='/'):
    """
    URI encode every byte.
    """
    return urllib_parse.quote_plus(text, safe=safe)

def hex_sha256_hash(text):
    """
    Secure Hash Algorithm (SHA) cryptographic hash function with lowercase base 16 encoding.
    """
    return hashlib.sha256(text).hexdigest()

def hmac_sha256(key, msg, hex=False):  # pylint: disable=redefined-builtin
    """
    Computes HMAC on a message by using the SHA256 algorithm with the provided signing key.
    """
    if hex:
        sig = hmac.new(key, msg.encode('utf-8'), hashlib.sha256).hexdigest()
    else:
        sig = hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()
    return sig

#################
# Signing Steps #
#################

def create_canonical_request(  # pylint: disable=too-many-arguments
        http_verb,
        canonical_uri,
        canonical_query_string,
        canonical_headers,
        signed_headers,
        hashed_payload):
    """
    Arranges the contents of the request (host, action, headers, &c.) into a standard canonical format.
    """
    canonical_request = http_verb + "\n" + \
                        canonical_uri + "\n" + \
                        canonical_query_string + "\n" + \
                        canonical_headers + "\n" + \
                        signed_headers + "\n" + \
                        hashed_payload

    return canonical_request

def derive_signing_key(key, date, region, service):
    """
    Derives a signing key by performing a succession of keyed hash operations (HMAC operations)
    on the request date, Region, and service, with the AWS secret access key as the key for the
    initial hashing operation.
    """
    date_key = hmac_sha256(("AWS4" + key).encode('utf-8'), date)
    date_region_key = hmac_sha256(date_key, region)
    date_region_service_key = hmac_sha256(date_region_key, service)
    signing_key = hmac_sha256(date_region_service_key, "aws4_request")
    return signing_key

def create_string_to_sign(timestamp, credential_scope, canonical_request):
    """
    Creates a string to sign with the canonical request and extra information such as the algorithm, 
    request date, credential scope, and the hash of the canonical request.
    """
    string_to_sign = ALGORITHM + "\n" + \
                     timestamp + "\n" + \
                     credential_scope + "\n" + \
                     hex_sha256_hash(canonical_request.encode('utf-8'))

    return string_to_sign

def calculate_signature(signing_key, string_to_sign):
    """
    Calculates the signature by performing a keyed hash operation on the string to sign.
    """
    signature = hmac_sha256(signing_key, string_to_sign, hex=True)
    return signature

def create_authorization_header(credential, scope, signed_headers, signature):
    """
    Creates the Authorization header for the request.
    """
    return ALGORITHM + " " + "Credential" + "=" + credential + "/" + scope + ", " \
                             "SignedHeaders" + "=" + signed_headers + ", " \
                             "Signature" + "=" + signature

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

def main():  # pylint: disable=too-many-locals
    """
    Main entrypoint.
    """

    # The keys to the kingdom 🗝️
    #
    # Note: when in doubt, use *temporary* credentials
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_request.html
    aws_access_key_id = os.environ['AWS_ACCESS_KEY_ID']
    aws_secret_access_key = os.environ['AWS_SECRET_ACCESS_KEY']
    aws_session_token = os.environ.get('AWS_SESSION_TOKEN')

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

    # The service ⚙️
    # https://aws.amazon.com/bedrock/
    service_name = "bedrock"

    # The model 🧠
    # https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html
    model_id = uri_encode("amazon.titan-embed-text-v2:0")

    # The destination AKA endpoint 🎯
    # https://docs.aws.amazon.com/general/latest/gr/bedrock.html
    # Note: endpoint prefix != service name
    host = f"bedrock-runtime.{region}.amazonaws.com"
    endpoint = f"https://{host}/model/{model_id}/invoke"

    # The request body AKA payload 📩
    # https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-titan-embed-text.html
    payload = '{"inputText": "Where\'s the love?"}'

    # The request date and time 📅🕒
    # Note: current UTC time in ISO 8601 format
    current_date = datetime.now(timezone.utc)
    timestamp = current_date.strftime(r"%Y%m%dT%H%M%SZ")
    datestamp = current_date.strftime(r"%Y%m%d")

    # STEP 1. Create a canonical request
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html#create-canonical-request

    # The HTTP method
    http_verb = "POST"

    # The URI-encoded version of the absolute path component URI
    canonical_uri = uri_encode(f"/model/{model_id}/invoke")

    # The URI-encoded query string parameters
    canonical_query_string = ""

    # A list of request headers with their values
    canonical_headers = lowercase("Host") + ":" + trim(host) + "\n" + \
                        lowercase("X-Amz-Date") + ":" + trim(timestamp) + "\n"

    # An alphabetically sorted, semicolon-separated list of lowercase request header names
    signed_headers = lowercase("Host") + ";" + lowercase("X-Amz-Date")

    # When using temporary credentials, we must include the security token
    # https://docs.aws.amazon.com/STS/latest/APIReference/CommonParameters.html
    if aws_session_token:
        canonical_headers += lowercase("X-Amz-Security-Token") + ":" + trim(aws_session_token) + "\n"
        signed_headers += ";" + lowercase("X-Amz-Security-Token")

    # A string created using the payload in the body of the HTTP request as input to the hash function
    hashed_payload = hex_sha256_hash(payload.encode('utf-8'))

    # We create the canonical request by concatenating the strings above, separated by newline characters
    canonical_request = create_canonical_request(
        http_verb,
        canonical_uri,
        canonical_query_string,
        canonical_headers,
        signed_headers,
        hashed_payload
    )

    print(f"\nCanonical Request\n-----\n{canonical_request}\n-----\n")

    # STEP 2. Create string to sign
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html#create-string-to-sign

    # The credential scope restricts the resulting signature to the specified Region and service
    credential_scope = f"{datestamp}/{region}/{service_name}/aws4_request"

    string_to_sign = create_string_to_sign(
        timestamp,
        credential_scope,
        canonical_request
    )

    print(f"\nString to Sign\n-----\n{string_to_sign}\n-----\n")

    # STEP 3. Derive a signing key
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html#derive-signing-key

    signing_key = derive_signing_key(
        aws_secret_access_key,
        datestamp,
        region,
        service_name
    )

    print(f"\n-----\n{signing_key}\n-----\n")

    # STEP 4. Calculate the signature
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html#calculate-signature

    signature = calculate_signature(
        signing_key,
        string_to_sign
    )

    print(f"\nSignature\n-----\n{signature}\n-----\n")

    # STEP 5. Add the signature to the request
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html#add-signature-to-request

    authorization = create_authorization_header(
        aws_access_key_id,
        credential_scope,
        signed_headers,
        signature
    )

    print(f"\nAuthorization\n-----\n{authorization}\n-----\n")

    # STEP 6. Send the request 🚀

    # Just as a reminder, here's some information about these headers:
    # https://docs.aws.amazon.com/STS/latest/APIReference/CommonParameters.html
    # https://docs.aws.amazon.com/IAM/latest/UserGuide/aws-signing-authentication-methods.html#aws-signing-authentication-methods-http
    headers = {
        "X-Amz-Date": timestamp,
        "Authorization": authorization,
    }

    # Don't forget the session token 🔔
    if aws_session_token:
        headers["X-Amz-Security-Token"] = aws_session_token

    # Finally, let's call the model 📞
    response = requests.request(
        method="POST",
        url=endpoint,
        headers=headers,
        data=payload,
        timeout=30
    )

    # and display the result ✔️
    print("Response: ", response.json())


if __name__ == '__main__':
    main()