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):