Cost Attribution for Generative AI with Amazon Bedrock Inference Profiles

The Challenge of Generative AI Cost Attribution

In enterprises adopting generative AI, a common problem arises when multiple internal teams share a single foundation model, especially through services like Amazon Bedrock. Traditional AWS billing often aggregates usage under one account or IAM role, making it difficult to discern which department is consuming the most resources. This lack of granular visibility hinders accurate chargebacks, budget setting, and identification of high-spending units, which is crucial for financial governance and resource optimization.

Solution: Amazon Bedrock Application Inference Profiles

The proposed solution leverages Amazon Bedrock's application inference profiles. An inference profile acts as a logical wrapper around a foundation model. The key architectural design is to associate each department or team with its own inference profile, which in turn is linked to specific AWS cost allocation tags (e.g., `Team=HR`). The application layer is responsible for identifying the user's department and routing their AI inference requests through the corresponding tagged inference profile. All profiles can point to the same underlying foundation model.

Architectural Flow for Cost Tracking

1. User authenticates at the application layer.
2. Application identifies the user's department.
3. Application routes the request to the department's specific, tagged inference profile in Amazon Bedrock.
4. All inference profiles utilize the same foundation model.
5. Amazon Bedrock records usage against each profile's `Team` tag.
6. AWS Cost Explorer displays costs grouped by department based on these tags.

import boto3
from botocore.exceptions import ClientError

client = boto3.client('bedrock-runtime', region_name='us-east-1')

# Map departments to their respective inference profile ARNs
DEPARTMENT_PROFILES = {
    'HR': 'arn:aws:bedrock:us-east-1:111122223333:application-inference-profile/abc123',
    'Accounting': 'arn:aws:bedrock:us-east-1:111122223333:application-inference-profile/def456',
    'IT': 'arn:aws:bedrock:us-east-1:111122223333:application-inference-profile/ghi789',
}

# Function to determine user's department (e.g., from session, OIDC/SAML claim)
def get_department_from_user_session():
    # Placeholder: Replace with actual logic to get the department
    return 'HR' 

department = get_department_from_user_session()

if department not in DEPARTMENT_PROFILES:
    raise ValueError(f"Unknown department: {department}")

try:
    response = client.converse(
        modelId=DEPARTMENT_PROFILES[department], # Use inference profile ARN as modelId
        messages=[{'role': 'user', 'content': [{'text': 'Your prompt here'}]}],
        inferenceConfig={'maxTokens': 300}
    )
    print(response)
except ClientError as e:
    print(f"Error invoking model: {e}")
    raise

This approach ensures that while the underlying AI model remains shared, the consumption metrics are disaggregated, providing a clear financial picture for each consuming team. The modification to the application is minimal, primarily involving passing the inference profile ARN instead of the raw model ID to the Bedrock API call.