Overview

This guide demonstrates how to use Pay-i's custom headers with AWS Bedrock's Python SDK to annotate your API requests. Due to AWS Bedrock's unique SDK architecture, it requires a special callback registration step before headers can be applied.

Important: Callback Registration Required

Unlike other provider SDKs, AWS Bedrock doesn't directly support custom headers in the same way. Instead, Pay-i provides a special registration function that hooks into Bedrock's request pipeline:

from payi.lib.helpers import register_bedrock_client_callbacks

# Required for Bedrock to properly handle Pay-i annotations
register_bedrock_client_callbacks(bedrock_client, 'InvokeModel')

You must register callbacks for each operation type you intend to use. Common operations include:

InvokeModel
InvokeModelWithResponseStream

Complete Example

Here's a complete example showing the entire process:

import json
import boto3
from payi.lib.helpers import register_bedrock_client_callbacks, create_headers

# Initialize Bedrock client
bedrock = boto3.client('bedrock-runtime')

# STEP 1: Register callbacks for each operation type you'll use
register_bedrock_client_callbacks(bedrock, 'InvokeModel')
register_bedrock_client_callbacks(bedrock, 'InvokeModelWithResponseStream')

# STEP 2: Create request annotations
headers = create_headers(
    use_case_name="text_generation",
    user_id="data_scientist_123",
    limit_ids=["aws_budget"],
    request_tags=["bedrock", "claude"]
)

# STEP 3: Create the request body
request_body = json.dumps({
    "prompt": "Human: Generate a product description for a smart water bottle.\n\nAssistant:",
    "temperature": 0.7,
    "max_tokens_to_sample": 500
})

# STEP 4: Make API call with annotations
response = bedrock.invoke_model(
    modelId="anthropic.claude-3-sonnet-20240229",
    body=request_body,
    extra_headers=headers  # Apply the annotations here
)

# Process the response
response_body = json.loads(response['body'].read())
print(response_body['completion'])

Streaming Example

For streaming responses, the pattern is similar but uses the streaming operation:

import json
import boto3
from payi.lib.helpers import register_bedrock_client_callbacks, create_headers

# Initialize Bedrock client
bedrock = boto3.client('bedrock-runtime')

# Register the streaming operation callback
register_bedrock_client_callbacks(bedrock, 'InvokeModelWithResponseStream')

# Create request annotations
headers = create_headers(
    use_case_name="streaming_generation",
    user_id="analyst_456",
    limit_ids=["streaming_budget"],
    request_tags=["bedrock", "streaming"]
)

# Create the request body
request_body = json.dumps({
    "prompt": "Human: Write a story about artificial intelligence.\n\nAssistant:",
    "temperature": 0.7,
    "max_tokens_to_sample": 1000
})

# Make streaming API call with annotations
response = bedrock.invoke_model_with_response_stream(
    modelId="anthropic.claude-3-sonnet-20240229",
    body=request_body,
    extra_headers=headers  # Apply the annotations here
)

# Process the streaming response
for event in response['body']:
    if 'chunk' in event:
        chunk = json.loads(event['chunk']['bytes'])
        if 'completion' in chunk:
            print(chunk['completion'], end='')

Why Is This Necessary?

This callback registration is necessary because:

The AWS SDK architecture doesn't directly expose header manipulation
Pay-i needs to intercept both request and response data
Different Bedrock operations require separate registrations

Once registered, you can use extra_headers with Bedrock just like with other providers, and Pay-i will handle the rest.

Common Troubleshooting

If you're not seeing tracking data or getting errors:

Missing Registration: Ensure you've registered the specific operation you're using
Registration Order: Register callbacks before making any API calls
Headers Format: Verify that create_headers() is being called correctly

For more information about custom headers and how to use them across different providers, see the Custom Headers documentation.