From 5a31792bf1df98412926e6ee6a332de65c342dc5 Mon Sep 17 00:00:00 2001 From: Lakindu Boteju Date: Thu, 12 Dec 2024 09:33:50 +0700 Subject: [PATCH] community: Add support for cross-region inference profile IDs in Bedrock Anthropic Claude token cost calculation (#28167) This change modifies the token cost calculation logic to support cross-region inference profile IDs for Anthropic Claude models. Instead of explicitly listing all regional variants of new inference profile IDs in the cost dictionaries, the code now extracts a base model ID from the input model ID (or inference profile ID), making it more maintainable and automatically supporting new regional variants. These inference profile IDs follow the format: `..` (e.g., `us.anthropic.claude-3-haiku-xxx`, `eu.anthropic.claude-3-sonnet-xxx`). Cross-region inference profiles are system-defined identifiers that enable distributing model inference requests across multiple AWS regions. They help manage unplanned traffic bursts and enhance resilience during peak demands without additional routing costs. References for Amazon Bedrock's cross-region inference profiles:- - https://docs.aws.amazon.com/bedrock/latest/userguide/cross-region-inference.html - https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-support.html --------- Co-authored-by: Erick Friis --- .../callbacks/bedrock_anthropic_callback.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/libs/community/langchain_community/callbacks/bedrock_anthropic_callback.py b/libs/community/langchain_community/callbacks/bedrock_anthropic_callback.py index 32c36f40768..6dd88639cbb 100644 --- a/libs/community/langchain_community/callbacks/bedrock_anthropic_callback.py +++ b/libs/community/langchain_community/callbacks/bedrock_anthropic_callback.py @@ -28,15 +28,25 @@ MODEL_COST_PER_1K_OUTPUT_TOKENS = { def _get_anthropic_claude_token_cost( prompt_tokens: int, completion_tokens: int, model_id: Union[str, None] ) -> float: + if model_id: + # The model ID can be a cross-region (system-defined) inference profile ID, + # which has a prefix indicating the region (e.g., 'us', 'eu') but + # shares the same token costs as the "base model". + # By extracting the "base model ID", by taking the last two segments + # of the model ID, we can map cross-region inference profile IDs to + # their corresponding cost entries. + base_model_id = model_id.split(".")[-2] + "." + model_id.split(".")[-1] + else: + base_model_id = None """Get the cost of tokens for the Claude model.""" - if model_id not in MODEL_COST_PER_1K_INPUT_TOKENS: + if base_model_id not in MODEL_COST_PER_1K_INPUT_TOKENS: raise ValueError( f"Unknown model: {model_id}. Please provide a valid Anthropic model name." "Known models are: " + ", ".join(MODEL_COST_PER_1K_INPUT_TOKENS.keys()) ) - return (prompt_tokens / 1000) * MODEL_COST_PER_1K_INPUT_TOKENS[model_id] + ( + return (prompt_tokens / 1000) * MODEL_COST_PER_1K_INPUT_TOKENS[base_model_id] + ( completion_tokens / 1000 - ) * MODEL_COST_PER_1K_OUTPUT_TOKENS[model_id] + ) * MODEL_COST_PER_1K_OUTPUT_TOKENS[base_model_id] class BedrockAnthropicTokenUsageCallbackHandler(BaseCallbackHandler):