diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index cc592bbcbc5..8442d7bca3c 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -24,6 +24,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from itertools import accumulate, chain, islice, takewhile +from math import ceil from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast @@ -496,6 +497,11 @@ _CODEX_REMINDER_MARKERS: Final = _DEFAULT_REMINDER_MARKERS + ( _TRUNCATION_MARKER: Final = "..." _TRUNCATION_HEAD_FRACTION: Final = 0.3 _MIN_QUOTED_TURN_CHARS: Final = 120 +# Reserve 25% for provider tokenizer differences and request framing. +_JEV_DEFAULT_INPUT_TOKENS: Final = 32768 +_JEV_INPUT_WINDOW_FRACTION: Final = 0.75 +# Bound tokenizer work and the provider's separate character-based size check. +_JEV_PRECUT_CHARS_PER_TOKEN: Final = 10 _CLASSIFIER_CIRCUIT_OPEN_SIGNAL: Final = "classifier-circuit-open" @@ -2312,14 +2318,22 @@ class ComplexityRouter(CustomLogger): ) ) timeout_s: Final = config.timeout_ms / 1000 - request: Final = build_jev_request( - prompt=self._classifier_context_payload(prompt, system_prompt, request_kwargs, messages), - system_prompt=None, - model=config.model, - instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, - criteria=criteria, - ) + instructions: Final = config.instructions or DEFAULT_JEV_INSTRUCTIONS try: + request: Final = build_jev_request( + prompt=await self._jev_state_within_budget( + prompt, + system_prompt, + request_kwargs, + messages, + question_text="\n".join((instructions, *(f"{name}: {text}" for name, text in criteria.items()))), + budget_tokens=self._jev_input_token_budget(config.deployment_name), + ), + system_prompt=None, + model=config.model, + instructions=instructions, + criteria=criteria, + ) response: Final = await asyncio.wait_for(client.evaluate(request, timeout_s, request_kwargs), timeout_s) answer: Final = response.answers.get("tier") if answer is None: @@ -2551,6 +2565,68 @@ class ComplexityRouter(CustomLogger): tier=tier, score=None, signals=("classifier-failed:default-model",), cause="default_model_fallback" ) + def _jev_input_token_budget(self, deployment_name: str | None) -> int: + """Use the smallest deployment window, defaulting unknown members to 32k.""" + if deployment_name is None: + return int(_JEV_DEFAULT_INPUT_TOKENS * _JEV_INPUT_WINDOW_FRACTION) + window, any_undeclared = self._group_window_facts(deployment_name) + if window is None: + return int(_JEV_DEFAULT_INPUT_TOKENS * _JEV_INPUT_WINDOW_FRACTION) + if any_undeclared: + return int(min(window, _JEV_DEFAULT_INPUT_TOKENS) * _JEV_INPUT_WINDOW_FRACTION) + return int(window * _JEV_INPUT_WINDOW_FRACTION) + + async def _jev_state_within_budget( + self, + prompt: str, + system_prompt: str | None, + request_kwargs: Mapping[str, object] | None, + messages: Sequence[Mapping[str, object]] | None, + question_text: str, + budget_tokens: int, + ) -> str: + """Fit state plus question, trimming caller constraints before the current ask. + + Keep both ends of each text and at least 120 characters. Recount after each + cut: token density can change. If fixed context cannot fit, let the existing + classifier error handler choose the configured fallback without an API call. + """ + import litellm + from litellm.litellm_core_utils.token_counter import offload_token_count + + quoted_system: Final = self._classifier_caller_constraints(system_prompt, request_kwargs) + precut_chars: Final = budget_tokens * _JEV_PRECUT_CHARS_PER_TOKEN + prompt_chars = min(len(prompt), precut_chars) # rebind-ok: each iteration reduces the current ask's cap + system_chars = min(len(quoted_system or ""), precut_chars) # rebind-ok: reduced after each cut + while True: + payload = self._classifier_context_payload( + prompt, + system_prompt, + request_kwargs, + messages, + prompt_chars=prompt_chars, + system_prompt_chars=system_chars, + ) + measured = f"{payload}\n{question_text}" + byte_bound = len(measured.encode()) + if byte_bound <= budget_tokens: + return payload + try: + tokens = await offload_token_count(litellm.token_counter)(text=measured) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # token_counter's custom_tokenizer is a bare dict + except Exception as exc: # noqa: BLE001 # use a conservative bound when the tokenizer fails + verbose_router_logger.debug("ComplexityRouter: JEV token count failed (%s)", type(exc).__name__) + tokens = byte_bound + if tokens <= budget_tokens: + return payload + overflow_fraction = (tokens - budget_tokens) / tokens + chars_to_remove = ceil(len(measured) * overflow_fraction) + 2 * len(_TRUNCATION_MARKER) + system_cut = min(chars_to_remove, max(system_chars - _MIN_QUOTED_TURN_CHARS, 0)) + prompt_cut = min(chars_to_remove - system_cut, max(prompt_chars - _MIN_QUOTED_TURN_CHARS, 0)) + if system_cut == 0 and prompt_cut == 0: + raise ValueError("JEV classifier instructions and fixed context exceed the input budget") + system_chars -= system_cut + prompt_chars -= prompt_cut + def _classifier_caller_constraints( self, system_prompt: str | None, request_kwargs: Mapping[str, object] | None ) -> str | None: @@ -2573,7 +2649,14 @@ class ComplexityRouter(CustomLogger): messages: Sequence[Mapping[str, object]] | None, *, encrypted_task: bool = False, + prompt_chars: int | None = None, + system_prompt_chars: int | None = None, ) -> str: + """The classifier's user message for this request. + + prompt_chars and system_prompt_chars cap the current ask and the quoted system text after the + prior turns are chosen, since a turn repeating the ask is excluded by matching its full text. + """ include_assistant: Final = self.config.classifier_context_include_assistant_turns marker_pairs: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) context_enabled: Final = bool(messages) and self.config.classifier_context_window_size > 0 @@ -2595,9 +2678,16 @@ class ComplexityRouter(CustomLogger): and len(tuple(islice(_iter_context_turns_newest_first(messages or (), include_assistant, marker_pairs), 2))) > 1 ) + caller_constraints: Final = self._classifier_caller_constraints(system_prompt, request_kwargs) + quoted_prompt: Final = _truncate(prompt, prompt_chars) if prompt_chars is not None else prompt + quoted_system: Final = ( + _truncate(caller_constraints, system_prompt_chars) + if caller_constraints is not None and system_prompt_chars is not None + else caller_constraints + ) return self._build_classifier_user_payload( - prompt="The delegated task in the following agent_message." if encrypted_task else prompt, - system_prompt=self._classifier_caller_constraints(system_prompt, request_kwargs), + prompt="The delegated task in the following agent_message." if encrypted_task else quoted_prompt, + system_prompt=quoted_system, prior_turns=prior_turns, messages=messages, has_prior_conversation=has_prior_conversation, diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 10a8f7929d9..830666fe4be 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -22,6 +22,7 @@ from litellm.router_strategy.complexity_router.jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevChoiceAnswer, + JevSystemOneRequest, JevSystemOneResponse, JevUsage, build_jev_request, @@ -976,3 +977,170 @@ def test_switching_to_deployed_classifier_does_not_copy_legacy_credentials() -> ) assert write.effective is not None assert write.effective["opensource_classifier_config"] == {"deployment_name": "registered"} + + +class _StateRecorder: + def __init__(self) -> None: + self.states: tuple[str, ...] = () + + async def evaluate( + self, request: JevSystemOneRequest, timeout_s: float, request_kwargs: Mapping[str, object] | None = None + ) -> JevSystemOneResponse: + self.states = (*self.states, request.state) + return JevSystemOneResponse(answers={"tier": _answer()}) + + +_DEFAULT_JEV_STATE_BUDGET: Final = 24576 + + +def _words(label: str, count: int) -> str: + return " ".join(f"{label}{i % 997}" for i in range(count)) + + +def _budget_router( + recorder: _StateRecorder, + model_list: list[dict[str, object]] | None = None, # mutable-ok: litellm.Router takes a list + jev_config: Mapping[str, object] | None = None, +) -> ComplexityRouter: + return ComplexityRouter( + "jev-budget", + litellm.Router(model_list=model_list or []), + {"classifier_type": "jev", "jev_classifier_config": dict(jev_config or {}), "tiers": {"SIMPLE": "cheap"}}, + jev_client=recorder, + derive_savings_baseline=False, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("system_prompt", ["caller constraints", f"RULES {_words('rule', 5000)} END"]) +async def test_jev_state_within_the_budget_is_sent_unchanged(system_prompt: str) -> None: + recorder: Final = _StateRecorder() + router: Final = _budget_router(recorder) + messages: Final = ( + {"role": "user", "content": "earlier question about the retry path"}, + {"role": "assistant", "content": "noted"}, + {"role": "user", "content": "rewrite it"}, + ) + outcome: Final = await router.aclassify("rewrite it", system_prompt=system_prompt, messages=messages) + assert outcome.cause == "jev_classifier" + assert recorder.states == (router._classifier_context_payload("rewrite it", system_prompt, None, messages),) + + +@pytest.mark.asyncio +async def test_jev_oversized_operator_instructions_use_fallback_without_sending_a_request() -> None: + recorder: Final = _StateRecorder() + router: Final = _budget_router(recorder, jev_config={"instructions": _words("rubric", 30_000)}) + outcome: Final = await router.aclassify("summarize this document") + assert outcome.cause == "heuristic_scorer" + assert outcome.classifier_failure_reason == "classifier_error" + assert recorder.states == () + + +@pytest.mark.asyncio +async def test_jev_oversized_fixed_history_uses_fallback_without_sending_a_request() -> None: + recorder: Final = _StateRecorder() + router: Final = _budget_router( + recorder, + model_list=[ + { + "model_name": "decision", + "litellm_params": {"model": "azure_ai/decision-1"}, + "model_info": {"max_input_tokens": 512}, + } + ], + jev_config={"deployment_name": "decision"}, + ) + outcome: Final = await router.aclassify( + "summarize this document", + messages=[{"role": "user", "content": _words("history", 2000)}], + ) + assert outcome.cause == "heuristic_scorer" + assert outcome.classifier_failure_reason == "classifier_error" + assert recorder.states == () + + +@pytest.mark.asyncio +async def test_jev_oversized_current_ask_is_cut_to_the_budget_keeping_both_ends() -> None: + recorder: Final = _StateRecorder() + ask: Final = f"ASK-HEAD {_words('ask', 60_000)} ASK-TAIL" + outcome: Final = await _budget_router(recorder).aclassify( + ask, system_prompt="caller constraints", messages=[{"role": "user", "content": ask}] + ) + state: Final = recorder.states[0] + assert outcome.cause == "jev_classifier" + assert litellm.token_counter(text=state) <= _DEFAULT_JEV_STATE_BUDGET + assert "caller constraints" in state + assert ("ASK-HEAD" in state, "ASK-TAIL" in state, len(state) < len(ask)) == (True, True, True) + + +@pytest.mark.asyncio +async def test_jev_oversized_system_prompt_is_cut_before_the_current_ask() -> None: + recorder: Final = _StateRecorder() + ask: Final = "rewrite the retry path and prove it cannot livelock" + system_prompt: Final = f"SYS-HEAD {_words('rule', 60_000)} SYS-TAIL" + outcome: Final = await _budget_router(recorder).aclassify(ask, system_prompt=system_prompt) + state: Final = recorder.states[0] + assert outcome.cause == "jev_classifier" + assert litellm.token_counter(text=state) <= _DEFAULT_JEV_STATE_BUDGET + assert f"Classify this message:\n{ask}" in state + assert ("SYS-HEAD" in state, "SYS-TAIL" in state, len(state) < len(system_prompt)) == (True, True, True) + + +@pytest.mark.asyncio +async def test_jev_cuts_the_system_prompt_to_its_floor_then_the_ask_when_both_are_oversized() -> None: + recorder: Final = _StateRecorder() + ask: Final = f"ASK-HEAD {_words('ask', 60_000)} ASK-TAIL" + outcome: Final = await _budget_router(recorder).aclassify( + ask, system_prompt=f"SYS-HEAD {_words('rule', 60_000)} SYS-TAIL" + ) + state: Final = recorder.states[0] + system_block: Final = state.split("\nClassify this message:\n")[0] + assert outcome.cause == "jev_classifier" + assert litellm.token_counter(text=state) <= _DEFAULT_JEV_STATE_BUDGET + assert ("SYS-HEAD" in system_block, "SYS-TAIL" in system_block, len(system_block) < 200) == (True, True, True) + assert ("ASK-HEAD" in state, "ASK-TAIL" in state) == (True, True) + + +@pytest.mark.asyncio +async def test_jev_state_budget_follows_the_decision_deployment_window() -> None: + recorder: Final = _StateRecorder() + router: Final = _budget_router( + recorder, + model_list=[ + { + "model_name": "decision", + "litellm_params": {"model": "azure_ai/decision-1"}, + "model_info": {"max_input_tokens": 4000}, + } + ], + jev_config={"deployment_name": "decision"}, + ) + outcome: Final = await router.aclassify(f"ASK-HEAD {_words('ask', 10_000)} ASK-TAIL") + state: Final = recorder.states[0] + assert outcome.cause == "jev_classifier" + assert litellm.token_counter(text=state) <= 3000 + assert ("ASK-HEAD" in state, "ASK-TAIL" in state) == (True, True) + + +@pytest.mark.asyncio +async def test_jev_state_budget_defaults_quietly_for_an_unmapped_decision_deployment( + caplog: pytest.LogCaptureFixture, +) -> None: + recorder: Final = _StateRecorder() + router: Final = _budget_router( + recorder, + model_list=[ + { + "model_name": "decision", + "litellm_params": {"model": "azure_ai/decision-1", "api_base": "https://decision.test"}, + } + ], + jev_config={"deployment_name": "decision"}, + ) + with caplog.at_level("WARNING"): + outcome: Final = await router.aclassify(f"ASK-HEAD {_words('ask', 60_000)} ASK-TAIL") + state: Final = recorder.states[0] + assert outcome.cause == "jev_classifier" + assert litellm.token_counter(text=state) <= _DEFAULT_JEV_STATE_BUDGET + assert ("ASK-HEAD" in state, "ASK-TAIL" in state) == (True, True) + assert caplog.records == []