From 49e4a357ba092b8740b2dc6ad9d8fbccf12fd854 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 2 Nov 2026 20:50:59 -0800 Subject: [PATCH] fix(router): count each distinct tokenizer once and steer only to an untried context window fallback Count a prompt once per distinct encoding and message framing instead of once per OpenAI model name, so an OpenAI-only group with several models runs one count. A context window fallback now keeps an uncertain overflow off a deployment only when its chain still holds an untried target other than the failing group. Complexity router admission counts through the router's bounded async counter. --- litellm/litellm_core_utils/token_counter.py | 15 +++ litellm/router_api.py | 16 +++ litellm/router_python_core.py | 18 ++- .../complexity_router/context_compaction.py | 27 +++-- litellm/router_utils/context_window_fit.py | 20 +++- .../router_utils/fallback_event_handlers.py | 4 +- .../test_context_compaction.py | 42 +++++++ tests/unit/test_router/test_router.py | 111 ++++++++++++++++++ 8 files changed, 238 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 78105ad3fab..3991f9f928a 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -716,6 +716,21 @@ class _MessageCountParams: return pdf_page_pricing(self._model) +@dataclass(frozen=True, slots=True) +class MessageCountFraming: + legacy_message_accounting: bool + tool_section: ChatTemplateToolSection | None + pdf_pricing: PdfPagePricing + + +def message_count_framing(model: str) -> MessageCountFraming: + return MessageCountFraming( + legacy_message_accounting=uses_legacy_message_accounting(model), + tool_section=chat_template_tool_section(model), + pdf_pricing=pdf_page_pricing(model), + ) + + def _is_openai_chat_model(model: str) -> bool: return model in litellm.open_ai_chat_completion_models or model in litellm.azure_llms diff --git a/litellm/router_api.py b/litellm/router_api.py index ad854d6a203..03ca1776142 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1518,6 +1518,22 @@ class RoutingContext: ) -> int: return self._router._count_pre_call_check_tokens(messages, input, request_kwargs) + async def acount_pre_call_check_tokens( + self, + model: str, + healthy_deployments: Sequence[Mapping[str, object]], + messages: Sequence[Mapping[str, object]] | None, + input: str | Sequence[object] | None, + request_kwargs: Mapping[str, object] | None, + ) -> Mapping[str | None, int] | None: + return await self._router._acount_pre_call_check_tokens( + model=model, + healthy_deployments=healthy_deployments, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ) + def update_kwargs_before_fallbacks( self, model: str, diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 59f6eeb236f..c708db99a44 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -295,6 +295,7 @@ from litellm.router_utils.context_window_fit import ( ContextWindowOverflow, context_window_overflow, pre_call_tokenizer, + token_count_key, ) from litellm.router_utils.cooldown_cache import CooldownCache from litellm.router_utils.cooldown_handlers import ( @@ -14886,8 +14887,8 @@ class PythonCore(DeprecatedRouterMembers): def _pre_call_check_tokenizer_models( self, model: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> frozenset[str | None]: - """The counts the prompt needs: one per OpenAI model and one stand-in count shared by every other deployment - that declares a context window. + """The models the prompt is counted for: each OpenAI model, and None for the stand-in shared by every other + 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. @@ -14902,24 +14903,29 @@ class PythonCore(DeprecatedRouterMembers): input: str | Sequence[object] | None, request_kwargs: Mapping[str, object] | None, ) -> Mapping[str | None, int]: - return MappingProxyType( + keys: Final = MappingProxyType( + {openai_model: token_count_key(openai_model) for openai_model in tokenizer_models} + ) + counted_models: Final = MappingProxyType({key: openai_model for openai_model, key in keys.items()}) + counts: Final = MappingProxyType( { - openai_model: self._count_pre_call_check_tokens( + key: self._count_pre_call_check_tokens( messages=messages, input=input, request_kwargs=request_kwargs, model=openai_model or "", custom_tokenizer=pre_call_tokenizer(openai_model), ) - for openai_model in tokenizer_models + for key, openai_model in counted_models.items() } ) + return MappingProxyType({openai_model: counts[key] for openai_model, key in keys.items()}) async def _acount_pre_call_check_tokens( self, model: str, healthy_deployments: Sequence[Mapping[str, object]], - messages: Sequence[Mapping[str, str]] | None, + messages: Sequence[Mapping[str, object]] | None, input: str | Sequence[object] | None, request_kwargs: Mapping[str, object] | None, ) -> Mapping[str | None, int] | None: diff --git a/litellm/router_strategy/complexity_router/context_compaction.py b/litellm/router_strategy/complexity_router/context_compaction.py index 04535bc7c68..c0376d1be40 100644 --- a/litellm/router_strategy/complexity_router/context_compaction.py +++ b/litellm/router_strategy/complexity_router/context_compaction.py @@ -384,13 +384,26 @@ async def _check_client_managed_admission( ) -> None: if not ctx.settings().enable_pre_call_checks: return - await asyncio.to_thread( - ctx.pre_call_checks, - model=str(payload["model"]), - healthy_deployments=list[object](_ITEMS.validate_python((deployment,))), - 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), + model: Final = str(payload["model"]) + healthy_deployments: Final = _ITEMS.validate_python((deployment,)) + messages: Final = _ITEMS.validate_python(payload["messages"]) if "messages" in payload else None + request_input: Final = _INPUT.validate_python(payload.get("input")) + request_kwargs: Final = _DICT.validate_python(payload) + token_counts: Final = await ctx.acount_pre_call_check_tokens( + model=model, + healthy_deployments=healthy_deployments, + messages=messages, + input=request_input, + request_kwargs=request_kwargs, + ) + ctx.pre_call_checks( + model=model, + healthy_deployments=list[object](healthy_deployments), + messages=messages, # pyright: ignore[reportArgumentType] # legacy annotation omits structured content + input=request_input, + request_kwargs=request_kwargs, + input_token_counts=token_counts, + skip_inline_token_count=True, ) diff --git a/litellm/router_utils/context_window_fit.py b/litellm/router_utils/context_window_fit.py index 5b589f6e07c..5d73bb43473 100644 --- a/litellm/router_utils/context_window_fit.py +++ b/litellm/router_utils/context_window_fit.py @@ -2,7 +2,12 @@ from collections.abc import Mapping from dataclasses import dataclass from typing import Final -from litellm.litellm_core_utils.token_counter import openai_tokenizer_encoding +from litellm.litellm_core_utils.token_counter import ( + MessageCountFraming, + message_count_framing, + openai_tokenizer_encoding, + openai_tokenizer_encoding_name, +) from litellm.rust_bridge.tokenizer import get_encoding from litellm.types.utils import SelectTokenizerResponse @@ -23,6 +28,12 @@ class ContextWindowOverflow: certain: bool +@dataclass(frozen=True, slots=True) +class TokenCountKey: + encoding: str + framing: MessageCountFraming + + def stand_in_tokenizer() -> SelectTokenizerResponse: return {"type": "openai_tokenizer", "tokenizer": get_encoding(STAND_IN_ENCODING)} @@ -33,6 +44,13 @@ def pre_call_tokenizer(openai_model: str | None) -> SelectTokenizerResponse: return {"type": "openai_tokenizer", "tokenizer": openai_tokenizer_encoding(openai_model)} +def token_count_key(openai_model: str | None) -> TokenCountKey: + return TokenCountKey( + encoding=STAND_IN_ENCODING if openai_model is None else openai_tokenizer_encoding_name(openai_model), + framing=message_count_framing(openai_model or ""), + ) + + def context_window_overflow( window: ContextWindow | None, token_counts: Mapping[str | None, int] | None ) -> ContextWindowOverflow | None: diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 9f510886938..35606a0f03c 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -621,7 +621,9 @@ def context_window_fallback_resolves() -> bool: chain, _ = get_fallback_model_group_for_lookup_groups( fallbacks=scope.fallbacks, lookup_groups=fallback_lookup_groups(scope.kwargs, scope.model_group) ) - return chain is not None + failed_group: Final = get_pre_routing_selection(scope.kwargs) or scope.model_group + remaining: Final = tuple(target for target in chain or () if fallback_attempt_key(target) != failed_group) + return bool(remaining) and has_unattempted_fallback_target(remaining, scope.kwargs) def _resolved_a_specific_chain( diff --git a/tests/unit/router_strategy/complexity_router/test_context_compaction.py b/tests/unit/router_strategy/complexity_router/test_context_compaction.py index 850accf574b..c58b05f7764 100644 --- a/tests/unit/router_strategy/complexity_router/test_context_compaction.py +++ b/tests/unit/router_strategy/complexity_router/test_context_compaction.py @@ -1,5 +1,6 @@ import asyncio import json +import threading from collections.abc import Iterator, Mapping from copy import deepcopy from functools import partial @@ -10,6 +11,8 @@ import pytest import respx import litellm +from litellm.constants import TOKEN_COUNTER_MAX_CONCURRENT_COUNTS +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.llms import compaction as native from litellm.router_api import RoutingContext from litellm.router_strategy.complexity_router.config import ContextCompactionConfig @@ -20,6 +23,7 @@ from litellm.router_strategy.complexity_router.context_compaction import ( compact_to_fit, compaction_executor, ) +from litellm.router_utils.context_window_fit import STAND_IN_OVERCOUNT_MARGIN, stand_in_tokenizer from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES from litellm.types.router import Deployment @@ -414,6 +418,44 @@ async def test_client_managed_overflow_keeps_context_window_admission( assert answer.call_count == 0 +async def test_client_managed_admission_counts_only_when_a_token_count_slot_is_free( + wire: tuple[respx.Route, respx.Route], +) -> None: + router: Final = make_router() + deployment: Final = router.get_deployment(model_id="pinned-answer") + assert deployment is not None + request: Final = {**history("chat"), "model": "small", "max_tokens": 64, "_context_compaction_state": CompactionState()} + await arm_compaction(request, ContextCompactionConfig(), ("large",)) + managed: Final = {**request, "context_management": {"edits": []}} + stand_in_count: Final = litellm.token_counter(custom_tokenizer=stand_in_tokenizer(), messages=request["messages"]) + assert stand_in_count > 512 * STAND_IN_OVERCOUNT_MARGIN + loop: Final = asyncio.get_running_loop() + holding: Final = asyncio.Queue[None]() + release: Final = threading.Event() + + def hold_count_slot() -> bool: + loop.call_soon_threadsafe(holding.put_nowait, None) + return release.wait(timeout=30) + + holders: Final = tuple( + asyncio.create_task(offload_token_count(hold_count_slot)()) for _ in range(TOKEN_COUNTER_MAX_CONCURRENT_COUNTS) + ) + try: + for _ in holders: + await holding.get() + admission: Final = asyncio.create_task( + compact_to_fit(RoutingContext(router), deployment.model_dump(), managed, "chat") + ) + finished_while_slots_held, _ = await asyncio.wait((admission,), timeout=0.2) + finally: + release.set() + assert await asyncio.gather(*holders) == [True] * TOKEN_COUNTER_MAX_CONCURRENT_COUNTS + assert not finished_while_slots_held + with pytest.raises(litellm.ContextWindowExceededError): + await admission + assert all(route.call_count == 0 for route in wire) + + @pytest.mark.parametrize("case", ["conflicting_defaults", "small_window", "capability_false", "capability_missing"]) async def test_automatic_compactor_skips_conflicts_and_requires_capacity_and_capability( wire: tuple[respx.Route, respx.Route], monkeypatch: pytest.MonkeyPatch, case: str, diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 692a8078a23..d9373b1d48f 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1,4 +1,5 @@ import ast, asyncio, importlib, time +import base64 import contextlib import copy import functools @@ -23,6 +24,7 @@ import httpx import openai import pytest import respx +from pypdf import PdfWriter from fastapi import HTTPException from opentelemetry import trace from pydantic import BaseModel, TypeAdapter @@ -37,6 +39,10 @@ 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.litellm_core_utils.token_counter import openai_tokenizer_encoding_name +from litellm.router import router_arguments +from litellm.router_config import build_router_config +from litellm.router_python_host import Endpoint, PythonHost from litellm.router_utils.context_window_fit import ( STAND_IN_OVERCOUNT_MARGIN, ContextWindow, @@ -30572,6 +30578,8 @@ def _context_window_fallback_router(**router_settings: object) -> Router: ({"context_window_fallbacks": [{"m": ["fallback"]}]}, {"disable_fallbacks": True}, "m"), ({"context_window_fallbacks": [{"m": ["fallback"]}], "max_fallbacks": 0}, {}, "m"), ({"context_window_fallbacks": [{"other": ["fallback"]}]}, {}, "m"), + ({"context_window_fallbacks": [{"m": []}]}, {}, "m"), + ({"context_window_fallbacks": [{"m": ["m"]}]}, {}, "m"), ({"fallbacks": [{"m": ["fallback"]}]}, {}, "m"), ({}, {}, "m"), ], @@ -30582,6 +30590,8 @@ def _context_window_fallback_router(**router_settings: object) -> Router: "fallbacks_disabled", "no_fallback_budget", "fallback_for_another_group", + "empty_fallback_chain", + "fallback_chain_names_only_the_failing_group", "generic_fallbacks_only", "no_fallback", ], @@ -30604,6 +30614,33 @@ def test_a_sync_completion_sends_an_uncertain_stand_in_overflow_to_its_context_w assert response.choices[0].message.content == "fallback" +@pytest.mark.parametrize( + "context_window_fallbacks", + [[{"a": ["b"]}], [{"a": ["b"]}, {"b": ["a"]}]], + ids=["fallback_target_has_no_chain_of_its_own", "fallback_target_falls_back_to_a_tried_group"], +) +async def test_a_context_window_fallback_target_keeps_its_own_uncertain_stand_in_overflow( + context_window_fallbacks: list[dict[str, list[str]]], +): + router: Final = Router( + model_list=[ + { + "model_name": group, + "litellm_params": {"model": model, "api_key": "my-fake-key", "mock_response": group}, + "model_info": {"id": group, "max_input_tokens": _uncertain_overflow_window()}, + } + for group, model in (("a", "openai/my-fake-model"), ("b", "openai/my-other-fake-model")) + ], + enable_pre_call_checks=True, + num_retries=0, + context_window_fallbacks=context_window_fallbacks, + ) + + response: Final = await router.acompletion(model="a", messages=_RUSSIAN_PROMPT) + + assert response.choices[0].message.content == "b" + + @pytest.mark.parametrize(("allowed_model_region", "kept"), [("eu", ["eu"]), (["eu"], [])]) def test_pre_call_checks_excludes_every_deployment_whose_region_cannot_match(allowed_model_region: object, kept): model_list: Final = [ @@ -30735,6 +30772,80 @@ async def test_acount_pre_call_check_tokens_counts_every_stand_in_deployment_wit } +class _TokenCounterCallsHost(PythonHost): + __slots__ = ("counted_models",) + + def __init__(self) -> None: + self.counted_models: Final[list[object]] = [] + + def invoke(self, op: str) -> Endpoint: + endpoint: Final = super().invoke(op) + if op != "token_counter": + return endpoint + + def counting_endpoint(**kwargs: object) -> object: + self.counted_models.append(kwargs.get("model")) + return endpoint(**kwargs) + + return counting_endpoint + + +def _counted_router(models: tuple[str, ...], host: PythonHost) -> tuple[Router, list[dict[str, object]]]: + model_list: Final = [ + { + "model_name": "m", + "litellm_params": {"model": model, "api_key": "my-fake-key"}, + "model_info": {"id": model, "max_input_tokens": 1_000_000}, + } + for model in models + ] + router_config: Final = build_router_config( + router_arguments(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + ) + return Router.from_config(router_config, host), model_list + + +async def test_acount_pre_call_check_tokens_counts_openai_models_that_share_an_encoding_once(): + models: Final = ("gpt-4o", "gpt-4.1", "gpt-4o-mini", "gpt-5", "gpt-5-mini") + host: Final = _TokenCounterCallsHost() + router, model_list = _counted_router(models, host) + + counts: Final = await router._acount_pre_call_check_tokens( + model="m", healthy_deployments=model_list, messages=_JAPANESE_PROMPT, input=None, request_kwargs=None + ) + + assert len(host.counted_models) == len({openai_tokenizer_encoding_name(model) for model in models}) < len(models) + assert counts == {model: litellm.token_counter(model=model, messages=_JAPANESE_PROMPT) for model in models} + + +def _blank_pdf_prompt() -> list[dict[str, object]]: + writer: Final = PdfWriter() + writer.add_blank_page(width=612, height=792) + pdf: Final = io.BytesIO() + writer.write(pdf) + file_data: Final = "data:application/pdf;base64," + base64.b64encode(pdf.getvalue()).decode() + return [{"role": "user", "content": [{"type": "file", "file": {"filename": "a.pdf", "file_data": file_data}}]}] + + +async def test_acount_pre_call_check_tokens_gives_each_openai_model_its_own_pdf_framing(): + prompt: Final = _blank_pdf_prompt() + listed_model, unlisted_model = "gpt-4o", "openai/gpt-9.9-preview" + assert openai_tokenizer_encoding_name(listed_model) == openai_tokenizer_encoding_name(unlisted_model) + host: Final = _TokenCounterCallsHost() + router, model_list = _counted_router((listed_model, unlisted_model, "cohere_chat/command-a-plus-05-2026"), host) + + counts: Final = await router._acount_pre_call_check_tokens( + model="m", healthy_deployments=model_list, messages=prompt, input=None, request_kwargs=None + ) + + expected: Final = { + None: litellm.token_counter(custom_tokenizer=stand_in_tokenizer(), messages=prompt), + **{model: litellm.token_counter(model=model, messages=prompt) for model in (listed_model, unlisted_model)}, + } + assert expected[listed_model] != expected[unlisted_model] + assert counts == expected + + def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails(): model_list: Final = [ {