Merge cf98a33374 into 74050e03c5

.github/workflows/test-unit-proxy-db.yml

@@ -124,6 +124,7 @@ jobs:
               tests/proxy_unit_tests/test_server_root_path.py
               tests/proxy_unit_tests/test_proxy_pass_user_config.py
               tests/proxy_unit_tests/test_proxy_token_counter.py
+              tests/proxy_unit_tests/test_strict_token_count.py
               tests/proxy_unit_tests/test_request_size_limit_middleware.py
               tests/proxy_unit_tests/test_multipart_bypass_repro.py
             workers: 4

litellm/proxy/proxy_server.py

@@ -12073,6 +12073,56 @@ def _get_provider_token_counter(
     return None, None, None
 
 
+def _deployment_wants_strict_count(deployment: Mapping[str, Any]) -> bool:
+    """Whether one configured deployment asks for exact token counts."""
+    info: Final = deployment.get("model_info")
+    if info is None:
+        return False
+    return bool(info.get("strict_token_count", False))
+
+
+def _is_strict_token_count_model(
+    llm_router: Router | None,
+    model_name: str | None,
+    model_info: ModelMapInfo | None,
+) -> bool:
+    """Whether this model requires an exact token count.
+
+    Prefers the selected deployment's `model_info`, then falls back to the
+    router's configuration for the requested model. The fallback matters
+    because deployment selection can fail for reasons unrelated to the
+    policy, and a strict model must not quietly return an estimate then.
+
+    Note the precedence when a model group is configured inconsistently: a
+    deployment carrying `strict_token_count: false` does **not** short-circuit
+    to permissive. It falls through to the router check, so the group is
+    strict if any of its deployments asks for it. The caller cannot choose
+    which deployment serves them, so an ambiguous configuration is read the
+    safe way. Set the flag on every deployment in a group to avoid relying on
+    this.
+    """
+    if model_info is not None and bool(model_info.get("strict_token_count", False)):
+        return True
+
+    if llm_router is None or model_name is None:
+        return False
+
+    # Strict if any configured deployment for this model asks for it: the
+    # caller cannot choose which deployment serves them, so the safe reading
+    # of a mixed configuration is the strict one.
+    try:
+        deployments: Final = llm_router.get_model_list(model_name=model_name)
+        if deployments is None:
+            return False
+        return any(_deployment_wants_strict_count(d) for d in deployments)
+    except (KeyError, AttributeError, TypeError, ValueError):
+        verbose_proxy_logger.debug(
+            "litellm.proxy.proxy_server._is_strict_token_count_model(): could not list deployments for %s",
+            model_name,
+        )
+        return False
+
+
 async def _try_provider_token_count(
     provider_counter: "BaseTokenCounter",
     custom_llm_provider: str | None,
@@ -12083,8 +12133,13 @@ async def _try_provider_token_count(
     request_model: str,
     tools: list | None = None,
     system: str | None = None,
+    strict_token_count: bool = False,
 ) -> Optional["TokenCountResponse"]:
-    """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting."""
+    """Attempt provider-specific token counting. Returns result on success, None to fall through to local counting.
+
+    When `strict_token_count` is True, a failed provider count raises instead of
+    falling through to a local estimate.
+    """
     if not provider_counter.should_use_token_counting_api(custom_llm_provider=custom_llm_provider):
         return None
     try:
@@ -12107,7 +12162,7 @@ async def _try_provider_token_count(
             code=status_code,
         )
     if result is not None and result.error is True:
-        if litellm.disable_token_counter is True:
+        if strict_token_count is True:
             raise ProxyException(
                 message=result.error_message or "Token counting failed",
                 type="token_counting_error",
@@ -12178,6 +12233,23 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
         litellm_model_name or request.model
     )  # use litellm model name, if it's not avalable then fallback to request.model
 
+    #########################################################
+    # Strict token counting.
+    #
+    # `litellm.disable_token_counter` turns a failed provider count into an
+    # error instead of a local estimate, but it applies proxy-wide.
+    # `strict_token_count` in a model's `model_info` opts a single model into
+    # that same behaviour, so a deployment can require an exact count for one
+    # model without giving up the local estimate for every other model.
+    #########################################################
+    # Resolved from the router's configuration rather than the selected
+    # deployment, so the policy still holds when deployment selection fails
+    # (every deployment cooling down, rate limited, ...). Failing open there
+    # would hand back the estimate the flag exists to refuse.
+    strict_token_count: Final = litellm.disable_token_counter is True or _is_strict_token_count_model(
+        llm_router=llm_router, model_name=request.model, model_info=model_info
+    )
+
     # Try provider-specific token counting first - only for non-direct requests (from provider endpoints)
     provider_counter: BaseTokenCounter | None = None
     custom_llm_provider: str | None = None
@@ -12198,12 +12270,13 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False)
             request_model=request.model,
             tools=tools,
             system=system,
+            strict_token_count=strict_token_count,
         )
         if result is not None:
             return result
 
-    # Check if token counter is disabled before fallback
-    if litellm.disable_token_counter is True:
+    # Check if strict token counting is required before falling back to an estimate
+    if strict_token_count is True:
         raise ProxyException(
             message="Token counting is disabled and no provider API result available",
             type="token_counting_disabled",

litellm/types/router.py

@@ -149,6 +149,12 @@ class ModelInfo(MirroredPricingParams):
     base_model: str | None = None  # specify if the base model is azure/gpt-3.5-turbo etc for accurate cost tracking
     tier: Literal["free", "paid"] | None = None
 
+    # Require an exact provider token count for this model. When the provider's
+    # token counting API is unavailable or does not support the model, fail the
+    # request instead of silently returning a local tiktoken estimate.
+    # This is the per-model form of the proxy-wide `litellm.disable_token_counter`.
+    strict_token_count: bool | None = None
+
     """
     Team Model Specific Fields
     """

tests/proxy_unit_tests/test_strict_token_count.py

@@ -0,0 +1,300 @@
+"""
+Tests for per-model strict token counting (`model_info.strict_token_count`).
+
+Context: https://github.com/BerriAI/litellm/issues/37102
+
+When a provider's token counting API cannot count a model, the proxy falls back
+to a local tiktoken estimate. That estimate can be materially lower than the
+real count, and `/v1/messages/count_tokens` returns it in a shape that is
+indistinguishable from an exact count.
+
+`litellm.disable_token_counter` already turns that fallback into an error, but
+it is proxy-wide. `strict_token_count` in a model's `model_info` opts a single
+model into the same behaviour.
+"""
+
+import os
+import sys
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import litellm
+import litellm.proxy.proxy_server as proxy_server
+from litellm.llms.bedrock.common_utils import BedrockError
+from litellm.proxy._types import ProxyException, TokenCountRequest
+from litellm.proxy.proxy_server import token_counter
+from litellm.router import Router
+
+# the error Bedrock CountTokens returns for a model it cannot count
+UNSUPPORTED_MODEL_ERROR = "This model doesn't support counting tokens."
+
+
+def _router(model_info=None):
+    deployment = {
+        "model_name": "claude-opus-5",
+        "litellm_params": {
+            "model": "bedrock/anthropic.claude-opus-5-20260101-v1:0",
+            "aws_region_name": "us-east-1",
+            "aws_access_key_id": "fake",
+            "aws_secret_access_key": "fake",
+        },
+    }
+    if model_info is not None:
+        deployment["model_info"] = model_info
+    return Router(model_list=[deployment])
+
+
+def _unsupported_count_tokens():
+    """Patch Bedrock CountTokens to reject the model, as it does in the report."""
+    return patch(
+        "litellm.llms.bedrock.count_tokens.handler.BedrockCountTokensHandler.handle_count_tokens_request",
+        new=AsyncMock(side_effect=BedrockError(status_code=400, message=UNSUPPORTED_MODEL_ERROR)),
+    )
+
+
+async def _count_tokens(router):
+    original_router = getattr(proxy_server, "llm_router", None)
+    setattr(proxy_server, "llm_router", router)
+    try:
+        return await token_counter(
+            request=TokenCountRequest(
+                model="claude-opus-5",
+                messages=[{"role": "user", "content": "hello " * 400}],
+            ),
+            call_endpoint=True,
+        )
+    finally:
+        setattr(proxy_server, "llm_router", original_router)
+
+
+@pytest.mark.asyncio
+async def test_strict_token_count_raises_instead_of_estimating():
+    """A model marked strict must fail rather than return a local estimate."""
+    router = _router(model_info={"strict_token_count": True})
+
+    with _unsupported_count_tokens():
+        with pytest.raises(ProxyException) as exc_info:
+            await _count_tokens(router)
+
+    assert exc_info.value.type == "token_counting_error"
+    assert UNSUPPORTED_MODEL_ERROR in exc_info.value.message
+
+
+@pytest.mark.asyncio
+async def test_without_strict_token_count_falls_back_to_estimate():
+    """Default behaviour is unchanged: fall back to the local estimate."""
+    router = _router()
+
+    with _unsupported_count_tokens():
+        response = await _count_tokens(router)
+
+    assert response.total_tokens > 0
+    # the estimate came from a local tokenizer, not the Bedrock API
+    assert response.tokenizer_type != "bedrock_api"
+
+
+@pytest.mark.asyncio
+async def test_strict_token_count_false_is_explicitly_permissive():
+    """`strict_token_count: false` must behave like the flag being absent."""
+    router = _router(model_info={"strict_token_count": False})
+
+    with _unsupported_count_tokens():
+        response = await _count_tokens(router)
+
+    assert response.total_tokens > 0
+
+
+@pytest.mark.asyncio
+async def test_strict_token_count_does_not_affect_other_models():
+    """
+    The point of the flag: one strict model must not make every other model
+    strict. `disable_token_counter` could not express this.
+    """
+    router = Router(
+        model_list=[
+            {
+                "model_name": "strict-model",
+                "litellm_params": {
+                    "model": "bedrock/anthropic.claude-opus-5-20260101-v1:0",
+                    "aws_region_name": "us-east-1",
+                    "aws_access_key_id": "fake",
+                    "aws_secret_access_key": "fake",
+                },
+                "model_info": {"strict_token_count": True},
+            },
+            {
+                "model_name": "lenient-model",
+                "litellm_params": {
+                    "model": "bedrock/anthropic.claude-opus-5-20260101-v1:0",
+                    "aws_region_name": "us-east-1",
+                    "aws_access_key_id": "fake",
+                    "aws_secret_access_key": "fake",
+                },
+            },
+        ]
+    )
+
+    original_router = getattr(proxy_server, "llm_router", None)
+    setattr(proxy_server, "llm_router", router)
+    try:
+        with _unsupported_count_tokens():
+            # strict model refuses, with the provider's reason preserved.
+            # Asserting the type matters: it distinguishes the provider-failure
+            # path from the later generic "no provider result" checkpoint.
+            with pytest.raises(ProxyException) as exc_info:
+                await token_counter(
+                    request=TokenCountRequest(
+                        model="strict-model",
+                        messages=[{"role": "user", "content": "hello " * 400}],
+                    ),
+                    call_endpoint=True,
+                )
+            assert exc_info.value.type == "token_counting_error"
+            assert UNSUPPORTED_MODEL_ERROR in exc_info.value.message
+
+            # the other model, same provider and same failure, still estimates
+            response = await token_counter(
+                request=TokenCountRequest(
+                    model="lenient-model",
+                    messages=[{"role": "user", "content": "hello " * 400}],
+                ),
+                call_endpoint=True,
+            )
+            assert response.total_tokens > 0
+    finally:
+        setattr(proxy_server, "llm_router", original_router)
+
+
+@pytest.mark.asyncio
+async def test_strict_read_from_the_selected_deployment():
+    """Strictness on the *selected deployment* is honoured on its own.
+
+    The router's configured list here does not carry the flag, so only the
+    deployment returned by selection does. This pins the `model_info` branch
+    independently of the router-config fallback.
+    """
+    router = _router()  # configured without strict_token_count
+
+    original = Router.async_get_available_deployment
+
+    async def _strict_deployment(self, *args, **kwargs):
+        deployment = await original(self, *args, **kwargs)
+        deployment = dict(deployment)
+        deployment["model_info"] = {
+            **(deployment.get("model_info") or {}),
+            "strict_token_count": True,
+        }
+        return deployment
+
+    with patch.object(Router, "async_get_available_deployment", new=_strict_deployment):
+        with _unsupported_count_tokens():
+            with pytest.raises(ProxyException) as exc_info:
+                await _count_tokens(router)
+
+    assert exc_info.value.type == "token_counting_error"
+    assert UNSUPPORTED_MODEL_ERROR in exc_info.value.message
+
+
+@pytest.mark.asyncio
+async def test_strict_survives_deployment_selection_failure():
+    """A strict model must not fall back to an estimate when routing fails.
+
+    Deployment selection can fail for reasons unrelated to the policy - every
+    deployment cooling down, rate limited, unhealthy. The selected deployment's
+    `model_info` is unavailable then, so resolving strictness only from it would
+    hand back exactly the estimate the flag exists to refuse.
+    """
+    router = _router(model_info={"strict_token_count": True})
+
+    with patch.object(
+        Router,
+        "async_get_available_deployment",
+        new=AsyncMock(side_effect=Exception("No deployments available - cooldown")),
+    ):
+        with pytest.raises(ProxyException) as exc_info:
+            await _count_tokens(router)
+
+    assert exc_info.value.type == "token_counting_disabled"
+
+
+@pytest.mark.asyncio
+async def test_non_strict_model_still_estimates_when_selection_fails():
+    """The failure path stays permissive for models that never opted in."""
+    router = _router()
+
+    with patch.object(
+        Router,
+        "async_get_available_deployment",
+        new=AsyncMock(side_effect=Exception("No deployments available - cooldown")),
+    ):
+        response = await _count_tokens(router)
+
+    assert response.total_tokens > 0
+
+
+@pytest.mark.asyncio
+async def test_mixed_group_is_strict_if_any_deployment_asks():
+    """A group configured inconsistently resolves to strict.
+
+    A deployment carrying `strict_token_count: false` does not make the group
+    permissive when a sibling asks for strict. The caller cannot choose which
+    deployment serves them, so the ambiguous configuration is read the safe
+    way. Documented on `_is_strict_token_count_model` because it is
+    non-obvious. See https://github.com/BerriAI/litellm/issues/37102.
+    """
+    router = Router(
+        model_list=[
+            {
+                "model_name": "mixed-model",
+                "litellm_params": {
+                    "model": "bedrock/anthropic.claude-opus-5-20260101-v1:0",
+                    "aws_region_name": "us-east-1",
+                    "aws_access_key_id": "fake",
+                    "aws_secret_access_key": "fake",
+                },
+                "model_info": {"strict_token_count": False},
+            },
+            {
+                "model_name": "mixed-model",
+                "litellm_params": {
+                    "model": "bedrock/anthropic.claude-opus-5-20260101-v1:0",
+                    "aws_region_name": "us-west-2",
+                    "aws_access_key_id": "fake",
+                    "aws_secret_access_key": "fake",
+                },
+                "model_info": {"strict_token_count": True},
+            },
+        ]
+    )
+
+    original_router = getattr(proxy_server, "llm_router", None)
+    setattr(proxy_server, "llm_router", router)
+    try:
+        with _unsupported_count_tokens():
+            with pytest.raises(ProxyException):
+                await token_counter(
+                    request=TokenCountRequest(
+                        model="mixed-model",
+                        messages=[{"role": "user", "content": "hello " * 400}],
+                    ),
+                    call_endpoint=True,
+                )
+    finally:
+        setattr(proxy_server, "llm_router", original_router)
+
+
+@pytest.mark.asyncio
+async def test_disable_token_counter_still_applies_proxy_wide():
+    """The existing proxy-wide flag must keep working for unmarked models."""
+    router = _router()
+    original = litellm.disable_token_counter
+    litellm.disable_token_counter = True
+    try:
+        with _unsupported_count_tokens():
+            with pytest.raises(ProxyException):
+                await _count_tokens(router)
+    finally:
+        litellm.disable_token_counter = original

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -37890,6 +37890,8 @@ export interface components {
             ptu_effective_from?: string | null;
             /** Ptu Effective To */
             ptu_effective_to?: string | null;
+            /** Strict Token Count */
+            strict_token_count?: boolean | null;
             /** Team Id */
             team_id?: string | null;
             /** Team Public Model Name */