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

Merge 8697474909 into 38b4d31447

Browse source
This commit is contained in:
ishaan-j-berri 2026-10-28 09:48:13 -04:00 • committed by GitHub
commit 941b0e752a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 267 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

108
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -25,6 +25,7 @@ import time
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
from functools import partial
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
@ -508,6 +509,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"
@ -2373,14 +2379,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:
@ -2612,6 +2626,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:
@ -2634,7 +2710,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
@ -2656,9 +2739,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,

168
tests/unit/router_strategy/complexity_router/test_jev_classifier.py
View file

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