From e7392252685216426aa4a4e0857eddc5f6e2961d Mon Sep 17 00:00:00 2001 From: aliabbas-muhammadi <143003167+aliabbas-muhammadi@users.noreply.github.com> Date: Tue, 18 Aug 2026 16:21:53 +1000 Subject: [PATCH] fix: mark /v1/messages/count_tokens local-fallback counts as estimates --- .../proxy/anthropic_endpoints/endpoints.py | 24 ++++++- .../test_proxy_token_counter.py | 11 ++- .../anthropic_endpoints/test_endpoints.py | 67 +++++++++++++++++++ 3 files changed, 98 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index f75899b91dc..f59bf0e10b3 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -287,8 +287,28 @@ async def count_tokens( elif isinstance(token_response, dict): _token_response_dict = token_response - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} + # Convert the internal response to Anthropic API format. + # + # Only a provider CountTokens API (e.g. Bedrock/Anthropic) returns an exact + # count, and every such path records the upstream reply in `original_response`. + # When it is absent the count came from litellm's local tokenizer and is an + # APPROXIMATION -- for models Bedrock's CountTokens API rejects (e.g. Claude + # Opus 5 / Sonnet 5) that tokenizer is miscalibrated and understates the count. + # Flag the estimate so callers doing context management do not treat it as + # authoritative and overrun the model's window. + # https://github.com/BerriAI/litellm/issues/37102 + is_estimate = _token_response_dict.get("original_response") is None + return { + "input_tokens": _token_response_dict.get("total_tokens", 0), + **( + { + "litellm_estimate": True, + "litellm_tokenizer_used": _token_response_dict.get("tokenizer_type"), + } + if is_estimate + else {} + ), + } except HTTPException: raise diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index 1079a5228a1..0e611fe75e9 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -282,7 +282,11 @@ async def test_anthropic_messages_count_tokens_endpoint(): assert isinstance(response, dict) assert "input_tokens" in response assert response["input_tokens"] == 15 - assert len(response) == 1 # Should only contain input_tokens + # No original_response on the mock => local tokenizer was used, so the + # count is surfaced as an estimate rather than an authoritative figure + # (issue #37102). + assert response["litellm_estimate"] is True + assert response["litellm_tokenizer_used"] == "openai_tokenizer" print("✅ Anthropic endpoint test passed!") @@ -354,7 +358,10 @@ async def test_anthropic_messages_count_tokens_with_non_anthropic_model(): assert isinstance(response, dict) assert "input_tokens" in response assert response["input_tokens"] == 12 - assert len(response) == 1 # Should only contain input_tokens + # Local tokenizer (no original_response) => flagged as an estimate + # (issue #37102). + assert response["litellm_estimate"] is True + assert response["litellm_tokenizer_used"] == "openai_tokenizer" print("✅ Non-Anthropic model test passed!") diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py index 69d90a8b59b..08962c2478f 100644 --- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py @@ -221,6 +221,73 @@ class TestStripTotalTokens(unittest.TestCase): assert response.usage == {"input_tokens": 100, "output_tokens": 50} +class TestCountTokensEstimateSignal: + """`/v1/messages/count_tokens` must not pass off a local-tokenizer estimate as + an authoritative count (issue #37102). + + When Bedrock's CountTokens API does not support a model (e.g. Claude Opus 5 / + Sonnet 5) litellm silently falls back to a local tokenizer, which understates + the count. The provider path records the upstream reply in `original_response`; + the local fallback leaves it None. The endpoint keys off that to flag estimates. + """ + + async def _call_endpoint(self, token_count_response): + from unittest.mock import AsyncMock, MagicMock + + import litellm.proxy.anthropic_endpoints.endpoints as ep + import litellm.proxy.proxy_server as proxy_server + + request_body = { + "model": "claude-opus-5", + "messages": [{"role": "user", "content": "hi"}], + } + with ( + patch.object(ep, "_read_request_body", new=AsyncMock(return_value=request_body)), + patch.object(proxy_server, "token_counter", new=AsyncMock(return_value=token_count_response)), + ): + return await ep.count_tokens(request=MagicMock(), user_api_key_dict=MagicMock()) + + @pytest.mark.asyncio + async def test_local_fallback_is_flagged_as_estimate(self): + """Bedrock CountTokens unsupported -> local tokenizer fallback -> the + endpoint must mark the count as an estimate instead of returning a bare + 200 that looks authoritative.""" + from litellm.types.utils import TokenCountResponse + + # What the internal token_counter returns after Bedrock rejects the model + # and it falls back to the local tokenizer: note original_response is None. + local_fallback = TokenCountResponse( + total_tokens=11208, + request_model="claude-opus-5", + model_used="claude-opus-5", + tokenizer_type="huggingface_tokenizer", + ) + + response = await self._call_endpoint(local_fallback) + + assert response["input_tokens"] == 11208 + assert response["litellm_estimate"] is True + assert response["litellm_tokenizer_used"] == "huggingface_tokenizer" + + @pytest.mark.asyncio + async def test_authoritative_provider_count_is_not_flagged(self): + """When Bedrock's CountTokens API actually answered (original_response set), + the count is exact and the response stays the bare Anthropic shape.""" + from litellm.types.utils import TokenCountResponse + + authoritative = TokenCountResponse( + total_tokens=26, + request_model="claude-sonnet-4-6", + model_used="claude-sonnet-4-6", + tokenizer_type="bedrock_api", + original_response={"inputTokens": 26}, + ) + + response = await self._call_endpoint(authoritative) + + assert response == {"input_tokens": 26} + + class TestStripTotalTokensFeatureFlag(unittest.TestCase): """The strip is gated behind `litellm.strip_anthropic_total_tokens`.