From f18056a85bfb100b312f95cf5b117acb61e38a90 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 04:57:00 -0700 Subject: [PATCH] fix(count_tokens): keep a local conversion 400 as a 400 on the Claude count routes The anthropic, azure_ai and Bedrock Mantle count handlers re-raise a BadRequestError raised before the provider call, and their token counters mark it rejected_before_provider_call, so a framing deployment answers the 400 instead of an estimate that leaves the attachment out, and Mantle no longer turns it into a 500 --- .../llms/anthropic/count_tokens/handler.py | 4 +- .../anthropic/count_tokens/token_counter.py | 11 +++++ .../anthropic/count_tokens/handler.py | 4 +- .../anthropic/count_tokens/token_counter.py | 11 +++++ .../count_tokens/bedrock_token_counter.py | 6 +++ .../bedrock/count_tokens/mantle_handler.py | 2 +- ...t_anthropic_count_tokens_transformation.py | 30 -------------- .../test_bedrock_token_counter.py | 18 ++++++++ tests/unit/proxy/test_proxy_token_counter.py | 41 +++++++++++++++++++ 9 files changed, 92 insertions(+), 35 deletions(-) diff --git a/litellm/llms/anthropic/count_tokens/handler.py b/litellm/llms/anthropic/count_tokens/handler.py index 4de88a287de..1b9b31d8ec8 100644 --- a/litellm/llms/anthropic/count_tokens/handler.py +++ b/litellm/llms/anthropic/count_tokens/handler.py @@ -116,8 +116,8 @@ class AnthropicCountTokensHandler(AnthropicCountTokensConfig): except AnthropicError: # Re-raise Anthropic exceptions as-is raise - except litellm.BadRequestError as e: - raise AnthropicError(status_code=400, message=e.message) from e + except litellm.BadRequestError: + raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code verbose_logger.error("HTTP error in CountTokens handler: %s", e) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index 91fc77773d1..65b6f56454b 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -143,6 +143,17 @@ class AnthropicTokenCounter(BaseTokenCounter): tokenizer_type="anthropic_api", original_response=result, ) + except litellm.BadRequestError as e: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=e.message, + status_code=400, + rejected_before_provider_call=True, + ) except (AnthropicError, AuthenticationError) as e: verbose_logger.warning("Anthropic CountTokens error: status=%s, message=%s", e.status_code, e.message) return TokenCountResponse( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index b443ccc7713..f4d426bc47b 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -128,8 +128,8 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): except AnthropicError: # Re-raise Anthropic exceptions as-is raise - except litellm.BadRequestError as e: - raise AnthropicError(status_code=400, message=e.message) from e + except litellm.BadRequestError: + raise except httpx.HTTPStatusError as e: # HTTP errors - preserve the actual status code verbose_logger.error("HTTP error in CountTokens handler: %s", e) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index af7cf9d81ce..da5ecc6c5db 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -147,6 +147,17 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): "original_response": result, } ) + except litellm.BadRequestError as e: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=e.message, + status_code=400, + rejected_before_provider_call=True, + ) except AnthropicError as e: verbose_logger.warning( "Azure AI Anthropic CountTokens API error: status=%s, message=%s", e.status_code, e.message diff --git a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py index 26bc5904e06..743eeaa014f 100644 --- a/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py +++ b/litellm/llms/bedrock/count_tokens/bedrock_token_counter.py @@ -49,6 +49,7 @@ class CountTokensFailure: status_code: int message: str tokenizer_type: str + rejected_before_provider_call: bool = False CountTokensOutcome = CountedTokens | CountTokensFailure @@ -120,6 +121,10 @@ class BedrockTokenCounter(BaseTokenCounter): client=self._client, ) reply: Final = _CountTokensReply.model_validate(result) + except litellm.BadRequestError as e: + return CountTokensFailure( + status_code=400, message=e.message, tokenizer_type=tokenizer_type, rejected_before_provider_call=True + ) except BedrockError as e: verbose_logger.debug( "%s CountTokens API error: status=%s, message=%s", tokenizer_type, e.status_code, e.message @@ -214,6 +219,7 @@ class BedrockTokenCounter(BaseTokenCounter): error=True, error_message=outcome.message, status_code=outcome.status_code, + rejected_before_provider_call=outcome.rejected_before_provider_call, ) case _: assert_never(outcome) diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 5cd26235ba2..d321f6b394b 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -99,7 +99,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): response=response, ) return _COUNT_RESPONSE.validate_json(response.content) - except BedrockError: + except (BedrockError, litellm.BadRequestError): raise except httpx.HTTPStatusError as e: raise BedrockError( diff --git a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py index 80f760b5243..0884969b0c6 100644 --- a/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/anthropic/test_anthropic_count_tokens_transformation.py @@ -14,7 +14,6 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.llms.anthropic.common_utils import AnthropicError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler from litellm.llms.anthropic.count_tokens.token_counter import AnthropicTokenCounter from litellm.llms.anthropic.count_tokens.transformation import ( @@ -526,35 +525,6 @@ def test_count_drops_a_leading_system_message_without_text() -> None: assert result == {"model": "claude-opus-5-5", "messages": [{"role": "user", "content": "hi"}]} -@pytest.mark.asyncio -@pytest.mark.usefixtures("httpx_transport_clients") -@pytest.mark.parametrize("handler_type", (AnthropicCountTokensHandler, AzureAIAnthropicCountTokensHandler)) -async def test_handler_answers_400_for_a_text_attachment_it_cannot_decode( - handler_type: type[AnthropicCountTokensHandler] | type[AzureAIAnthropicCountTokensHandler], -) -> None: - text_attachment: Final = "data:text/plain;charset=bogus-cs;base64," + b64encode(b"hello").decode() - messages: Final[list[dict[str, JsonValue]]] = [ - {"role": "user", "content": [{"type": "image_url", "image_url": {"url": text_attachment}}]} - ] - api_base: Final = "https://gateway.example" - handler: Final = handler_type() - with respx.mock: - route: Final = respx.post(f"{api_base}/v1/messages/count_tokens").mock( - return_value=httpx.Response(200, json={"input_tokens": 7}) - ) - with pytest.raises(AnthropicError) as exc_info: - await ( - handler.handle_count_tokens_request( - model="claude-opus-5-5", messages=messages, api_base=api_base, auth_header={"x-api-key": "test-key"} - ) if isinstance(handler, AnthropicCountTokensHandler) else handler.handle_count_tokens_request( - model="claude-opus-5-5", messages=messages, api_base=api_base, api_key="test-key" - ) - ) - - assert exc_info.value.status_code == 400 - assert "bogus-cs" in exc_info.value.message - assert not route.called - @pytest.mark.asyncio async def test_handler_sends_the_leading_system_run_as_system_not_as_a_message(httpx_transport_clients): """The wire body is what the provider judges: ``system`` carries the prompt and no message has 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 cec7950b859..075e05c586b 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 @@ -293,6 +293,24 @@ async def test_claude_count_rejects_a_lone_responses_tool_without_calling_bedroc assert bedrock.posted_hosts() == () +@pytest.mark.asyncio +async def test_claude_count_rejects_a_text_attachment_it_cannot_decode_without_calling_bedrock() -> None: + bedrock: Final = _Bedrock(runtime=(200, {"inputTokens": 412}), mantle=(200, {"input_tokens": 2177})) + text_attachment: Final = "data:text/plain;charset=bogus-cs;base64," + base64.b64encode(b"hello").decode() + + result: Final = await _counter(bedrock).count_tokens( + model_to_use="bedrock_mantle/anthropic.claude-opus-5-5", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": text_attachment}}]}], + contents=None, + deployment=DEPLOYMENT, + ) + + assert result is not None + assert (result.error, result.status_code, result.rejected_before_provider_call) == (True, 400, True) + assert "bogus-cs" in (result.error_message or "") + assert bedrock.posted_hosts() == () + + _CLAUDE_REPLY: Final = { "id": "msg_1", "type": "message", diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 43c017468ff..6caa41e9995 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2040,6 +2040,47 @@ async def test_framing_keeps_the_400_for_blocks_the_provider_counter_cannot_conv assert "search_result" in raised.value.message +@pytest.mark.parametrize( + "litellm_params", + [ + {"model": "anthropic/claude-opus-5-5", "api_key": "sk-unused-framed"}, + { + "model": "azure_ai/claude-opus-5-5", + "api_key": "sk-unused-framed", + "api_base": "https://framed.services.ai.azure.com/anthropic", + }, + ], + ids=["anthropic", "azure_ai"], +) +@pytest.mark.asyncio +async def test_framing_keeps_the_400_for_a_text_attachment_the_claude_counter_cannot_decode( + monkeypatch: pytest.MonkeyPatch, litellm_params: dict[str, str] +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + {"model_name": "framed", "litellm_params": litellm_params, "model_info": {"token_count_framing": 100}} + ] + ), + ) + text_attachment: Final = "data:text/plain;charset=bogus-cs;base64," + base64.b64encode(b"hello").decode() + + with pytest.raises(ProxyException) as raised: + await count_request_tokens( + request=TokenCountRequest( + model="framed", + messages=[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": text_attachment}}]}], + ), + call_endpoint=True, + messages_format="anthropic", + ) + + assert raised.value.code == "400" + assert "bogus-cs" in raised.value.message + + @pytest.mark.asyncio @pytest.mark.parametrize( "unreadable_body",