scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(router): count pre-call context-window tokens with each deployment's tokenizer

Browse source 
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.
This commit is contained in:
mateo-berri 2026-10-29 03:05:57 -07:00
parent bad98c4b02
commit 020ded24fe
7 changed files with 374 additions and 109 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/litellm_core_utils/token_counter.py
View file

@ -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(

202
litellm/router.py
View file

@ -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,

15
litellm/router_strategy/complexity_router/context_compaction.py
View file

@ -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], ...]]:

47
litellm/router_utils/context_window_fit.py Normal file
View file

@ -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,
)

10
tests/unit/litellm_core_utils/test_token_counter.py
View file

@ -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")

138
tests/unit/test_router/test_router.py
View file

@ -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 = [
{

20
tests/unit/test_router_model_cost_isolation.py
View file

@ -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: