From de896efc8ec22a14e8b7032688b79615088a6d07 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Wed, 21 Oct 2026 01:51:08 +0000 Subject: [PATCH] fix(count_tokens): fall back to Mantle on Bedrock 404 (Bedrock part of #46430) (cherry picked from commit 08a24d391af6ba0fbd4e99313857fbb05754ddf2) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../count_tokens/bedrock_token_counter.py | 5 +- .../test_bedrock_mantle_count_tokens_wire.py | 134 ++++++++++++++++++ .../test_bedrock_token_counter.py | 38 ++++- 3 files changed, 174 insertions(+), 3 deletions(-) diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 18a7b362d85..dec441565d8 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -20,6 +20,7 @@ from litellm.types.utils import LiteLLMPydanticObjectBase, LlmProviders, TokenCo RUNTIME_TOKENIZER_TYPE: Final = "bedrock_api" MANTLE_TOKENIZER_TYPE: Final = "bedrock_mantle_api" +_RUNTIME_FALLBACK_STATUSES: Final = frozenset({400, 404}) class _CountTokensReply(LiteLLMPydanticObjectBase): @@ -46,14 +47,14 @@ CountTokensOutcome = CountedTokens | CountTokensFailure def _runtime_rejected_claude_model(outcome: CountTokensOutcome, resolved_model: str) -> bool: return ( isinstance(outcome, CountTokensFailure) - and outcome.status_code == 400 + and outcome.status_code in _RUNTIME_FALLBACK_STATUSES and is_mantle_claude_model(resolved_model) ) class BedrockTokenCounter(BaseTokenCounter): """Token counter for AWS Bedrock: bedrock-runtime CountTokens first, and Anthropic's count_tokens - on bedrock-mantle for the Claude models bedrock-runtime answers 400 for""" + on bedrock-mantle for Claude models bedrock-runtime answers 400 or 404 for""" def __init__( self, diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index f9d4f7f91bb..77349a1f206 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -451,6 +451,140 @@ def test_gemini_count_tokens_route_reports_the_runtime_total_for_a_model_bedrock assert mantle.drain() == () +def test_runtime_404_on_claude_falls_back_to_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + system: Final = f"{_SYSTEM} status=404" + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert _payload(response) == {"input_tokens": _MANTLE_COUNT}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) + + +def test_runtime_404_and_mantle_404_surface_provider_error(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + system: Final = f"{_SYSTEM} status=404" + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(_rejecting(404)), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = _count(gateway, {**_bare(model), "system": system, "tools": _TOOLS}) + assert response.status_code == 404, response.text + assert _REJECTION in response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == (_mantle_body(system=system, tools=_TOOLS),) + + +def test_runtime_404_on_nova_does_not_call_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: + gateway: Final = counting_proxy.gateway + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url, _NOVA) + body: Final = {**_bare(model), "messages": [{"role": "user", "content": "status=404"}]} + response: Final = _count(gateway, body) + assert response.status_code == 404, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert mantle.drain() == () + + +def test_google_count_tokens_forwards_system_and_tools_to_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + system_text: Final = "Use the provided tool" + tool_schema: Final = {"type": "object", "properties": {"location": {"type": "string"}}} + body: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "status=404"}]}], + "systemInstruction": {"parts": [{"text": system_text}]}, + "tools": [ + { + "functionDeclarations": [ + {"name": "get_weather", "description": "Look up weather", "parameters": tool_schema} + ] + } + ], + } + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ( + _mantle_body( + messages=[{"role": "user", "content": "status=404"}], + system=[{"type": "text", "text": system_text}], + tools=[ + { + "name": "get_weather", + "description": "Look up weather", + "input_schema": tool_schema, + } + ], + ), + ) + + +def test_google_count_tokens_normalizes_uppercase_parameters_for_mantle( + counting_proxy: OwnedProxy, mantle_port: int +) -> None: + gateway: Final = counting_proxy.gateway + body: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "status=404"}]}], + "tools": [ + { + "functionDeclarations": [ + { + "name": "get_weather", + "parameters": { + "type": "OBJECT", + "properties": {"location": {"type": "STRING"}}, + "required": ["location"], + }, + } + ] + } + ], + } + with ( + wire_server(_runtime) as runtime, + wire_server(_mantle(), port=mantle_port) as mantle, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, runtime.url) + response: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", body) + assert _payload(response) == {"totalTokens": _MANTLE_COUNT, "promptTokensDetails": []}, response.text + assert len(_runtime_count_targets(runtime)) == 1 + assert _mantle_bodies(mantle) == ( + _mantle_body( + messages=[{"role": "user", "content": "status=404"}], + tools=[ + { + "name": "get_weather", + "input_schema": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + ], + ), + ) + + def test_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with ( diff --git a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py index a3dd8021f73..8a21027872a 100644 --- a/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/test_litellm/llms/bedrock/count_tokens/test_bedrock_token_counter.py @@ -93,8 +93,10 @@ async def test_bedrock_runtime_count_is_kept_when_it_answers() -> None: @pytest.mark.parametrize( ("model_to_use", "runtime"), ( + ("global.anthropic.claude-opus-4-8", (401, {"Message": "not authorized to perform: bedrock:CountTokens"})), ("global.anthropic.claude-opus-4-8", (403, {"Message": "not authorized to perform: bedrock:CountTokens"})), ("amazon.nova-pro-v1:0", (400, UNSUPPORTED)), + ("amazon.nova-pro-v1:0", (404, UNSUPPORTED)), ), ) async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( @@ -120,7 +122,7 @@ async def test_other_bedrock_runtime_failures_are_not_retried_on_mantle( @pytest.mark.asyncio async def test_mantle_failure_is_reported_with_its_status() -> None: bedrock: Final = _Bedrock( - runtime=(400, UNSUPPORTED), + runtime=(404, UNSUPPORTED), mantle=(404, {"type": "error", "error": {"type": "not_found_error", "message": "does not exist"}}), ) @@ -138,3 +140,37 @@ async def test_mantle_failure_is_reported_with_its_status() -> None: assert result.tokenizer_type == "bedrock_mantle_api" assert "does not exist" in (result.error_message or "") assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + + +@pytest.mark.asyncio +async def test_claude_runtime_404_retries_with_system_and_tools_on_mantle() -> None: + bedrock: Final = _Bedrock(runtime=(404, UNSUPPORTED), mantle=(200, {"input_tokens": 2177})) + tools: Final = [ + { + "name": "lookup", + "description": "Find a record", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + } + ] + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="global.anthropic.claude-opus-4-8", + messages=MESSAGES, + contents=None, + deployment=DEPLOYMENT, + request_model="claude-opus-4-8", + system="You are a terse assistant.", + tools=tools, + ) + + assert result is not None + assert result.error is False + assert result.total_tokens == 2177 + assert bedrock.posted_hosts() == (RUNTIME_HOST, MANTLE_HOST) + mantle_body: Final = _JSON_BODY.validate_json(bedrock.requests[1].content) + assert mantle_body == { + "model": "anthropic.claude-opus-4-8", + "messages": MESSAGES, + "system": "You are a terse assistant.", + "tools": tools, + }