This commit is contained in:
Priyansh Nandwana 2026-08-26 22:50:20 +03:00 committed by GitHub
commit 106b050291
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 115 additions and 67 deletions

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, mantle_omits_max_tokens
class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
@ -103,6 +103,8 @@ class BedrockMantleChatConfig(BedrockMantleAuthMixin, OpenAILikeChatConfig):
base_params.append("reasoning_effort")
except Exception as e:
verbose_logger.debug("BedrockMantleChatConfig: error checking reasoning support: %s", e)
if mantle_omits_max_tokens(model, litellm.model_cost):
return [param for param in base_params if param != "max_tokens"]
return base_params
def get_model_response_iterator(

View file

@ -136,6 +136,21 @@ def mantle_supports_responses(model: str | None, model_cost: dict) -> bool:
return entry.get("mode") == "responses"
def mantle_omits_max_tokens(model: str | None, model_cost: dict) -> bool:
"""Whether a Bedrock Mantle model rejects the OpenAI ``max_tokens`` param.
Data-driven from the model's price-map omit_max_tokens_param flag (overridable
via register_model / proxy model_info), matching mantle_base_segment. The
google gemma-4-* family carries that flag: Mantle answers max_tokens with
``unsupported_parameter``, and its native-API name max_output_tokens is not a
substitute on this OpenAI-compatible route, so omitting it is the only thing
that succeeds. As above there is deliberately NO model-name match, so a new
model with the same quirk is a JSON change rather than a code change.
"""
entry: Final = model_cost.get(f"bedrock_mantle/{model}", {})
return entry.get("omit_max_tokens_param") is True
def mantle_base_segment(model: str | None, model_cost: dict) -> str:
"""Return the base path segment for a Bedrock Mantle model's OpenAI surface.

View file

@ -49739,6 +49739,7 @@
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"omit_max_tokens_param": true,
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -49758,6 +49759,7 @@
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"omit_max_tokens_param": true,
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -49777,6 +49779,7 @@
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"omit_max_tokens_param": true,
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",

View file

@ -304,6 +304,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
]
supported_endpoints: list[str] | None
use_openai_responses_path: bool | None
omit_max_tokens_param: bool | None
tpm: int | None
rpm: int | None
provider_specific_entry: dict[str, float] | None

View file

@ -49739,6 +49739,7 @@
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"omit_max_tokens_param": true,
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -49758,6 +49759,7 @@
"max_output_tokens": 256000,
"max_tokens": 256000,
"mode": "chat",
"omit_max_tokens_param": true,
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",
@ -49777,6 +49779,7 @@
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"omit_max_tokens_param": true,
"use_openai_responses_path": true,
"supported_endpoints": [
"/v1/chat/completions",

View file

@ -47,14 +47,8 @@ class TestBedrockMantleProviderRegistration:
assert len(litellm.bedrock_mantle_models) > 0
assert "bedrock_mantle/openai.gpt-oss-120b" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-oss-20b" in litellm.bedrock_mantle_models
assert (
"bedrock_mantle/openai.gpt-oss-safeguard-120b"
in litellm.bedrock_mantle_models
)
assert (
"bedrock_mantle/openai.gpt-oss-safeguard-20b"
in litellm.bedrock_mantle_models
)
assert "bedrock_mantle/openai.gpt-oss-safeguard-120b" in litellm.bedrock_mantle_models
assert "bedrock_mantle/openai.gpt-oss-safeguard-20b" in litellm.bedrock_mantle_models
class TestBedrockMantleConfig:
@ -108,9 +102,7 @@ class TestBedrockMantleConfig:
cfg._get_openai_compatible_provider_info(
None,
None,
litellm_params=GenericLiteLLMParams(
aws_region_name="us-east-1.api.aws.attacker.example/"
),
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1.api.aws.attacker.example/"),
)
def test_get_llm_provider_rejects_malicious_aws_region_name(self, monkeypatch):
@ -123,14 +115,10 @@ class TestBedrockMantleConfig:
litellm.get_llm_provider(
model="openai.gpt-5.5",
custom_llm_provider="bedrock_mantle",
litellm_params=GenericLiteLLMParams(
aws_region_name="us-east-1.api.aws.attacker.example/"
),
litellm_params=GenericLiteLLMParams(aws_region_name="us-east-1.api.aws.attacker.example/"),
)
def test_get_llm_provider_uses_aws_region_name_for_responses(
self, monkeypatch, local_cost_map
):
def test_get_llm_provider_uses_aws_region_name_for_responses(self, monkeypatch, local_cost_map):
from litellm.types.router import GenericLiteLLMParams
monkeypatch.delenv("BEDROCK_MANTLE_REGION", raising=False)
@ -167,18 +155,14 @@ class TestBedrockMantleConfig:
monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2")
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
cfg = BedrockMantleChatConfig()
api_base, _ = cfg._get_openai_compatible_provider_info(
None, None, model="openai.gpt-oss-120b"
)
api_base, _ = cfg._get_openai_compatible_provider_info(None, None, model="openai.gpt-oss-120b")
assert api_base == "https://bedrock-mantle.us-east-2.api.aws/v1"
@pytest.mark.parametrize(
"model_id",
["google.gemma-4-31b", "google.gemma-4-26b-a4b", "google.gemma-4-e2b"],
)
def test_chat_base_for_gemma_4_uses_openai_v1(
self, monkeypatch, local_cost_map, model_id
):
def test_chat_base_for_gemma_4_uses_openai_v1(self, monkeypatch, local_cost_map, model_id):
# The chat-config bug the Gemma 4 cards exposed: gemma-4-* is served on the
# /openai/v1 base, not the hardcoded /v1. Driven by the price-map
# use_openai_responses_path flag (loaded by local_cost_map). Fails before
@ -186,22 +170,16 @@ class TestBedrockMantleConfig:
monkeypatch.setenv("BEDROCK_MANTLE_REGION", "us-east-2")
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
cfg = BedrockMantleChatConfig()
api_base, _ = cfg._get_openai_compatible_provider_info(
None, None, model=model_id
)
api_base, _ = cfg._get_openai_compatible_provider_info(None, None, model=model_id)
assert api_base == "https://bedrock-mantle.us-east-2.api.aws/openai/v1"
def test_chat_base_explicit_api_base_wins_over_derived(
self, monkeypatch, local_cost_map
):
def test_chat_base_explicit_api_base_wins_over_derived(self, monkeypatch, local_cost_map):
# An explicit api_base must not be overridden by the data-driven default,
# even for a model whose default differs (gemma-4 -> openai/v1).
monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False)
custom_base = "https://bedrock-mantle.us-west-2.api.aws/v1"
cfg = BedrockMantleChatConfig()
api_base, _ = cfg._get_openai_compatible_provider_info(
custom_base, None, model="google.gemma-4-31b"
)
api_base, _ = cfg._get_openai_compatible_provider_info(custom_base, None, model="google.gemma-4-31b")
assert api_base == custom_base
def test_api_key_from_env(self, monkeypatch):
@ -244,9 +222,7 @@ class TestBedrockMantleChatAuth:
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
signer = BaseAWSLLM()
signer.get_credentials = MagicMock(
side_effect=AssertionError("SigV4 must not run when a Bearer token exists")
)
signer.get_credentials = MagicMock(side_effect=AssertionError("SigV4 must not run when a Bearer token exists"))
return signer
def test_bearer_token_skips_sigv4(self, monkeypatch):
@ -363,9 +339,7 @@ class TestBedrockMantleChatAuth:
assert "/eu-west-1/bedrock/aws4_request" in headers["Authorization"]
def test_sigv4_scope_matches_api_base_when_aws_region_name_disagrees(
self, monkeypatch
):
def test_sigv4_scope_matches_api_base_when_aws_region_name_disagrees(self, monkeypatch):
# If a caller (e.g. proxy) passes a stale api_base in one region and an
# aws_region_name in a different region, the SigV4 credential scope must
# match the URL host or Bedrock rejects the request with 401. Without the
@ -439,9 +413,7 @@ class TestBedrockMantleChatAuth:
):
monkeypatch.delenv(var, raising=False)
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIAEXAMPLE")
monkeypatch.setenv(
"AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0"
)
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "c2VjcmV0LXRlc3Qtc2VjcmV0LXRlc3Qtc2VjcmV0")
monkeypatch.setenv("AWS_REGION", "us-east-2")
requests = []
@ -471,9 +443,7 @@ class TestBedrockMantleChatAuth:
request=httpx.Request("POST", url),
)
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post
):
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post):
response = litellm.completion(
model="bedrock_mantle/openai.gpt-oss-120b",
messages=[{"role": "user", "content": "hello"}],
@ -518,9 +488,7 @@ class TestBedrockMantleProjectHeader:
def mock_post(self, url, data=None, headers=None, **kwargs):
raw_body = data.decode("utf-8") if isinstance(data, bytes) else data
requests.append(
{"headers": headers or {}, "body": json.loads(raw_body or "{}")}
)
requests.append({"headers": headers or {}, "body": json.loads(raw_body or "{}")})
return httpx.Response(
status_code=200,
json={
@ -544,9 +512,7 @@ class TestBedrockMantleProjectHeader:
request=httpx.Request("POST", url),
)
with patch(
"litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post
):
with patch("litellm.llms.custom_httpx.http_handler.HTTPHandler.post", mock_post):
response = litellm.completion(
model="bedrock_mantle/openai.gpt-oss-120b",
messages=[{"role": "user", "content": "hello"}],
@ -562,16 +528,12 @@ class TestBedrockMantleProjectHeader:
class TestBedrockMantleProviderResolution:
def test_get_llm_provider_resolves_correctly(self):
model, provider, _, _ = litellm.get_llm_provider(
"bedrock_mantle/openai.gpt-oss-120b"
)
model, provider, _, _ = litellm.get_llm_provider("bedrock_mantle/openai.gpt-oss-120b")
assert provider == "bedrock_mantle"
assert model == "openai.gpt-oss-120b"
def test_get_llm_provider_20b(self):
model, provider, _, _ = litellm.get_llm_provider(
"bedrock_mantle/openai.gpt-oss-20b"
)
model, provider, _, _ = litellm.get_llm_provider("bedrock_mantle/openai.gpt-oss-20b")
assert provider == "bedrock_mantle"
assert model == "openai.gpt-oss-20b"
@ -617,9 +579,7 @@ class TestBedrockMantlePricing:
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "true")
litellm.add_known_models()
info_120b = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-120b")
info_safeguard = litellm.get_model_info(
"bedrock_mantle/openai.gpt-oss-safeguard-120b"
)
info_safeguard = litellm.get_model_info("bedrock_mantle/openai.gpt-oss-safeguard-120b")
assert info_safeguard["max_output_tokens"] > info_120b["max_output_tokens"]
def test_reasoning_support(self, monkeypatch):
@ -643,9 +603,7 @@ class TestBedrockMantlePricing:
("google.gemma-4-e2b", 4e-08, 8e-08, 128000),
],
)
def test_gemma_4_bedrock_mantle_model_metadata(
local_cost_map, model_id, input_cost, output_cost, max_tokens
):
def test_gemma_4_bedrock_mantle_model_metadata(local_cost_map, model_id, input_cost, output_cost, max_tokens):
full_model_name = f"bedrock_mantle/{model_id}"
info = litellm.get_model_info(full_model_name)
@ -659,10 +617,7 @@ def test_gemma_4_bedrock_mantle_model_metadata(
assert info["supports_tool_choice"] is True
assert info["supports_vision"] is True
assert (
litellm.supports_parallel_function_calling(
model=full_model_name, custom_llm_provider="bedrock_mantle"
)
is False
litellm.supports_parallel_function_calling(model=full_model_name, custom_llm_provider="bedrock_mantle") is False
)
@ -682,3 +637,72 @@ def test_gemma_4_models_register_under_bedrock_mantle(local_cost_map, model_id):
resolved_model, provider, _, _ = litellm.get_llm_provider(full_model_name)
assert provider == "bedrock_mantle"
assert resolved_model == model_id
class TestBedrockMantleMaxTokensOmission:
"""
Regression coverage for https://github.com/BerriAI/litellm/issues/36970:
Mantle answers max_tokens with `unsupported_parameter` for the gemma-4-*
family, and the native-API name max_output_tokens is not a substitute on the
OpenAI-compatible route, so the param must not reach Bedrock. Driven by the
price-map omit_max_tokens_param flag, never by a model-name match.
"""
@pytest.mark.parametrize(
"model_id",
["google.gemma-4-31b", "google.gemma-4-26b-a4b", "google.gemma-4-e2b"],
)
def test_gemma_4_does_not_support_max_tokens(self, local_cost_map, model_id):
cfg = BedrockMantleChatConfig()
assert "max_tokens" not in cfg.get_supported_openai_params(model_id)
def test_non_gemma_model_still_supports_max_tokens(self, local_cost_map):
cfg = BedrockMantleChatConfig()
assert "max_tokens" in cfg.get_supported_openai_params("openai.gpt-oss-120b")
def test_max_tokens_dropped_for_gemma_4_when_drop_params(self, local_cost_map):
optional_params = litellm.utils.get_optional_params(
model="google.gemma-4-31b",
custom_llm_provider="bedrock_mantle",
max_tokens=50,
drop_params=True,
)
assert "max_tokens" not in optional_params
def test_max_tokens_raises_for_gemma_4_without_drop_params(self, local_cost_map):
with pytest.raises(litellm.UnsupportedParamsError) as exc_info:
litellm.utils.get_optional_params(
model="google.gemma-4-31b",
custom_llm_provider="bedrock_mantle",
max_tokens=50,
drop_params=False,
)
assert "max_tokens" in str(exc_info.value)
def test_max_tokens_forwarded_for_non_gemma_model(self, local_cost_map):
optional_params = litellm.utils.get_optional_params(
model="openai.gpt-oss-120b",
custom_llm_provider="bedrock_mantle",
max_tokens=50,
drop_params=False,
)
assert optional_params["max_tokens"] == 50
def test_flag_is_data_driven_not_name_matched(self, local_cost_map, monkeypatch):
from litellm.llms.bedrock_mantle.common_utils import mantle_omits_max_tokens
monkeypatch.setitem(
litellm.model_cost,
"bedrock_mantle/vendor.some-future-model",
{"omit_max_tokens_param": True},
)
cfg = BedrockMantleChatConfig()
assert mantle_omits_max_tokens("vendor.some-future-model", litellm.model_cost) is True
assert "max_tokens" not in cfg.get_supported_openai_params("vendor.some-future-model")
assert mantle_omits_max_tokens("google.gemma-4-31b-not-a-real-model", litellm.model_cost) is False