diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py index e32df0bbdbb..e150791b09e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py @@ -1,11 +1,13 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import TypeAdapter +import litellm from litellm.anthropic_beta_gated_fields import strip_beta_gated_request_body from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_anthropic_image_obj, @@ -25,7 +27,9 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation ) from litellm.llms.bedrock.common_utils import ( BEDROCK_INVOKE_TRACE_FIELD, + BedrockModelInfo, apply_bedrock_invoke_structured_output, + bedrock_capability_model, bedrock_chat_uses_native_structured_output, bedrock_cleaned_output_schema, bedrock_supports_tool_search, @@ -35,6 +39,7 @@ from litellm.llms.bedrock.common_utils import ( normalize_custom_field_on_tools, normalize_tool_input_schema_types_for_bedrock_invoke, requests_opt_in_native_output, + strip_bedrock_routing_prefix, strip_unsupported_bedrock_invoke_output_config_keys, tools_without_eager_input_streaming, ) @@ -59,6 +64,19 @@ _NON_DEFAULT_PARAMS: Final = TypeAdapter(Mapping[str, object]) _LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object]) +def _anthropic_config_max_tokens() -> int | None: + return litellm.AnthropicConfig.max_tokens + + +def _mapped_max_tokens(model: str) -> int | None: + row: Final = object_mapping(litellm.model_cost.get(model)) or {} # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # cost map rows are untyped JSON + match row.get("max_output_tokens"): + case int() as max_tokens: + return max_tokens + case _: + return None + + class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): """ Reference: @@ -79,9 +97,18 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): anthropic_version: str = "bedrock-2023-05-31" - def __init__(self, model_info: Mapping[str, object] | None = None) -> None: + def __init__( + self, + model_info: Mapping[str, object] | None = None, + read_configured_max_tokens: Callable[[], int | None] = _anthropic_config_max_tokens, + ) -> None: super().__init__() self._deployment_model_info: Final = model_info + self._read_configured_max_tokens: Final = read_configured_max_tokens + + @classmethod + def get_config(cls, *, model: str | None = None) -> dict[str, object]: # mutable-ok: overrides BaseConfig's dict + return {key: value for key, value in super().get_config(model=model).items() if key != "max_tokens"} @property def custom_llm_provider(self) -> str | None: @@ -241,6 +268,36 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): headers=headers, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return ( + BedrockModelInfo.get_base_model( + bedrock_capability_model(strip_bedrock_routing_prefix(model), litellm_params.get("base_model")) + ), + ) + + def _default_max_tokens(self, model: str, litellm_params: Mapping[str, object]) -> int: + configured_max_tokens: Final = self._read_configured_max_tokens() + match deployment_model_info(litellm_params).get("max_output_tokens"): + case int() as deployment_max_tokens: + return ( + deployment_max_tokens + if configured_max_tokens is None + else min(deployment_max_tokens, configured_max_tokens) + ) + case _ if configured_max_tokens is not None: + return configured_max_tokens + case _: + return next( + ( + max_tokens + for max_tokens in map( + _mapped_max_tokens, self._default_max_tokens_models(model, litellm_params) + ) + if max_tokens is not None + ), + DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + ) + def _build_bedrock_anthropic_request_base( self, model: str, @@ -257,13 +314,21 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): model=model, output_config=filtered_params["output_config"], ) - filtered_params = self._normalize_bedrock_tool_search_tools(filtered_params) + normalized_params: Final = self._normalize_bedrock_tool_search_tools(filtered_params) + default_max_tokens: Final = ( + None + if "max_tokens" in normalized_params + else self._default_max_tokens(model, _LITELLM_PARAMS.validate_python(litellm_params)) + ) + request_params: Final = ( + normalized_params if default_max_tokens is None else {**normalized_params, "max_tokens": default_max_tokens} + ) anthropic_request: Final = AnthropicConfig.transform_request( self, model=model, messages=messages, - optional_params=filtered_params, + optional_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -377,7 +442,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): "data": image_chunk["data"], } - def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict: + def _normalize_bedrock_tool_search_tools(self, optional_params: dict) -> dict[str, object]: """ Convert tool search entries to the format supported by the Bedrock Invoke API. """ diff --git a/litellm/llms/bedrock/chat/mantle/transformation.py b/litellm/llms/bedrock/chat/mantle/transformation.py index b5a352bab48..cbde6f918bc 100644 --- a/litellm/llms/bedrock/chat/mantle/transformation.py +++ b/litellm/llms/bedrock/chat/mantle/transformation.py @@ -54,6 +54,9 @@ class AmazonMantleConfig(AmazonAnthropicClaudeConfig): region=region, ) + def _default_max_tokens_models(self, model: str, litellm_params: Mapping[str, object]) -> tuple[str, ...]: + return (f"bedrock_mantle/{model}", *super()._default_max_tokens_models(model, litellm_params)) + def validate_environment( self, headers: dict, diff --git a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index d4f9533fcdc..32b6c4f1690 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -70,10 +70,18 @@ def _chat_messages(marker: str) -> list[JsonValue]: return [{"role": "user", "content": _prompt(marker)}] -def _outbound(backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: +def _default_max_tokens(gateway: Gateway, backend: str) -> int: + match gateway.get("/public/litellm_model_cost_map").get(f"bedrock_mantle/{backend}"): + case {"max_output_tokens": int() as mapped}: + return mapped + case _: + return DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + + +def _outbound(gateway: Gateway, backend: str, messages: list[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: return { "messages": messages, - "max_tokens": DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + "max_tokens": _default_max_tokens(gateway, backend), "anthropic_version": _BRIDGE_VERSION, "model": backend, **fields, @@ -290,7 +298,7 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ gateway: Gateway, backend: str ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -314,7 +322,7 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -355,7 +363,7 @@ def test_chat_request_fields_are_translated_to_the_native_messages_shape( outbound_fields: dict[str, JsonValue], ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -377,7 +385,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ marker: Final = uuid4().hex backend: Final = _mantle_claude_where(gateway, "supports_forced_tool_use", disabled=False) expected: Final = _outbound( - backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice + gateway, backend, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], tool_choice=outbound_choice ) reply: Final = _tool_reply(marker, backend, "get_weather") with wire_server(_bearer_peer(expected, reply)) as wire, gateway.scenario() as scenario: @@ -407,7 +415,7 @@ def test_chat_tools_reach_mantle_as_input_schema_and_the_tool_use_block_returns_ def test_chat_stream_relays_a_native_tool_use_block_as_tool_call_chunks(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], tools=[_WEATHER_TOOL_OUTBOUND], stream=True) with wire_server(_bearer_peer(expected, _tool_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -439,7 +447,7 @@ def test_chat_tool_result_turn_reaches_mantle_as_tool_use_and_tool_result_blocks marker: Final = uuid4().hex tool_call: Final = f"toolu_{marker}" expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [ _user_turn(marker), { @@ -484,7 +492,7 @@ def test_chat_json_schema_response_format_is_sent_as_a_forced_tool_and_returned_ ) -> None: marker: Final = uuid4().hex expected: Final = _outbound( - _HAIKU, + gateway, _HAIKU, [_user_turn(marker)], tools=[{"name": "json_tool_call", "input_schema": _CITY_SCHEMA}], tool_choice={"name": "json_tool_call", "type": "tool"}, @@ -537,7 +545,7 @@ def _sigv4_peer(expected: Mapping[str, JsonValue], reply: Reply) -> Callable[[Re def test_chat_completion_signs_the_native_messages_request_with_the_deployment_aws_keys(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_sigv4_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model( model=f"bedrock_mantle/{_HAIKU}", @@ -556,7 +564,7 @@ def test_chat_completion_signs_the_native_messages_request_with_the_deployment_a def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantle_price(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/us-east-2/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -577,7 +585,7 @@ def test_region_prefix_in_a_mantle_claude_id_is_routing_only_and_keeps_the_mantl def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -606,7 +614,7 @@ def test_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> N def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) with OpenAI( @@ -645,7 +653,7 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -674,7 +682,7 @@ async def test_async_openai_sdk_chat_completion_on_a_mantle_claude_id(gateway: G async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) async with AsyncOpenAI( @@ -714,7 +722,7 @@ def test_identical_chat_request_is_served_from_the_response_cache_and_logged_as_ gateway: Gateway, ) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) request: Final[dict[str, JsonValue]] = {"model": model, "messages": _chat_messages(marker)} @@ -802,7 +810,7 @@ def test_open_weight_mantle_id_keeps_the_openai_compatible_chat_route_and_its_sa ) def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(gateway: Gateway, suffix: str) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url + suffix, api_key=_API_KEY) response: Final = gateway.request( @@ -814,7 +822,7 @@ def test_api_base_with_a_route_suffix_still_posts_to_the_native_messages_path(ga def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)]) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request("POST", "/v1/responses", {"model": model, "input": _prompt(marker)}) @@ -839,7 +847,7 @@ def test_responses_api_on_a_mantle_claude_id_bridges_to_native_messages(gateway: def test_responses_api_stream_on_a_mantle_claude_id_bridges_to_the_native_stream(gateway: Gateway) -> None: marker: Final = uuid4().hex - expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) + expected: Final = _outbound(gateway, _HAIKU, [_user_turn(marker)], stream=True) with wire_server(_bearer_peer(expected, _text_stream(marker))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{_HAIKU}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -996,7 +1004,7 @@ def test_unsupported_claude_param_is_dropped_under_drop_params( ) -> None: marker: Final = uuid4().hex backend: Final = _backend_without(gateway, flag) - expected: Final = _outbound(backend, [_user_turn(marker)], **outbound_fields) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)], **outbound_fields) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = scenario.model(model=f"bedrock_mantle/{backend}", api_base=wire.url, api_key=_API_KEY) response: Final = gateway.request( @@ -1114,7 +1122,7 @@ def _wildcard_prefix(gateway: Gateway, scenario: Scenario, api_base: str) -> str def test_uppercase_claude_id_through_a_mantle_wildcard_is_sent_to_native_messages(gateway: Gateway) -> None: marker: Final = uuid4().hex backend: Final = _HAIKU.upper() - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) with wire_server(_bearer_peer(expected, _text_reply(marker, backend))) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" response: Final = gateway.request( @@ -1129,7 +1137,7 @@ def test_five_kilobyte_claude_id_through_a_mantle_wildcard_reaches_mantle_and_it ) -> None: marker: Final = uuid4().hex backend: Final = f"{_HAIKU}-{'x' * 5000}" - expected: Final = _outbound(backend, [_user_turn(marker)]) + expected: Final = _outbound(gateway, backend, [_user_turn(marker)]) unknown: Final = _error_reply(404, "not_found_error", "model: synthetic unknown model") with wire_server(_bearer_peer(expected, unknown)) as wire, gateway.scenario() as scenario: model: Final = f"{_wildcard_prefix(gateway, scenario, wire.url)}/{backend}" diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py index d40abb4860b..4850a402103 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude_transformation.py @@ -14,6 +14,7 @@ from pydantic import BaseModel, ConfigDict # Ensure the project root is on the import path so `litellm` can be imported when # tests are executed from any working directory. import litellm +from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude_transformation import ( AmazonAnthropicClaudeConfig, ) @@ -2010,6 +2011,123 @@ def test_invoke_uses_opt_in_native_output_is_false_when_native_output_did_not_co ) +BEDROCK_HAIKU_45_BASE_ID: Final = "anthropic.claude-haiku-4-5-20251001-v1:0" + + +@pytest.mark.parametrize( + "model", + [ + f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + f"invoke/{BEDROCK_HAIKU_45_BASE_ID}", + ], +) +def test_bedrock_invoke_claude_default_max_tokens_comes_from_base_model(local_model_cost_map, model: str): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=model, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"default max_tokens should be the base model's cap, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_explicit_max_tokens_is_kept(local_model_cost_map): + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 123}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == 123, f"explicit max_tokens must not be overridden, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_takes_the_deployments_own_cap(local_model_cost_map): + deployment_cap: Final = 8192 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != deployment_cap + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == deployment_cap, f"the deployment's own cap should win, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_ignores_a_sibling_deployments_cap(local_model_cost_map): + regional_id: Final = f"us.{BEDROCK_HAIKU_45_BASE_ID}" + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + litellm.register_model(model_cost={regional_id: {"max_output_tokens": 50}}, persist_across_reloads=False) + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model=f"invoke/{regional_id}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"a sibling's cap on the shared id leaked, got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_yields_to_anthropic_config(local_model_cost_map): + configured_max_tokens: Final = 1000 + assert litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) != configured_max_tokens + + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == configured_max_tokens, f"AnthropicConfig(max_tokens) was ignored, got {body['max_tokens']}" + + +@pytest.mark.parametrize(("configured_max_tokens", "deployment_cap"), [(100000, 50), (1000, 8192)]) +def test_bedrock_invoke_claude_default_max_tokens_is_the_lower_of_anthropic_config_and_the_deployments_cap( + local_model_cost_map, configured_max_tokens: int, deployment_cap: int +): + body: Final = AmazonAnthropicClaudeConfig( + read_configured_max_tokens=lambda: configured_max_tokens + ).transform_request( + model=f"invoke/us.{BEDROCK_HAIKU_45_BASE_ID}", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"model_info": {"max_output_tokens": deployment_cap}}, + headers={}, + ) + + assert body["max_tokens"] == min(configured_max_tokens, deployment_cap), f"got {body['max_tokens']}" + + +def test_bedrock_invoke_claude_default_max_tokens_reads_an_inference_profile_arns_base_model(local_model_cost_map): + mapped_cap: Final = litellm.get_max_tokens(BEDROCK_HAIKU_45_BASE_ID) + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonAnthropicClaudeConfig().transform_request( + model="invoke/anthropic/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={"base_model": BEDROCK_HAIKU_45_BASE_ID}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"an ARN should take its base_model's cap, got {body['max_tokens']}" + + @pytest.mark.usefixtures("local_model_cost_map", "local_beta_headers_config") @pytest.mark.parametrize("drop_params", [True, False]) def test_bedrock_chat_invoke_adds_the_binding_beta_so_block_binding_is_kept( diff --git a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py index 154be9f0370..56d6b039788 100644 --- a/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py +++ b/tests/unit/llms/bedrock/chat/mantle/test_bedrock_chat_mantle_transformation.py @@ -2,6 +2,7 @@ import base64 import json import uuid from types import SimpleNamespace +from typing import Final import httpx import pytest @@ -243,3 +244,21 @@ async def test_mantle_claude_async_completion_streams_a_wrapped_json_schema_answ chunks = tuple([chunk async for chunk in stream]) assert _streamed_answer(chunks) == BARE_STREAMED_ANSWER + + +def test_mantle_claude_default_max_tokens_comes_from_the_mantle_row(local_model_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.bedrock.chat.mantle.transformation import AmazonMantleConfig + + mapped_cap: Final = litellm.get_max_tokens("bedrock_mantle/anthropic.claude-haiku-4-5") + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + + body: Final = AmazonMantleConfig().transform_request( + model="mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert body["max_tokens"] == mapped_cap, f"Mantle should send its mapped cap, got {body['max_tokens']}" diff --git a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py index f252d9886d3..6005acffc0a 100644 --- a/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py +++ b/tests/unit/llms/bedrock_mantle/test_bedrock_mantle_transformation.py @@ -1096,6 +1096,25 @@ class TestBedrockMantleClaudeChatRoute: ) assert response._hidden_params["response_cost"] != pytest.approx(_row_cost("anthropic.claude-opus-5-5", 10, 5)) + def test_claude_completion_without_max_tokens_sends_the_mantle_rows_output_limit(self, monkeypatch, local_cost_map): + from litellm.constants import DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + mapped_cap: Final = litellm.model_cost["bedrock_mantle/anthropic.claude-haiku-4-5"]["max_output_tokens"] + assert mapped_cap != DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, "the model's mapped cap must differ from the fallback" + monkeypatch.setenv("BEDROCK_MANTLE_API_KEY", "mantle-key") + monkeypatch.delenv("BEDROCK_MANTLE_API_BASE", raising=False) + handler: Final = Mock(side_effect=_anthropic_message) + + litellm.completion( + model="bedrock_mantle/anthropic.claude-haiku-4-5", + messages=[{"role": "user", "content": "hello"}], + aws_region_name="us-east-1", + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))), + ) + + assert json.loads(handler.call_args.args[0].content)["max_tokens"] == mapped_cap + def test_claude_streaming_completion_uses_native_messages_endpoint(self, monkeypatch, local_cost_map): from litellm.llms.custom_httpx.http_handler import HTTPHandler