From ef3c50f2d99ae840dc10f12135b7e1191d2ddba3 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Tue, 20 Oct 2026 17:46:21 +0000 Subject: [PATCH] fix(count_tokens): fall back to Mantle on Bedrock 404 (Bedrock part of #46430) Backport of #46430 to stable/1.104.x. Cherry-picked from 08a24d391af6ba0fbd4e99313857fbb05754ddf2 (main), Bedrock changes only. 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 | 46 +++++++++++++++++++ .../test_bedrock_token_counter.py | 38 ++++++++++++++- 3 files changed, 86 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 65a6aee027c..66fcd53b948 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 the 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..8d8750a1da9 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,52 @@ 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_responses_input_tokens_counts_through_mantle(counting_proxy: OwnedProxy, mantle_port: int) -> None: gateway: Final = counting_proxy.gateway with ( diff --git a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py b/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py index a3dd8021f73..8a21027872a 100644 --- a/tests/unit/llms/bedrock/count_tokens/test_bedrock_token_counter.py +++ b/tests/unit/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, + }