diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 3725e0f5805..e7ce6fa2c67 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7dad7abd210..a45c3a181c8 100644 --- a/litellm/proxy/proxy_server.py +++ b/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", diff --git a/litellm/types/router.py b/litellm/types/router.py index a3335be2b2b..0de91daee4f 100644 --- a/litellm/types/router.py +++ b/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 """ diff --git a/tests/proxy_unit_tests/test_strict_token_count.py b/tests/proxy_unit_tests/test_strict_token_count.py new file mode 100644 index 00000000000..6d30f907582 --- /dev/null +++ b/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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 24ed24512dc..8f9f93897fe 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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 */