feat(count_tokens): reject unreliable local estimates

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/common_utils/token_count_settings.py

@@ -1,7 +1,7 @@
-from collections.abc import Callable, Mapping
+from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass, replace
 from types import MappingProxyType
-from typing import Final, Literal
+from typing import Final, Literal, TypeGuard
 
 from pydantic import BaseModel, ConfigDict, ValidationError
 
@@ -77,6 +77,67 @@ class ProviderCountFailure:
         )
 
 
+@dataclass(frozen=True, slots=True)
+class _UnreliableLocalEstimateRule:
+    provider: str
+    model_substring: str
+    trigger: Literal["tools", "no_tools"]
+
+
+_UNRELIABLE_LOCAL_ESTIMATE_RULES: Final[tuple[_UnreliableLocalEstimateRule, ...]] = (
+    _UnreliableLocalEstimateRule("bedrock", "nova-2-lite", "tools"),
+    _UnreliableLocalEstimateRule("azure_ai", "deepseek-v4", "tools"),
+    _UnreliableLocalEstimateRule("azure_ai", "grok-4", "no_tools"),
+)
+
+
+def _content_blocks(message: object) -> tuple[object, ...]:
+    message_mapping: Final = as_str_mapping(message)
+    if message_mapping is None:
+        return ()
+    content: Final = message_mapping.get("content")
+    if not _is_object_sequence(content):
+        return ()
+    return tuple(content)
+
+
+def _is_object_sequence(value: object) -> TypeGuard[Sequence[object]]:
+    return isinstance(value, Sequence) and not isinstance(value, (str, bytes))
+
+
+def _is_document_content_block(block: object) -> bool:
+    block_mapping: Final = as_str_mapping(block)
+    return block_mapping is not None and block_mapping.get("type") in {"document", "file"}
+
+
+def _has_document_content_block(messages: Sequence[object] | None) -> bool:
+    return any(
+        any(_is_document_content_block(block) for block in _content_blocks(message)) for message in messages or ()
+    )
+
+
+def unreliable_local_estimate(
+    custom_llm_provider: str | None,
+    model: str,
+    messages: Sequence[object] | None,
+    tools: Sequence[object] | None,
+) -> str | None:
+    if _has_document_content_block(messages):
+        return f"{model} has no server-side count_tokens and LiteLLM cannot count document content locally"
+    trigger: Final = "tools" if tools else "no_tools"
+    matched_rule: Final = next(
+        (
+            rule
+            for rule in _UNRELIABLE_LOCAL_ESTIMATE_RULES
+            if custom_llm_provider == rule.provider
+            and rule.model_substring in model.casefold()
+            and rule.trigger == trigger
+        ),
+        None,
+    )
+    return None if matched_rule is None else f"Local token estimate for {model} is known to be off by more than 5%"
+
+
 def cost_map_entry(model: str, custom_llm_provider: str | None = None) -> Mapping[str, object]:
     try:
         return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)

litellm/proxy/proxy_server.py

@@ -151,6 +151,9 @@ from litellm.proxy.common_utils.token_count_settings import (
     cost_map_entry,
     resolve_token_count_settings,
 )
+from litellm.proxy.common_utils.token_count_settings import (
+    unreliable_local_estimate as _unreliable_local_estimate,
+)
 from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog
 from litellm.router_utils.access_windows import access_windows_config_error
 from litellm.router_utils.add_retry_fallback_headers import (
@@ -14110,6 +14113,21 @@ async def count_request_tokens(
     if isinstance(ladder, TokenCountResponse):
         return ladder
 
+    if call_endpoint is True:
+        unreliable_reason: Final = _unreliable_local_estimate(
+            custom_llm_provider=custom_llm_provider,
+            model=model_to_use,
+            messages=messages,
+            tools=tools,
+        )
+        if unreliable_reason is not None:
+            raise ProxyException(
+                message=unreliable_reason,
+                type="invalid_request_error",
+                param="model",
+                code=400,
+            )
+
     # Check if token counter is disabled before fallback
     if litellm.disable_token_counter is True:
         raise ProxyException(

tests/unit/proxy/test_proxy_token_counter.py

@@ -58,6 +58,7 @@ from litellm.llms.bedrock.count_tokens.kimi_k3 import (
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler
 from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter
+from litellm.proxy import proxy_server
 from litellm.proxy._types import ProxyException, TokenCountRequest
 from litellm.proxy.anthropic_endpoints.endpoints import (
     count_tokens as anthropic_count_tokens,
@@ -67,6 +68,7 @@ from litellm.proxy.common_utils.token_count_settings import (
     TokenCountModel,
     TokenCountSettings,
     cost_map_entry,
+    unreliable_local_estimate,
 )
 from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request
 from litellm.proxy.proxy_server import (
@@ -2637,3 +2639,169 @@ async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_need
         local.total_tokens,
         local.tokenizer_type,
     )
+
+
+@pytest.mark.parametrize(
+    ("provider", "model", "messages", "tools", "expected"),
+    [
+        ("bedrock", "meta.llama4-maverick", [{"content": [{"type": "document"}]}], None, "document content"),
+        ("bedrock", "us.amazon.nova-2-lite", [], [{"name": "search"}], "more than 5%"),
+        ("azure_ai", "DeepSeek-V4-Pro", [], [{"name": "search"}], "more than 5%"),
+        ("azure_ai", "grok-4-1-fast-reasoning", [], None, "more than 5%"),
+    ],
+)
+def test_unreliable_local_estimate_returns_reason(
+    provider: str,
+    model: str,
+    messages: list[object],
+    tools: list[object] | None,
+    expected: str,
+) -> None:
+    reason: Final = unreliable_local_estimate(provider, model, messages, tools)
+
+    assert reason is not None and expected in reason
+
+
+@pytest.mark.parametrize(
+    ("provider", "model", "messages", "tools"),
+    [
+        ("bedrock", "meta.llama4-maverick", [{"content": "text"}], None),
+        ("bedrock", "us.amazon.nova-2-lite", [], None),
+        ("azure_ai", "deepseek-v4-pro", [], None),
+        ("azure_ai", "grok-4-1-fast-reasoning", [], [{"name": "search"}]),
+        ("azure", "gpt-4.1", [{"content": "text"}], None),
+    ],
+)
+def test_unreliable_local_estimate_preserves_local_count(
+    provider: str, model: str, messages: list[object], tools: list[object] | None
+) -> None:
+    assert unreliable_local_estimate(provider, model, messages, tools) is None
+
+
+@pytest.mark.parametrize(
+    ("resolved_model", "messages", "tools", "call_endpoint", "reject"),
+    [
+        ("bedrock/meta.llama4-maverick", [{"role": "user", "content": [{"type": "document"}]}], None, True, True),
+        ("bedrock/meta.llama4-maverick", [{"role": "user", "content": "text"}], None, True, False),
+        ("bedrock/us.amazon.nova-2-lite", [{"role": "user", "content": "text"}], [{"name": "search"}], True, True),
+        ("bedrock/us.amazon.nova-2-lite", [{"role": "user", "content": "text"}], None, True, False),
+        ("azure_ai/deepseek-v4-pro", [{"role": "user", "content": "text"}], [{"name": "search"}], True, True),
+        ("azure_ai/deepseek-v4-pro", [{"role": "user", "content": "text"}], None, True, False),
+        ("azure_ai/grok-4-1-fast-reasoning", [{"role": "user", "content": "text"}], None, True, True),
+        ("azure_ai/grok-4-1-fast-reasoning", [{"role": "user", "content": "text"}], [{"name": "search"}], True, False),
+        (
+            "bedrock/meta.llama4-maverick",
+            [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "document",
+                            "source": {
+                                "type": "base64",
+                                "media_type": "application/pdf",
+                                "data": "JVBERi0xLjQK",
+                            },
+                        }
+                    ],
+                }
+            ],
+            None,
+            False,
+            False,
+        ),
+        ("azure/gpt-4.1", [{"role": "user", "content": "text"}], None, True, False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_count_request_tokens_rejects_only_known_unreliable_estimates(
+    monkeypatch: pytest.MonkeyPatch,
+    resolved_model: str,
+    messages: list[dict[str, object]],
+    tools: list[dict[str, object]] | None,
+    call_endpoint: bool,
+    reject: bool,
+) -> None:
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "count-target",
+                "litellm_params": {"model": resolved_model, "api_key": "test-key", "api_base": "https://count.test"},
+            }
+        ]
+    )
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    failure: Final = ProviderCountFailure(status_code=400, message="unsupported")
+    with patch.object(
+        proxy_server,
+        "_provider_token_count",
+        new=AsyncMock(return_value=failure),
+    ):
+        if reject:
+            with pytest.raises(ProxyException) as raised:
+                await count_request_tokens(
+                    request=TokenCountRequest(model="count-target", messages=messages, tools=tools),
+                    call_endpoint=call_endpoint,
+                    messages_format="anthropic",
+                )
+            assert raised.value.code == "400"
+        else:
+            response: Final = await count_request_tokens(
+                request=TokenCountRequest(model="count-target", messages=messages, tools=tools),
+                call_endpoint=call_endpoint,
+                messages_format="anthropic",
+            )
+            assert response.count_source == "local_estimate"
+
+
+@pytest.mark.asyncio
+async def test_document_block_keeps_successful_provider_count(monkeypatch: pytest.MonkeyPatch) -> None:
+    router: Final = Router(
+        model_list=[
+            {
+                "model_name": "count-target",
+                "litellm_params": {
+                    "model": "bedrock/us.anthropic.claude-sonnet-4-6",
+                    "api_key": "test-key",
+                    "api_base": "https://count.test",
+                },
+            }
+        ]
+    )
+    monkeypatch.setattr(proxy_server, "llm_router", router)
+    provider_result: Final = TokenCountResponse(
+        total_tokens=17,
+        request_model="count-target",
+        model_used="anthropic.claude-sonnet-4-6",
+        tokenizer_type="openai_tokenizer",
+        count_source="provider",
+    )
+    with patch.object(
+        proxy_server,
+        "_provider_token_count",
+        new=AsyncMock(return_value=provider_result),
+    ):
+        result: Final = await count_request_tokens(
+            request=TokenCountRequest(
+                model="count-target",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "document",
+                                "source": {
+                                    "type": "base64",
+                                    "media_type": "application/pdf",
+                                    "data": "JVBERi0xLjQK",
+                                },
+                            }
+                        ],
+                    }
+                ],
+            ),
+            call_endpoint=True,
+            messages_format="anthropic",
+        )
+
+    assert result is provider_result