From 020ded24fef71163b71a628ce1d4ba79fd70a44d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 03:05:57 -0700 Subject: [PATCH] fix(router): count pre-call context-window tokens with each deployment's tokenizer With pre-call checks on, every prompt was counted with gpt-3.5-turbo's tokenizer, so a prompt the provider accepts could get a local context-window 400. Each deployment now counts with its own model's tokenizer when litellm has one, and otherwise with an o200k_base stand-in that refuses locally only past 2x the window. A group with a deployment that fits still routes to it. --- litellm/litellm_core_utils/token_counter.py | 51 ++++- litellm/router.py | 202 +++++++++++------- .../complexity_router/context_compaction.py | 15 +- litellm/router_utils/context_window_fit.py | 47 ++++ .../litellm_core_utils/test_token_counter.py | 10 + tests/unit/test_router/test_router.py | 138 +++++++++--- .../unit/test_router_model_cost_isolation.py | 20 +- 7 files changed, 374 insertions(+), 109 deletions(-) create mode 100644 litellm/router_utils/context_window_fit.py diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index f42e4eae08f..a690455dc28 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -984,7 +984,12 @@ def _get_exact_count_function( return count_tokens if tokenizer_json["type"] == "openai_tokenizer": - encoding: Final = openai_tokenizer_encoding(model) + encoding: Final = ( + tokenizer_json["tokenizer"] + if custom_tokenizer is not None + and isinstance(tokenizer_json["tokenizer"], (tiktoken.Encoding, OpenAIEncoding)) + else openai_tokenizer_encoding(model) + ) def encode_length(text: str) -> int: return _encoding_count(encoding, text) @@ -1008,14 +1013,42 @@ def openai_tokenizer_encoding_name(model: str) -> str: """The tiktoken encoding name for `model`, without loading the encoding.""" from litellm.utils import print_verbose - model_to_use: Final = _fix_model_name(model) + own_encoding: Final = _own_openai_encoding_name(model) + if own_encoding is not None: + return own_encoding + print_verbose("Warning: model not found. Using cl100k_base encoding.") + return "cl100k_base" + + +def _own_openai_encoding_name(model: str) -> str | None: + model_to_use: Final = _known_openai_model_name(model) + if model_to_use is None: + return None if "gpt-4o" in model_to_use: return "o200k_base" try: return tiktoken.encoding_name_for_model(model_to_use) except KeyError: - print_verbose("Warning: model not found. Using cl100k_base encoding.") - return "cl100k_base" + return None + + +def counts_with_own_tokenizer(model: str) -> bool: + """Whether `token_counter` counts `model` with that model's own tokenizer rather than a stand-in encoding. + + The bundled Anthropic tokenizer is Claude 2's, which overcounts current Claude models on some scripts.""" + from litellm.utils import select_tokenizer, uses_anthropic_tokenizer + + if uses_anthropic_tokenizer(model): + return False + return select_tokenizer(model)["type"] == "huggingface_tokenizer" or _own_openai_encoding_name(model) is not None + + +def own_tokenizer_model(model: str) -> str | None: + """The name `token_counter` counts `model` by with that model's own tokenizer, trying it without its provider + prefix too, or None when litellm has no tokenizer of the model's own.""" + provider, _, bare_model = model.partition("/") + candidates: Final = (model, bare_model) if bare_model and provider in litellm.provider_list else (model,) + return next((candidate for candidate in candidates if counts_with_own_tokenizer(candidate)), None) def uses_legacy_message_accounting(model: str) -> bool: @@ -1025,13 +1058,17 @@ def uses_legacy_message_accounting(model: str) -> bool: def _fix_model_name(model: str) -> str: """We normalize some model names to others""" + known_model: Final = _known_openai_model_name(model) + return known_model if known_model is not None else "gpt-3.5-turbo" + + +def _known_openai_model_name(model: str) -> str | None: if model in litellm.azure_llms: # azure llms use gpt-35-turbo instead of gpt-3.5-turbo 🙃 return model.replace("-35", "-3.5") - elif model in litellm.open_ai_chat_completion_models: + if model in litellm.open_ai_chat_completion_models: return model - else: - return "gpt-3.5-turbo" + return None def _count_image_tokens( diff --git a/litellm/router.py b/litellm/router.py index b63130c3313..59642bcb653 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -119,6 +119,7 @@ from litellm.litellm_core_utils.batch_utils import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, coerce_token_limit, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -162,7 +163,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, mask_sensitive_structure, ) -from litellm.litellm_core_utils.token_counter import offload_token_count +from litellm.litellm_core_utils.token_counter import counts_with_own_tokenizer, offload_token_count, own_tokenizer_model from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception from litellm.llms.base_llm.passthrough.transformation import replace_path_segment @@ -256,6 +257,13 @@ from litellm.router_utils.common_utils import ( without_model_capability_flags, without_router_only_kwargs, ) +from litellm.router_utils.context_window_fit import ( + ContextWindow, + ContextWindowOverflow, + PreCallTokenCount, + context_window_overflow, + stand_in_tokenizer, +) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( DEFAULT_COOLDOWN_TIME_SECONDS, @@ -460,6 +468,7 @@ from litellm.types.utils import ( LlmProviders, ModelInfo, ModelResponseStream, + SelectTokenizerResponse, StandardLoggingPayload, StandardLoggingRoutingDecision, Usage, @@ -14301,8 +14310,10 @@ class Router: def _count_pre_call_check_tokens( self, messages: Sequence[Mapping[str, object]] | None, - input: str | list[object] | None, + input: str | Sequence[object] | None, request_kwargs: Mapping[str, object] | None = None, + model: str = "", + custom_tokenizer: SelectTokenizerResponse | None = None, ) -> int: """ Count input tokens for context-window pre-call checks. @@ -14333,7 +14344,9 @@ class Router: system_message: Final = anthropic_system_to_openai_message(extras.get("system")) if messages is not None: counted_messages: Final = (system_message, *messages) if system_message is not None else messages - return litellm.token_counter(messages=counted_messages, tools=tools) + return litellm.token_counter( + model=model, custom_tokenizer=custom_tokenizer, messages=counted_messages, tools=tools + ) if input is not None: from openai.types.responses.response_create_params import ResponseInputParam @@ -14348,13 +14361,15 @@ class Router: only_since_compaction=True, ) return litellm.token_counter( + model=model, + custom_tokenizer=custom_tokenizer, messages=cast(list, input_messages), # cast-ok: transformed chat messages tools=tools, ) raise ValueError("Either messages or input must be provided to count tokens") - def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None: - """The deployment's declared context window, or None when it declares none or cannot be resolved.""" + def _deployment_context_window(self, model: str, deployment: Mapping[str, object]) -> ContextWindow | None: + """The deployment's model and declared context window, or None when it declares none or cannot be resolved.""" try: model_info: Final = self.get_router_model_info( deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts @@ -14362,22 +14377,59 @@ class Router: ) except Exception as e: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others verbose_router_logger.debug( - "litellm.router.py::_deployment_max_input_tokens: skipping deployment. Got - %s", e + "litellm.router.py::_deployment_context_window: skipping deployment. Got - %s", e ) return None max_input_tokens: Final = model_info.get("max_input_tokens") - return max_input_tokens if isinstance(max_input_tokens, int) else None + if not isinstance(max_input_tokens, int): + return None + return ContextWindow( + tokenizer_model=own_tokenizer_model(model_info["key"]) or model_info["key"], + max_input_tokens=max_input_tokens, + ) - def _pre_call_checks_need_token_count( + def _pre_call_check_tokenizer_models( self, model: str, healthy_deployments: Sequence[Mapping[str, object]] - ) -> bool: - """Whether any healthy deployment declares a context window that a token count could exceed. + ) -> frozenset[str]: + """The models whose tokenizers must count the prompt: one per deployment that declares a context window. Resolves each deployment the way ``_pre_call_checks`` does, so one unmappable deployment cannot hide a later one that does declare a limit. """ - return any( - self._deployment_max_input_tokens(model, deployment) is not None for deployment in healthy_deployments + windows: Final = tuple(self._deployment_context_window(model, deployment) for deployment in healthy_deployments) + return frozenset(window.tokenizer_model for window in windows if window is not None) + + def _pre_call_check_token_count( + self, + tokenizer_model: str, + messages: Sequence[Mapping[str, object]] | None, + input: str | Sequence[object] | None, + request_kwargs: Mapping[str, object] | None, + ) -> PreCallTokenCount: + own_tokenizer: Final = counts_with_own_tokenizer(tokenizer_model) + return PreCallTokenCount( + tokens=self._count_pre_call_check_tokens( + messages=messages, + input=input, + request_kwargs=request_kwargs, + model=tokenizer_model, + custom_tokenizer=None if own_tokenizer else stand_in_tokenizer(), + ), + own_tokenizer=own_tokenizer, + ) + + def _pre_call_check_token_counts( + self, + tokenizer_models: frozenset[str], + messages: Sequence[Mapping[str, object]] | None, + input: str | Sequence[object] | None, + request_kwargs: Mapping[str, object] | None, + ) -> Mapping[str, PreCallTokenCount]: + return MappingProxyType( + { + tokenizer_model: self._pre_call_check_token_count(tokenizer_model, messages, input, request_kwargs) + for tokenizer_model in tokenizer_models + } ) async def _acount_pre_call_check_tokens( @@ -14387,7 +14439,7 @@ class Router: messages: Sequence[Mapping[str, str]] | None, input: str | Sequence[object] | None, request_kwargs: Mapping[str, object] | None, - ) -> int | None: + ) -> Mapping[str, PreCallTokenCount] | None: """Count input tokens off the event loop, so a multi-MB prompt cannot stall the proxy. Returns None when no deployment limits its context window, and when counting fails. The @@ -14398,11 +14450,13 @@ class Router: if messages is None and input is None: return None try: - if not self._pre_call_checks_need_token_count(model, healthy_deployments): + tokenizer_models: Final = self._pre_call_check_tokenizer_models(model, healthy_deployments) + if not tokenizer_models: return None - return await offload_token_count(self._count_pre_call_check_tokens)( - messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter - input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter + return await offload_token_count(self._pre_call_check_token_counts)( + tokenizer_models=tokenizer_models, + messages=messages, + input=input, request_kwargs=request_kwargs, ) except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request @@ -14411,22 +14465,33 @@ class Router: ) return None - def _resolve_pre_call_check_input_tokens( + @staticmethod + def _context_window_exceeded_line(deployment: Mapping[str, object], overflow: ContextWindowOverflow) -> str: + model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + litellm_params: Final = as_str_mapping(deployment.get("litellm_params")) or {} + deployment_model: Final = ( + model_info.get("base_model") or litellm_params.get("base_model") or litellm_params.get("model") + ) + return f"Model={deployment_model}, Max Input Tokens={overflow.max_input_tokens}, Got={overflow.tokens}" + + def _resolve_pre_call_check_token_counts( self, messages: Sequence[Mapping[str, str]] | None, input: str | Sequence[object] | None, request_kwargs: Mapping[str, object] | None, - input_token_count: int | None, + token_counts: Mapping[str, PreCallTokenCount] | None, + tokenizer_models: frozenset[str], count_inline: bool, - ) -> int | None: - if input_token_count is not None: - return input_token_count + ) -> Mapping[str, PreCallTokenCount] | None: + if token_counts is not None: + return token_counts if not count_inline: return None try: - return self._count_pre_call_check_tokens( - messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter - input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter + return self._pre_call_check_token_counts( + tokenizer_models=tokenizer_models, + messages=messages, + input=input, request_kwargs=request_kwargs, ) except Exception as e: @@ -14585,7 +14650,7 @@ class Router: messages: list[dict[str, str]] | None = None, input: str | list | None = None, request_kwargs: dict | None = None, - input_token_count: int | None = None, + input_token_counts: Mapping[str, PreCallTokenCount] | None = None, skip_inline_token_count: bool = False, request_deployments: Sequence[DeploymentTypedDict] | None = None, ): @@ -14603,25 +14668,23 @@ class Router: # We only pop from the list, not modify deployment dicts - 100x+ faster on hot path (every request) _returned_deployments = list(healthy_deployments) - invalid_model_indices: Final = set[int]() # Use set for O(1) membership checks - - _context_window_error = False - _potential_error_str = "" - has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( request_kwargs ) - deployment_max_input_tokens: Final = tuple( - self._deployment_max_input_tokens(model, deployment) for deployment in _returned_deployments + context_windows: Final = tuple( + self._deployment_context_window(model, deployment) for deployment in _returned_deployments ) - input_tokens: Final = self._resolve_pre_call_check_input_tokens( + tokenizer_models: Final = frozenset(window.tokenizer_model for window in context_windows if window is not None) + token_counts: Final = self._resolve_pre_call_check_token_counts( messages=messages, input=cast(str | list[object] | None, input), # cast-ok: forwarded to the resolver request_kwargs=cast(Mapping[str, object] | None, request_kwargs), # cast-ok: forwarded to the resolver - input_token_count=input_token_count, - count_inline=has_countable_input - and not skip_inline_token_count - and any(limit is not None for limit in deployment_max_input_tokens), + token_counts=input_token_counts, + tokenizer_models=tokenizer_models, + count_inline=has_countable_input and not skip_inline_token_count and bool(tokenizer_models), + ) + overflows: Final = tuple( + context_window_overflow(window, token_counts if has_countable_input else None) for window in context_windows ) request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} requested_params: Final[frozenset[str]] = ( @@ -14629,40 +14692,35 @@ class Router: if request_kwargs is not None and litellm.drop_params is False else frozenset() ) - - for idx, deployment in enumerate(_returned_deployments): - # Cache nested dict access to avoid repeated temporary dict allocations - _litellm_params = deployment.get("litellm_params", {}) - _model_info = deployment.get("model_info", {}) - - max_input_tokens = deployment_max_input_tokens[idx] - if ( - max_input_tokens is not None - and has_countable_input - and input_tokens is not None - and input_tokens > max_input_tokens - ): - invalid_model_indices.add(idx) - _context_window_error = True - deployment_model = ( - _model_info.get("base_model") or _litellm_params.get("base_model") or _litellm_params.get("model") + allowed_model_region: Final = request_settings.get("allowed_model_region") + region_excluded: Final = ( + frozenset() + if not isinstance(allowed_model_region, str) + else frozenset( + idx + for idx, deployment in enumerate(_returned_deployments) + if not is_region_allowed( + litellm_params=LiteLLM_Params.model_validate(deployment.get("litellm_params", {})), + allowed_model_region=allowed_model_region, ) - _potential_error_str += ( - f"Model={deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" - ) - continue - - ## REGION CHECK ## - if request_kwargs is not None and request_kwargs.get("allowed_model_region") is not None: - allowed_model_region = request_kwargs.get("allowed_model_region") - - if allowed_model_region is not None: - if not is_region_allowed( - litellm_params=LiteLLM_Params.model_validate(_litellm_params), - allowed_model_region=allowed_model_region, - ): - invalid_model_indices.add(idx) - continue + ) + ) + some_deployment_fits: Final = any( + overflow is None for idx, overflow in enumerate(overflows) if idx not in region_excluded + ) + window_excluded: Final = MappingProxyType( + { + idx: overflow + for idx, overflow in enumerate(overflows) + if overflow is not None and (overflow.certain or some_deployment_fits) + } + ) + _context_window_error: Final = bool(window_excluded) + _potential_error_str: Final = "".join( + self._context_window_exceeded_line(_returned_deployments[idx], window_excluded[idx]) + for idx in sorted(window_excluded) + ) + invalid_model_indices: Final = frozenset(window_excluded) | region_excluded typed_deployments: Final = tuple( cast(DeploymentTypedDict, deployment) # cast-ok: router deployments are plain dicts @@ -15214,7 +15272,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, - input_token_count=await self._acount_pre_call_check_tokens( + input_token_counts=await self._acount_pre_call_check_tokens( model=model, healthy_deployments=deployments_to_check, messages=messages, diff --git a/litellm/router_strategy/complexity_router/context_compaction.py b/litellm/router_strategy/complexity_router/context_compaction.py index 6a1a63ea4b6..cc6da7f420e 100644 --- a/litellm/router_strategy/complexity_router/context_compaction.py +++ b/litellm/router_strategy/complexity_router/context_compaction.py @@ -31,6 +31,7 @@ from litellm.router_strategy.complexity_router.config import ContextCompactionCo if TYPE_CHECKING: from litellm.router import Router + from litellm.router_utils.context_window_fit import PreCallTokenCount Surface: TypeAlias = Literal["chat", "messages", "responses"] _SURFACES: Final[Mapping[str, Surface]] = MappingProxyType( @@ -387,11 +388,23 @@ async def _check_client_managed_admission( messages=_ITEMS.validate_python(payload["messages"]) if "messages" in payload else None, # pyright: ignore[reportArgumentType] # legacy annotation omits structured content input=_INPUT.validate_python(payload.get("input")), request_kwargs=_DICT.validate_python(payload), - input_token_count=await _count(router, payload), + input_token_counts=await _admission_token_counts(router, deployment, payload), skip_inline_token_count=True, ) +async def _admission_token_counts( + router: Router, deployment: Mapping[str, object], payload: Mapping[str, object] +) -> Mapping[str, PreCallTokenCount]: + return await asyncio.to_thread( + router._pre_call_check_token_counts, # pyright: ignore[reportPrivateUsage] # shared Router admission counter + tokenizer_models=router._pre_call_check_tokenizer_models(str(payload["model"]), (deployment,)), # pyright: ignore[reportPrivateUsage] # shared Router admission counter + messages=_ITEMS.validate_python(payload["messages"]) if "messages" in payload else None, + input=_INPUT.validate_python(payload.get("input")), + request_kwargs=payload, + ) + + def _portable_history( payload: Mapping[str, object], surface: Surface ) -> tuple[tuple[Mapping[str, object], ...], tuple[Mapping[str, object], ...], tuple[Mapping[str, object], ...]]: diff --git a/litellm/router_utils/context_window_fit.py b/litellm/router_utils/context_window_fit.py new file mode 100644 index 00000000000..2ebba266896 --- /dev/null +++ b/litellm/router_utils/context_window_fit.py @@ -0,0 +1,47 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from litellm.rust_bridge.tokenizer import get_encoding +from litellm.types.utils import SelectTokenizerResponse + +STAND_IN_ENCODING: Final = "o200k_base" +STAND_IN_OVERCOUNT_MARGIN: Final = 2 + + +@dataclass(frozen=True, slots=True) +class ContextWindow: + tokenizer_model: str + max_input_tokens: int + + +@dataclass(frozen=True, slots=True) +class PreCallTokenCount: + tokens: int + own_tokenizer: bool + + +@dataclass(frozen=True, slots=True) +class ContextWindowOverflow: + max_input_tokens: int + tokens: int + certain: bool + + +def stand_in_tokenizer() -> SelectTokenizerResponse: + return {"type": "openai_tokenizer", "tokenizer": get_encoding(STAND_IN_ENCODING)} + + +def context_window_overflow( + window: ContextWindow | None, token_counts: Mapping[str, PreCallTokenCount] | None +) -> ContextWindowOverflow | None: + """How a prompt overflows a deployment's window, or None when it fits or was not counted. A stand-in count is + only an estimate, so it proves an overflow only past STAND_IN_OVERCOUNT_MARGIN times the window.""" + count: Final = token_counts.get(window.tokenizer_model) if window is not None and token_counts is not None else None + if window is None or count is None or count.tokens <= window.max_input_tokens: + return None + return ContextWindowOverflow( + max_input_tokens=window.max_input_tokens, + tokens=count.tokens, + certain=count.own_tokenizer or count.tokens > window.max_input_tokens * STAND_IN_OVERCOUNT_MARGIN, + ) diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index b02e0732817..60e3bdbad76 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -120,6 +120,16 @@ def test_token_counter_default_encoding_matches_cl100k(): assert token_counter_new(model=None, text="hello world") == expected +def test_token_counter_counts_with_a_custom_openai_tokenizer_encoding(): + japanese: Final = "日本語のテキストを数えます。" * 50 + o200k: Final = tiktoken.get_encoding("o200k_base") + assert len(o200k.encode(japanese)) != len(tiktoken.get_encoding("cl100k_base").encode(japanese)) + + assert token_counter_new( + custom_tokenizer={"type": "openai_tokenizer", "tokenizer": o200k}, text=japanese + ) == len(o200k.encode(japanese)) + + def test_token_counter_text_over_chunk_boundary_stays_close_to_tiktoken(): text = ("The quick brown fox jumps over the lazy dog. " * 30)[:1025] encoding = tiktoken.get_encoding("cl100k_base") diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 300c9cff6ff..f356b47c2ae 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -36,6 +36,12 @@ from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError, ModifyResponseException from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.router_utils.context_window_fit import ( + STAND_IN_OVERCOUNT_MARGIN, + ContextWindow, + PreCallTokenCount, + stand_in_tokenizer, +) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.plumbing.context import set_request_root_span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -7916,7 +7922,7 @@ def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 5} ) calls = [] @@ -7950,7 +7956,7 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 5} ) calls = [] @@ -7966,7 +7972,7 @@ def test_pre_call_checks_uses_precounted_tokens(monkeypatch): model="m", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], - input_token_count=1000, + input_token_counts={"gpt-3.5-turbo": PreCallTokenCount(tokens=1000, own_tokenizer=True)}, ) assert calls == [] @@ -7984,7 +7990,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 1_000_000} ) counting_threads = [] @@ -7998,7 +8004,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo original_pre_call_checks = router._pre_call_checks def spy(**kwargs): - counts_passed_in.append(kwargs.get("input_token_count")) + counts_passed_in.append(kwargs.get("input_token_counts")) return original_pre_call_checks(**kwargs) monkeypatch.setattr(router, "_pre_call_checks", spy) @@ -8013,7 +8019,7 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo ) assert len(result) == 1 - assert counts_passed_in == [42] + assert counts_passed_in == [{"gpt-3.5-turbo": PreCallTokenCount(tokens=42, own_tokenizer=True)}] assert len(counting_threads) == 1 assert counting_threads[0] is not threading.current_thread() @@ -8021,12 +8027,12 @@ async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(mo @pytest.mark.parametrize( "model_info,expected", [ - ({"max_input_tokens": 100}, True), - ({"max_input_tokens": None}, False), - ({}, False), + ({"key": "gpt-3.5-turbo", "max_input_tokens": 100}, frozenset({"gpt-3.5-turbo"})), + ({"key": "gpt-3.5-turbo", "max_input_tokens": None}, frozenset()), + ({"key": "gpt-3.5-turbo"}, frozenset()), ], ) -def test_pre_call_checks_need_token_count(monkeypatch, model_info, expected): +def test_pre_call_check_tokenizer_models(monkeypatch, model_info, expected): """Only a deployment that declares an integer context window makes a token count worth taking.""" router = litellm.Router( model_list=[ @@ -8039,7 +8045,7 @@ def test_pre_call_checks_need_token_count(monkeypatch, model_info, expected): deployments = [ {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, ] - assert router._pre_call_checks_need_token_count("m", deployments) is expected + assert router._pre_call_check_tokenizer_models("m", deployments) == expected def test_deployment_max_input_tokens_survives_an_unmappable_deployment(monkeypatch): @@ -8058,16 +8064,18 @@ def test_deployment_max_input_tokens_survives_an_unmappable_deployment(monkeypat def flaky_model_info(deployment, received_model_name, id=None): if deployment["model_info"]["id"] == "unmapped": raise ValueError("This model isn't mapped yet.") - return {"max_input_tokens": 100} + return {"key": "gpt-3.5-turbo", "max_input_tokens": 100} monkeypatch.setattr(router, "get_router_model_info", flaky_model_info) unmapped = {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "unmapped"}} mapped = {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "mapped"}} - assert router._deployment_max_input_tokens("m", unmapped) is None - assert router._deployment_max_input_tokens("m", mapped) == 100 - assert router._pre_call_checks_need_token_count("m", [unmapped, mapped]) is True + assert router._deployment_context_window("m", unmapped) is None + assert router._deployment_context_window("m", mapped) == ContextWindow( + tokenizer_model="gpt-3.5-turbo", max_input_tokens=100 + ) + assert router._pre_call_check_tokenizer_models("m", [unmapped, mapped]) == frozenset({"gpt-3.5-turbo"}) def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monkeypatch): @@ -8082,7 +8090,7 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 5} ) calls = [] @@ -8097,7 +8105,7 @@ def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monke model="m", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], - input_token_count=None, + input_token_counts=None, skip_inline_token_count=True, ) @@ -8117,7 +8125,7 @@ async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypa enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 5} ) counting_threads = [] @@ -8154,7 +8162,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 5} ) deployments = [ @@ -8176,7 +8184,7 @@ async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypat ran.append("count") await task - assert count is not None and count > 0 + assert count is not None and count["gpt-3.5-turbo"].tokens > 0 assert ran == ["competitor", "count"] @@ -8222,7 +8230,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 1} ) deployments = [ @@ -8249,7 +8257,7 @@ def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 1} ) deployments = [ @@ -8293,7 +8301,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): assert with_instructions_tokens > input_only_tokens monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": input_only_tokens} ) with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( @@ -8364,7 +8372,7 @@ def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwarg messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": prompt_only_tokens} ) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 @@ -8405,7 +8413,7 @@ def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): messages = [{"role": "user", "content": "hi"}] messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) - monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) + monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": messages_only_tokens}) assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 with pytest.raises(litellm.ContextWindowExceededError): @@ -8524,7 +8532,7 @@ def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): enable_pre_call_checks=True, ) monkeypatch.setattr( - router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} + router, "get_router_model_info", lambda **kwargs: {"key": "gpt-3.5-turbo", "max_input_tokens": 5} ) counted: list[dict] = [] @@ -30143,6 +30151,84 @@ def test_pre_call_checks_filters_deployment_when_token_count_exceeds_limit(): assert [d["model_info"]["id"] for d in filtered] == ["large"] +_JAPANESE_PROMPT: Final = [{"role": "user", "content": "日本語のテキストを数えます。" * 50}] +_RUSSIAN_PROMPT: Final = [{"role": "user", "content": "Привет, как дела? Это тестовый текст. " * 50}] + + +def _stand_in_tokens(messages: list[dict[str, str]]) -> int: + return litellm.token_counter(custom_tokenizer=stand_in_tokenizer(), messages=messages) + + +def _single_deployment(model: str, max_input_tokens: int) -> list[dict[str, object]]: + return [ + { + "model_name": "m", + "litellm_params": {"model": model, "api_key": "my-fake-key"}, + "model_info": {"id": "only", "max_input_tokens": max_input_tokens}, + } + ] + + +def _kept_ids(model_list: list[dict[str, object]], messages: list[dict[str, str]]) -> list[str]: + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + return [ + d["model_info"]["id"] + for d in router._pre_call_checks(model="m", healthy_deployments=model_list, messages=messages) + ] + + +@pytest.mark.parametrize("model", ["openai/my-fake-model", "claude-sonnet-5-5"]) +def test_pre_call_checks_leaves_a_stand_in_overcount_within_the_margin_to_the_provider(model): + tokens: Final = _stand_in_tokens(_RUSSIAN_PROMPT) + smallest_window_within_margin: Final = -(-tokens // STAND_IN_OVERCOUNT_MARGIN) + + assert _kept_ids(_single_deployment(model, tokens - 1), _RUSSIAN_PROMPT) == ["only"] + assert _kept_ids(_single_deployment(model, smallest_window_within_margin), _RUSSIAN_PROMPT) == ["only"] + + +def test_pre_call_checks_refuses_a_stand_in_count_past_the_margin(): + tokens: Final = _stand_in_tokens(_RUSSIAN_PROMPT) + largest_window_past_margin: Final = (tokens - 1) // STAND_IN_OVERCOUNT_MARGIN + + with pytest.raises(litellm.ContextWindowExceededError): + _kept_ids(_single_deployment("openai/my-fake-model", largest_window_past_margin), _RUSSIAN_PROMPT) + + +def test_pre_call_checks_counts_an_openai_model_with_its_own_encoding(): + own_tokens: Final = litellm.token_counter(model="gpt-5.6-luna", messages=_JAPANESE_PROMPT) + assert own_tokens < litellm.token_counter(messages=_JAPANESE_PROMPT) + + assert _kept_ids(_single_deployment("gpt-5.6-luna", own_tokens), _JAPANESE_PROMPT) == ["only"] + with pytest.raises(litellm.ContextWindowExceededError): + _kept_ids(_single_deployment("gpt-5.6-luna", own_tokens - 1), _JAPANESE_PROMPT) + + +def test_pre_call_checks_counts_a_provider_prefixed_openai_model_with_its_own_encoding(): + own_tokens: Final = litellm.token_counter(model="gpt-5.6-luna", messages=_JAPANESE_PROMPT) + + with pytest.raises(litellm.ContextWindowExceededError): + _kept_ids(_single_deployment("openai/gpt-5.6-luna", own_tokens - 1), _JAPANESE_PROMPT) + assert _kept_ids(_single_deployment("openai/gpt-5.6-luna", own_tokens), _JAPANESE_PROMPT) == ["only"] + + +def test_pre_call_checks_steers_off_a_stand_in_overcount_when_another_deployment_fits(): + tokens: Final = _stand_in_tokens(_RUSSIAN_PROMPT) + model_list: Final = [ + { + "model_name": "m", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "small", "max_input_tokens": tokens - 1}, + }, + { + "model_name": "m", + "litellm_params": {"model": "openai/my-other-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "large", "max_input_tokens": tokens}, + }, + ] + + assert _kept_ids(model_list, _RUSSIAN_PROMPT) == ["large"] + + def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails(): model_list: Final = [ { diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 20b0055f729..9d1a16dffb4 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -20,6 +20,7 @@ import sys import threading from pathlib import Path from types import MappingProxyType +from collections.abc import Mapping from typing import Final, cast from unittest.mock import Mock, patch @@ -35,6 +36,7 @@ from litellm.litellm_core_utils.ptu_pricing import ptu_config_error from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS +from litellm.router_utils.context_window_fit import PreCallTokenCount from litellm.types.router import ( Deployment, DeploymentTypedDict, @@ -383,9 +385,14 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( assert listing is not None assert listing.max_input_tokens == 8192 assert router.get_configured_token_limits("local") == (8192, 8192) - assert router._deployment_max_input_tokens("local", router.model_list[1]) == 4096 + window: Final = router._deployment_context_window("local", router.model_list[1]) + assert window is not None + assert window.max_input_tokens == 4096 allowed: Final = router._pre_call_checks( - model="local", healthy_deployments=router.model_list, input="prompt", input_token_count=5000 + model="local", + healthy_deployments=router.model_list, + input="prompt", + input_token_counts={window.tokenizer_model: PreCallTokenCount(tokens=5000, own_tokenizer=True)}, ) assert [deployment["model_info"]["id"] for deployment in allowed] == ["one"] assert router.model_list[0]["model_info"].get("max_input_tokens") is None @@ -4162,6 +4169,13 @@ def _capacity_test_router(route: str, base: str, base_location: str | None, expl ) +def _tokens_for_every_model(router: Router, tokens: int) -> Mapping[str, PreCallTokenCount]: + return { + model: PreCallTokenCount(tokens=tokens, own_tokenizer=True) + for model in router._pre_call_check_tokenizer_models("capacity-test", router.model_list) + } + + def _deployments_allowed_for_2048_input_tokens(router: Router) -> int: try: return len( @@ -4169,7 +4183,7 @@ def _deployments_allowed_for_2048_input_tokens(router: Router) -> int: model="capacity-test", healthy_deployments=router.model_list, input="synthetic input", - input_token_count=2048, + input_token_counts=_tokens_for_every_model(router, 2048), ) ) except litellm.ContextWindowExceededError: