scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(bedrock_mantle): price GovCloud regions from the regional cost row and accept region-prefixed model names

Browse source
This commit is contained in:
mateo-berri 2026-09-04 18:22:04 -07:00
parent 59d42d36e6
commit ec3a5c793c
8 changed files with 250 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/cost_calculator.py
View file

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

3
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -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"

15
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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(

6
litellm/llms/bedrock_mantle/chat/transformation.py
View file

@ -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

12
litellm/llms/bedrock_mantle/common_utils.py
View file

@ -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"

75
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -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

78
tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_transformation.py
View file

@ -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."""

53
tests/test_litellm/test_cost_calculator.py
View file

@ -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/<region>/<model> 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/<region>/<model> 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."""