diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cb7da32f857..376523a0e9b 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -465,7 +465,8 @@ def cost_per_token( else: model_with_provider = f"{custom_llm_provider}/{model}" if region_name is not None: - model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{model}" + bare_model: Final = model[len(_prov_prefix) :] if model_is_str and model.startswith(_prov_prefix) else model + model_with_provider_and_region: Final = f"{custom_llm_provider}/{region_name}/{bare_model}" if model_with_provider_and_region in model_cost_ref: # use region based pricing, if it's available model_with_provider = model_with_provider_and_region else: @@ -754,6 +755,7 @@ def _select_model_name_for_cost_calc( custom_pricing: bool | None = None, custom_llm_provider: str | None = None, router_model_id: str | None = None, + region_name: str | None = None, ) -> str | None: """ 1. If custom pricing is true, return received model name @@ -775,8 +777,8 @@ def _select_model_name_for_cost_calc( provider_response_model: Final = _get_hidden_str_for_cost_calc(hidden_params, "provider_response_model") explicit_pricing: Final = custom_pricing is True or base_model is not None priced_from_response: Final = provider_response_model is not None or completion_response_model is not None - region_name: Final = ( - _get_hidden_str_for_cost_calc(hidden_params, "region_name") + priced_region: Final = ( + _get_hidden_str_for_cost_calc(hidden_params, "region_name") or region_name if not explicit_pricing and priced_from_response else None ) @@ -813,8 +815,10 @@ def _select_model_name_for_cost_calc( and custom_llm_provider is not None and not _model_contains_known_llm_provider(return_model) ): # add provider prefix if not already present, to match model_cost - provider_prefix: Final = custom_llm_provider if region_name is None else f"{custom_llm_provider}/{region_name}" - return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", region_name) + provider_prefix: Final = ( + custom_llm_provider if priced_region is None else f"{custom_llm_provider}/{priced_region}" + ) + return_model = _strip_unregistered_leading_segments(f"{provider_prefix}/{return_model}", priced_region) return return_model @@ -1281,6 +1285,7 @@ def completion_cost( custom_pricing=custom_pricing, base_model=base_model, router_model_id=router_model_id, + region_name=region_name, ) potential_model_names: Final = [ @@ -1842,6 +1847,7 @@ def response_cost_calculator( data_residency: str | None = None, # for OpenAI regional-processing uplift (e.g. "eu", "us") ### VERTEX LOCATION ### vertex_location: str | None = None, # for Vertex AI regional-endpoint uplift (e.g. "us-east5", "global") + region_name: str | None = None, ) -> float: """ Returns @@ -1875,6 +1881,7 @@ def response_cost_calculator( service_tier=service_tier, data_residency=data_residency, vertex_location=vertex_location, + region_name=region_name, ) return response_cost except Exception as e: diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index ce51fb19970..02681d8b499 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -601,12 +601,15 @@ def _get_openai_compatible_provider_info( dynamic_api_key, ) = litellm.GroqChatConfig()._get_openai_compatible_provider_info(api_base, api_key) elif custom_llm_provider == "bedrock_mantle": + from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix + ( api_base, dynamic_api_key, ) = litellm.BedrockMantleChatConfig()._get_openai_compatible_provider_info( api_base, api_key, litellm_params=litellm_params, model=model ) + model = split_mantle_region_prefix(model)[1] # rebind-ok: the prefix is routing only, not a Mantle model id elif custom_llm_provider == "nvidia_nim": # nvidia_nim is openai compatible, we just need to set this to custom_openai and have the api_base be https://api.endpoints.anyscale.com/v1 api_base = api_base or get_secret("NVIDIA_NIM_API_BASE") or "https://integrate.api.nvidia.com/v1" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index df579f6df5b..088dea4d78d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -414,6 +414,17 @@ def _resolve_vertex_location_for_cost( return VertexBase.get_vertex_region(configured_location, model) +def _resolve_mantle_region_for_cost( + custom_llm_provider: str | None, + litellm_params: Mapping[str, object] | None, +) -> str | None: + if custom_llm_provider != "bedrock_mantle": + return None + from litellm.llms.bedrock_mantle.common_utils import resolve_mantle_region + + return resolve_mantle_region(litellm_params or MappingProxyType({})) + + def _provider_response_id(source: object) -> str | None: candidate: Final = source.get("id") if isinstance(source, dict) else getattr(source, "id", None) return candidate if isinstance(candidate, str) and candidate else None @@ -1711,6 +1722,10 @@ class Logging(LiteLLMLoggingBaseClass): optional_params=self.optional_params, model=litellm_model_name or self.model, ), + "region_name": _resolve_mantle_region_for_cost( + custom_llm_provider=self.model_call_details.get("custom_llm_provider", None), + litellm_params=self.model_call_details.get("litellm_params"), + ), } except Exception as e: # error creating kwargs for cost calculation debug_info = StandardLoggingModelCostFailureDebugInformation( diff --git a/litellm/llms/bedrock_mantle/chat/transformation.py b/litellm/llms/bedrock_mantle/chat/transformation.py index 64d7ef2bed6..31edeb7ec51 100644 --- a/litellm/llms/bedrock_mantle/chat/transformation.py +++ b/litellm/llms/bedrock_mantle/chat/transformation.py @@ -25,7 +25,7 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams from ...openai_like.chat.transformation import OpenAILikeChatConfig -from ..common_utils import mantle_base_segment +from ..common_utils import mantle_base_segment, split_mantle_region_prefix class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): @@ -52,8 +52,10 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): litellm_params: GenericLiteLLMParams | None = None, model: str | None = None, ) -> tuple[str | None, str | None]: + prefix_region, base_model = split_mantle_region_prefix(model) if model else (None, None) region: Final = ( (litellm_params.aws_region_name if litellm_params else None) + or prefix_region or get_secret_str("BEDROCK_MANTLE_REGION") or get_secret_str("AWS_REGION_NAME") or get_secret_str("AWS_REGION") @@ -66,7 +68,7 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig): api_base = ( api_base or get_secret_str("BEDROCK_MANTLE_API_BASE") - or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(model, litellm.model_cost)}" + or f"https://bedrock-mantle.{region}.api.aws/{mantle_base_segment(base_model, litellm.model_cost)}" ) dynamic_api_key: Final = self._resolve_bearer_token(api_key) return api_base, dynamic_api_key diff --git a/litellm/llms/bedrock_mantle/common_utils.py b/litellm/llms/bedrock_mantle/common_utils.py index 850738bc320..952642d3cd9 100644 --- a/litellm/llms/bedrock_mantle/common_utils.py +++ b/litellm/llms/bedrock_mantle/common_utils.py @@ -24,6 +24,7 @@ from botocore.exceptions import ( ) from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.common_utils import _get_all_bedrock_regions from litellm.secret_managers.main import get_secret_str BEDROCK_MANTLE_DEFAULT_REGION: Final = "us-east-1" @@ -36,6 +37,13 @@ def resolve_mantle_bearer_token(api_key: str | None) -> str | None: return api_key or get_secret_str("BEDROCK_MANTLE_API_KEY") or get_secret_str("AWS_BEARER_TOKEN_BEDROCK") +def split_mantle_region_prefix(model: str) -> tuple[str | None, str]: + head, sep, tail = model.partition("/") + if sep and head in _get_all_bedrock_regions(): + return head, tail + return None, model + + def resolve_mantle_region(params: Mapping[str, object]) -> str: region: Final = params.get("aws_region_name") if isinstance(region, str) and region: @@ -130,7 +138,7 @@ def mantle_supports_responses(model: str | None, model_cost: dict) -> bool: gpt-oss substring), so a substring gate would be wrong. A model absent from model_cost simply has no signal and returns False (chat-completions emulation). """ - entry: Final = model_cost.get(f"bedrock_mantle/{model}", {}) + entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {} if "/v1/responses" in (entry.get("supported_endpoints") or []): return True return entry.get("mode") == "responses" @@ -147,5 +155,5 @@ def mantle_base_segment(model: str | None, model_cost: dict) -> str: the base for the model's whole OpenAI-compatible surface, so both the chat and responses configs derive from it -- there is no separate model-name rule. """ - entry: Final = model_cost.get(f"bedrock_mantle/{model}", {}) + entry: Final = model_cost.get(f"bedrock_mantle/{split_mantle_region_prefix(model)[1]}", {}) if model else {} return "openai/v1" if entry.get("use_openai_responses_path") is True else "v1" diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index f1de7390b5b..60eba8a5497 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -5629,6 +5629,81 @@ def test_resolve_vertex_location_for_cost_default_region(monkeypatch): assert _resolve("vertex_ai", None, None, "gemini-3.5-flash") == "us-central1" +def test_resolve_mantle_region_for_cost(monkeypatch): + """Bedrock Mantle requests resolve the served region the way dispatch does (explicit + aws_region_name, then the api_base host, then the default); other providers get None.""" + from litellm.litellm_core_utils.litellm_logging import _resolve_mantle_region_for_cost + + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION_NAME", "AWS_REGION"): + monkeypatch.delenv(var, raising=False) + + assert _resolve_mantle_region_for_cost("bedrock", {"aws_region_name": "us-gov-west-1"}) is None + assert _resolve_mantle_region_for_cost(None, {"aws_region_name": "us-gov-west-1"}) is None + assert _resolve_mantle_region_for_cost("bedrock_mantle", {"aws_region_name": "us-gov-west-1"}) == "us-gov-west-1" + assert ( + _resolve_mantle_region_for_cost( + "bedrock_mantle", + {"api_base": "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"}, + ) + == "us-gov-west-1" + ) + assert _resolve_mantle_region_for_cost("bedrock_mantle", None) == "us-east-1" + + +def test_response_cost_calculator_prices_mantle_calls_on_the_served_region(monkeypatch): + """ + Mantle responses carry no region of their own (the OpenAI-compatible transform rebuilds the + response, and streams never had one), so the logging layer must price them from the region + the deployment was served in: an explicit aws_region_name or the api_base host, both of which + must select the GovCloud row over the commercial one. + """ + from datetime import datetime + + from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") + monkeypatch.setattr(litellm, "model_cost", get_model_cost_map(url="")) + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION_NAME", "AWS_REGION"): + monkeypatch.delenv(var, raising=False) + + def cost_with(litellm_params): + logging_obj = LitellmLogging( + model="xai.grok-4.3", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=datetime.now(), + litellm_call_id="mantle-region", + function_id="f", + ) + logging_obj.update_environment_variables( + model="xai.grok-4.3", + user="", + optional_params={}, + litellm_params=litellm_params, + custom_llm_provider="bedrock_mantle", + ) + response = ModelResponse( + id="resp-1", + model="xai.grok-4.3", + choices=[{"message": {"role": "assistant", "content": "hello"}, "index": 0, "finish_reason": "stop"}], + usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + ) + return logging_obj._response_cost_calculator(result=response) + + commercial = litellm.model_cost["bedrock_mantle/xai.grok-4.3"] + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + expected_commercial = 38 * commercial["input_cost_per_token"] + 20 * commercial["output_cost_per_token"] + expected_gov = 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + assert expected_gov != expected_commercial + + assert cost_with({"api_base": ""}) == pytest.approx(expected_commercial) + assert cost_with({"aws_region_name": "us-gov-west-1"}) == pytest.approx(expected_gov) + assert cost_with( + {"api_base": "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions"} + ) == pytest.approx(expected_gov) + + def test_response_cost_calculator_prices_proxy_vertex_calls_on_the_configured_location(monkeypatch): """ Proxy-shaped logging objects (created before the router picks a deployment) carry the diff --git a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index b88c27e64b9..ee81ec482b0 100644 --- a/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -146,6 +146,27 @@ class TestBedrockMantleConfig: # /openai/v1 base per the AWS model card. assert api_base == "https://bedrock-mantle.us-east-2.api.aws/openai/v1" + def test_region_prefixed_model_routes_to_that_region(self, monkeypatch, local_cost_map): + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(var, raising=False) + cfg = BedrockMantleChatConfig() + api_base, _ = cfg._get_openai_compatible_provider_info(None, None, model="us-gov-west-1/xai.grok-4.3") + assert api_base == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1" + + def test_aws_region_name_param_beats_model_region_prefix(self, monkeypatch, local_cost_map): + from litellm.types.router import GenericLiteLLMParams + + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(var, raising=False) + cfg = BedrockMantleChatConfig() + api_base, _ = cfg._get_openai_compatible_provider_info( + None, + None, + litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1"), + model="us-gov-west-1/xai.grok-4.3", + ) + assert api_base == "https://bedrock-mantle.us-east-1.api.aws/openai/v1" + def test_default_api_base_fallback_to_us_east_1(self, monkeypatch): monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False) monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) @@ -681,6 +702,63 @@ class TestBedrockMantleProviderResolution: assert model == "openai.gpt-oss-20b" + def test_get_llm_provider_strips_region_prefix(self, monkeypatch, local_cost_map): + for var in ("BEDROCK_MANTLE_REGION", "BEDROCK_MANTLE_API_BASE", "AWS_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(var, raising=False) + model, provider, _, api_base = litellm.get_llm_provider("bedrock_mantle/us-gov-west-1/xai.grok-4.3") + assert provider == "bedrock_mantle" + assert model == "xai.grok-4.3" + assert api_base == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1" + + def test_completion_region_prefixed_model_sends_bare_model_to_that_region(self, monkeypatch, local_cost_map): + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + for var in ( + "BEDROCK_MANTLE_API_KEY", + "AWS_BEARER_TOKEN_BEDROCK", + "BEDROCK_MANTLE_API_BASE", + "BEDROCK_MANTLE_REGION", + "AWS_REGION_NAME", + "AWS_REGION", + "AWS_PROFILE", + ): + monkeypatch.delenv(var, raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0") + + requests = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1733529600, + "model": "xai.grok-4.3", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + }, + request=request, + ) + + response = litellm.completion( + model="bedrock_mantle/us-gov-west-1/xai.grok-4.3", + messages=[{"role": "user", "content": "hello"}], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + sent = requests[0] + assert str(sent.url) == "https://bedrock-mantle.us-gov-west-1.api.aws/openai/v1/chat/completions" + assert json.loads(sent.content)["model"] == "xai.grok-4.3" + assert "/us-gov-west-1/bedrock/aws4_request" in sent.headers["Authorization"] + assert response._hidden_params["response_cost"] == pytest.approx( + 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + ) + + class TestBedrockMantlePricing: """Tests that verify Bedrock Mantle uses correct AWS Bedrock pricing, not OpenAI pricing.""" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 2046695f151..ca2549f60a8 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4142,6 +4142,59 @@ def test_select_model_name_applies_region_to_private_provider_response_model(_lo assert selected == "bedrock/us-east-1/anthropic.claude-v2:1" +def test_completion_cost_region_name_prices_mantle_on_the_regional_row(_local_model_cost_map): + """completion_cost(region_name=...) must price a Bedrock Mantle call from the + bedrock_mantle// row when one exists, for the bare and the provider-prefixed + model alike, and keep the flat row for regions without their own row.""" + + response = litellm.ModelResponse( + id="x", + choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + model="xai.grok-4.3", + usage={"prompt_tokens": 38, "completion_tokens": 20, "total_tokens": 58}, + ) + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + flat = litellm.model_cost["bedrock_mantle/xai.grok-4.3"] + expected_gov = 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + expected_flat = 38 * flat["input_cost_per_token"] + 20 * flat["output_cost_per_token"] + assert expected_gov != expected_flat + + for model in ("xai.grok-4.3", "bedrock_mantle/xai.grok-4.3"): + assert litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock_mantle", + region_name="us-gov-west-1", + ) == pytest.approx(expected_gov) + assert litellm.completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock_mantle", + region_name="eu-west-1", + ) == pytest.approx(expected_flat) + assert litellm.completion_cost( + completion_response=response, model="xai.grok-4.3", custom_llm_provider="bedrock_mantle" + ) == pytest.approx(expected_flat) + + +def test_cost_per_token_region_name_applies_to_provider_prefixed_model(_local_model_cost_map): + """A provider-prefixed model must still find its bedrock_mantle// row instead of + composing the region key with the provider segment twice.""" + + prompt_cost, completion_cost = litellm.cost_per_token( + model="bedrock_mantle/xai.grok-4.3", + prompt_tokens=38, + completion_tokens=20, + custom_llm_provider="bedrock_mantle", + region_name="us-gov-west-1", + ) + gov = litellm.model_cost["bedrock_mantle/us-gov-west-1/xai.grok-4.3"] + + assert prompt_cost + completion_cost == pytest.approx( + 38 * gov["input_cost_per_token"] + 20 * gov["output_cost_per_token"] + ) + + def test_select_model_name_keeps_base_model_free_of_region(_local_model_cost_map): """An explicit base_model keeps pricing on that model's own key even when the request carries a region with different regional rates, so the private provider model never widens region pricing."""