fix(router): budget classifier input by provider limits (#49306)

This commit is contained in:
tin-berri 2026-10-30 14:38:14 -07:00 • committed by GitHub
parent 35d3a1784b
commit ec4f66bb94
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 589 additions and 55 deletions

View file

@ -218,7 +218,7 @@ def supported_decisions_providers() -> tuple[str, ...]:
)
def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
def get_decisions_provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig:
provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None)
provider_config: Final = (
None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider)
@ -408,7 +408,7 @@ def _provider_ir_request(
return supported
def _outbound_headers(
def decisions_outbound_headers(
extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object]
) -> Mapping[str, str]:
forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
@ -466,7 +466,7 @@ def prepare_call(
api_base=call_api_base,
api_key=api_key,
)
provider_config: Final = _provider_config(upstream_model, provider)
provider_config: Final = get_decisions_provider_config(upstream_model, provider)
if not upstream_model:
raise litellm.BadRequestError(
message="A model name is required for the Decisions API",
@ -478,7 +478,7 @@ def prepare_call(
except ValueError as error:
raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error
try:
headers: Final = _outbound_headers(extra_headers, provider, kwargs)
headers: Final = decisions_outbound_headers(extra_headers, provider, kwargs)
request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs)
request: Final = _fields_without_unsupported_media(
fields, provider_config, canonical_model, provider, kwargs

View file

@ -21,6 +21,11 @@ _ROUTING_SCALAR_FIELDS: Final = frozenset(
"classifier_cost",
"classifier_failure_reason",
"classifier_error_type",
"classifier_input_budget_tokens",
"classifier_input_tokens",
"classifier_input_count_method",
"classifier_input_truncated",
"classifier_input_budget_status",
"classifier_confidence",
"classifier_primary_rule",
"classifier_capability_boundary",

View file

@ -12,7 +12,7 @@ from litellm.llms.azure_ai.common_utils import (
AZURE_ENTRA_LITELLM_PARAM_KEYS,
get_azure_ai_auth_headers,
)
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits
_FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$")
_AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"})
@ -26,6 +26,16 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig):
api_base_env = ("AZURE_AI_API_BASE",)
api_key_required = False
def get_input_limits(
self, model: str, custom_llm_provider: str, headers: Mapping[str, str]
) -> DecisionsInputLimits:
limits: Final = super().get_input_limits(model, custom_llm_provider, headers)
return (
limits
if limits.max_request_tokens is not None
else super().get_input_limits("Microsoft-Decision-1", custom_llm_provider, headers)
)
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://<resource>.services.ai.azure.com"

View file

@ -1,7 +1,7 @@
import json
from abc import ABC
from collections.abc import Callable, Mapping, Sequence
from dataclasses import replace
from dataclasses import dataclass, replace
from itertools import chain
from types import MappingProxyType
from typing import Final
@ -511,6 +511,12 @@ def _is_decisions_model(model: str, custom_llm_provider: str) -> bool:
return is_decisions_model_mode(model_mode)
@dataclass(frozen=True, slots=True)
class DecisionsInputLimits:
max_request_tokens: int | None = None
max_question_tokens: int | None = None
class BaseDecisionsConfig(ABC):
path: str = "/v1/systemone"
api_key_env: tuple[str, ...] = ()
@ -533,6 +539,19 @@ class BaseDecisionsConfig(ABC):
def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None:
return None
def get_input_limits(
self, model: str, custom_llm_provider: str, headers: Mapping[str, str]
) -> DecisionsInputLimits:
try:
limit: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get(
"max_input_tokens"
)
except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model
return DecisionsInputLimits()
return DecisionsInputLimits(
max_request_tokens=limit if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0 else None
)
def missing_api_base_message(self, custom_llm_provider: str) -> str:
return f"api_base is required for Decisions provider '{custom_llm_provider}'"

View file

@ -6,7 +6,7 @@ from types import MappingProxyType
from typing import Final, Literal
from litellm.constants import TYPESAFE_DANGER_ZONE_HEADER_PREFIX
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, image_media_type
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits, image_media_type
from litellm.types.decisions import (
DecisionsAsset,
DecisionsIRMessages,
@ -19,6 +19,7 @@ from litellm.types.decisions import (
)
_IMAGES_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}images"
_LONG_CONTEXT_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}long-context"
_ASSET_FORMATS: Final[Mapping[str, Literal["png", "jpg"]]] = MappingProxyType(
{"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg"}
)
@ -84,6 +85,16 @@ class TypeSafeDecisionsConfig(BaseDecisionsConfig):
def get_default_api_base(self) -> str | None:
return "https://api.typesafe.ai"
def get_input_limits(
self, model: str, custom_llm_provider: str, headers: Mapping[str, str]
) -> DecisionsInputLimits:
long_context_values: Final = tuple(
value.strip().lower() for name, value in headers.items() if name.lower() == _LONG_CONTEXT_HEADER
)
if long_context_values and all(value == "true" for value in long_context_values):
return DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144)
return DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768)
def supports_images_in(
self, decision_input: DecisionsIRState | DecisionsIRMessages, headers: Mapping[str, str]
) -> bool:

View file

@ -1216,6 +1216,9 @@ class RoutingContext:
def registry(self) -> RouterRegistry:
return self._router.registry()
def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]:
return self._router._params_sent_to_deployment(request, deployment)
def settings(self) -> RouterSettings:
return self._router.settings()

View file

@ -126,10 +126,13 @@ from .jev_classifier import (
DEFAULT_JEV_INSTRUCTIONS,
HttpJevClassifierClient,
JevClassifierClient,
JevSystemOneRequest,
JevVerdict,
RouterJevClassifierClient,
build_jev_request,
fit_jev_request,
jev_classifier_cost,
jev_input_targets,
)
from .llm_v2 import LLM_V2_PROMPT_VERSION, LLMV2Decision, LLMV2TaskContext, LLMV2Verdict, llm_v2_response_format
from .stall_detector import detect_stalled_task
@ -1166,6 +1169,7 @@ class ClassificationOutcome(NamedTuple):
heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None
classifier_failure_reason: ClassifierFailureReason | None = None
classifier_error_type: str | None = None
classifier_input_telemetry: StandardLoggingRoutingDecision | None = None
def _with_classifier_failure(
@ -1207,23 +1211,24 @@ def _with_classifier_forecast(
decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome
) -> StandardLoggingRoutingDecision:
"""Attach validated forecasts and their applied policy to the routing decision."""
with_input: Final[StandardLoggingRoutingDecision] = {**decision, **(outcome.classifier_input_telemetry or {})}
if outcome.heuristic_v2_forecast is not None:
return {**decision, "heuristic_v2_forecast": outcome.heuristic_v2_forecast}
return {**with_input, "heuristic_v2_forecast": outcome.heuristic_v2_forecast}
if outcome.jev_verdict is not None:
forecasted_decision: Final[StandardLoggingRoutingDecision] = {
**decision,
**with_input,
"classifier_probabilities": outcome.jev_verdict.probabilities,
"classifier_confidence": outcome.jev_verdict.confidence,
}
return forecasted_decision
if outcome.llm_v2_forecast is not None:
return _with_llm_v2_forecast(decision, outcome.llm_v2_forecast)
return _with_llm_v2_forecast(with_input, outcome.llm_v2_forecast)
forecast: Final = outcome.capability_forecast
if forecast is None:
return decision
return with_input
verdict: Final = forecast.verdict
enriched: Final[StandardLoggingRoutingDecision] = {
**decision,
**with_input,
"classifier_crux": verdict.crux,
"classifier_primary_rule": verdict.primary_rule,
"classifier_capability_boundary": verdict.capability_boundary,
@ -1974,19 +1979,21 @@ class ComplexityRouter(CustomLogger):
`score` because the score is only interpretable against the boundaries that
mapped it to a tier.
"""
decision: Final = StandardLoggingRoutingDecision(
router_model_name=self.model_name,
router_type="complexity",
routed_model=routed_model,
cause=cause,
conversation_continuing=conversation_continuing,
**(
{"classifier_failure_reason": classifier_failure_reason}
if classifier_failure_reason is not None
else {}
),
**({"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {}),
failure: Final[StandardLoggingRoutingDecision] = (
{"classifier_failure_reason": classifier_failure_reason} if classifier_failure_reason is not None else {}
)
error: Final[StandardLoggingRoutingDecision] = (
{"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {}
)
decision: Final[StandardLoggingRoutingDecision] = {
"router_model_name": self.model_name,
"router_type": "complexity",
"routed_model": routed_model,
"cause": cause,
"conversation_continuing": conversation_continuing,
**failure,
**error,
}
if (baseline := self.savings_baseline) is not None:
decision["savings_baseline_model"] = baseline.model
if baseline.deployment_id is not None:
@ -2371,16 +2378,6 @@ class ComplexityRouter(CustomLogger):
system_prompt,
failure_reason="unsupported_input",
)
breaker: Final = self._classifier_circuit_breaker
permit: Final = breaker.acquire_permit() if breaker is not None else None
if breaker is not None and permit is None:
return self._classifier_failure_outcome(
"jev classifier circuit is open",
prompt,
system_prompt,
failure_reason="circuit_open",
signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL,
)
criteria: Final[Mapping[str, str]] = (
MappingProxyType(
{
@ -2395,13 +2392,40 @@ 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,
from litellm.litellm_core_utils.token_counter import offload_token_count
def build_request(current_prompt: str, caller_constraints: str | None) -> JevSystemOneRequest:
return build_jev_request(
prompt=self._classifier_context_payload(
current_prompt, caller_constraints, request_kwargs, messages, current_ask_for_history=prompt
),
system_prompt=None,
model=config.model,
instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS,
criteria=criteria,
)
budgeted: Final = await offload_token_count(fit_jev_request)(
build_request, prompt, system_prompt, jev_input_targets(self.litellm_router_instance, config), _truncate
)
request: Final = budgeted.request
if request is None:
return self._classifier_failure_outcome(
"classifier input cannot fit the provider budget",
prompt,
system_prompt,
failure_reason="unsupported_input",
)._replace(classifier_input_telemetry=budgeted.telemetry)
breaker: Final = self._classifier_circuit_breaker
permit: Final = breaker.acquire_permit() if breaker is not None else None
if breaker is not None and permit is None:
return self._classifier_failure_outcome(
"jev classifier circuit is open",
prompt,
system_prompt,
failure_reason="circuit_open",
signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL,
)._replace(classifier_input_telemetry=budgeted.telemetry)
try:
response: Final = await asyncio.wait_for(client.evaluate(request, timeout_s, request_kwargs), timeout_s)
answer: Final = response.answers.get("tier")
@ -2449,6 +2473,7 @@ class ComplexityRouter(CustomLogger):
cause="jev_classifier",
classifier_cost=verdict.cost,
jev_verdict=verdict,
classifier_input_telemetry=budgeted.telemetry,
)
except asyncio.CancelledError:
if breaker is not None and permit is not None:
@ -2463,7 +2488,7 @@ class ComplexityRouter(CustomLogger):
system_prompt,
failure_reason=_classifier_failure_reason(e),
error_type=type(e).__name__,
)
)._replace(classifier_input_telemetry=budgeted.telemetry)
def _classifier_failure_outcome(
self,
@ -2654,6 +2679,7 @@ class ComplexityRouter(CustomLogger):
messages: Sequence[Mapping[str, object]] | None,
*,
encrypted_task: bool = False,
current_ask_for_history: str | None = None,
) -> str:
include_assistant: Final = self.config.classifier_context_include_assistant_turns
marker_pairs: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
@ -2661,7 +2687,7 @@ class ComplexityRouter(CustomLogger):
prior_turns: Final = (
_extract_prior_turns(
messages,
current_ask=prompt,
current_ask=prompt if current_ask_for_history is None else current_ask_for_history,
window_size=self.config.classifier_context_window_size,
budget_chars=self.config.classifier_context_budget_chars,
per_turn_chars=self.config.classifier_context_per_turn_chars,
@ -4983,16 +5009,19 @@ class ComplexityRouter(CustomLogger):
return PreRoutingHookResponse(
model=fallback_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=fallback_model,
conversation_continuing=conversation_continuing,
cause=outcome.cause,
signals=outcome.signals,
classifier_failure_reason=outcome.classifier_failure_reason,
classifier_error_type=outcome.classifier_error_type,
escalation_keyword=escalation_keyword,
escalated=False,
tier_litellm_params=fallback_tier_params,
routing_decision=_with_classifier_forecast(
self._build_routing_decision(
routed_model=fallback_model,
conversation_continuing=conversation_continuing,
cause=outcome.cause,
signals=outcome.signals,
classifier_failure_reason=outcome.classifier_failure_reason,
classifier_error_type=outcome.classifier_error_type,
escalation_keyword=escalation_keyword,
escalated=False,
tier_litellm_params=fallback_tier_params,
),
outcome,
),
litellm_params=fallback_tier_params,
)

View file

@ -1,4 +1,5 @@
from collections.abc import Mapping
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias
@ -25,10 +26,13 @@ from litellm.llms.pass_through.typesafe_logging_handler import TypeSafePassthrou
from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS
from litellm.types.decisions import DecisionsResponse
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN
from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, StandardLoggingRoutingDecision
if TYPE_CHECKING:
from litellm.router import Router
from litellm.router_api import RoutingContext
from litellm.router_strategy.complexity_router.config import OpenSourceClassifierConfig
from litellm.types.router import Deployment
JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)]
ClassifierProvider: TypeAlias = Literal["typesafe", "laya", "bespoke", "databricks"]
@ -52,6 +56,155 @@ class JevSystemOneRequest(LiteLLMBaseModel):
questions: Mapping[str, JevChoiceQuestion]
@dataclass(frozen=True, slots=True)
class JevInputTarget:
model: str
max_tokens: int | None
@dataclass(frozen=True, slots=True)
class JevBudgetedRequest:
request: JevSystemOneRequest | None
telemetry: StandardLoggingRoutingDecision
def jev_input_targets(router: "Router", config: "OpenSourceClassifierConfig") -> tuple[JevInputTarget, ...]:
from litellm.router_api import RoutingContext
from litellm.types.router import Deployment
if config.deployment_name is None:
provider: Final = "typesafe" if config.provider == "jev" else config.provider
return (_jev_input_target({"model": f"{provider}/{config.model}"}, None),)
context: Final = RoutingContext(router)
deployments: Final = context.registry().deployments(model_name=config.deployment_name) or ()
return tuple(
_jev_deployment_input_target(context, Deployment.model_validate(deployment)) for deployment in deployments
)
def _jev_deployment_input_target(context: "RoutingContext", deployment: "Deployment") -> JevInputTarget:
params: Final = context.deployment_params({}, deployment)
explicit_limit: Final = deployment.model_info.model_dump().get("max_input_tokens")
return _jev_input_target(params, explicit_limit)
def _jev_input_target(params: Mapping[str, object], explicit_limit: object) -> JevInputTarget:
from litellm.decisions.call import decisions_outbound_headers, get_decisions_provider_config
from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata
qualified_model: Final = str(params.get("model") or "")
override: Final = params.get("custom_llm_provider")
try:
model, provider = resolve_model_and_provider_for_metadata(
qualified_model, override if isinstance(override, str) else None
)
headers: Final = decisions_outbound_headers(
TypeAdapter[Mapping[str, str] | None](Mapping[str, str] | None).validate_python(
params.get("extra_headers")
),
provider,
params,
)
limits: Final = get_decisions_provider_config(model, provider).get_input_limits(model, provider, headers)
candidates: Final = tuple(
value
for value in (limits.max_request_tokens, limits.max_question_tokens, explicit_limit)
if isinstance(value, int) and not isinstance(value, bool) and value > 0
)
return JevInputTarget(f"{provider}/{model}", min(candidates) if candidates else None)
except Exception: # noqa: BLE001 # unavailable metadata must not turn an existing classifier into a request failure
return JevInputTarget(qualified_model, None)
def _jev_request_count(request: JevSystemOneRequest, models: Sequence[str], budget: int) -> tuple[int, str]:
from litellm._lazy_imports import get_token_counter_new
payload: Final = request.model_dump_json(exclude={"model"})
byte_count: Final = len(payload.encode("utf-8"))
if byte_count <= budget or litellm.disable_token_counter:
return byte_count, "utf8_bytes"
try:
counts: Final = tuple(get_token_counter_new()(model=model, text=payload) for model in models)
if counts and all(count > 0 for count in counts):
return max(counts), "estimated_tokens"
except Exception: # noqa: BLE001 # tokenization is best-effort; byte counting remains available
return byte_count, "utf8_bytes"
return byte_count, "utf8_bytes"
def fit_jev_request(
build_request: Callable[[str, str | None], JevSystemOneRequest],
prompt: str,
system_prompt: str | None,
targets: Sequence[JevInputTarget],
truncate: Callable[[str, int], str],
) -> JevBudgetedRequest:
known: Final = tuple(target for target in targets if target.max_tokens is not None)
if not known:
return JevBudgetedRequest(
build_request(prompt, system_prompt),
{"classifier_input_budget_status": "unknown_limits", "classifier_input_truncated": False},
)
budget: Final = max(1, int(min(target.max_tokens for target in known if target.max_tokens is not None) * 0.75))
models: Final = tuple(dict.fromkeys(target.model for target in known))
initial_cap: Final = max(120, budget * 8)
bounded_prompt: Final = truncate(prompt, initial_cap)
bounded_system: Final = truncate(system_prompt, initial_cap) if system_prompt is not None else None
return _fit_jev_components(
build_request,
bounded_prompt,
bounded_system,
models,
budget,
truncate,
len(known) != len(targets),
bounded_prompt != prompt or bounded_system != system_prompt,
)
def _fit_jev_components(
build_request: Callable[[str, str | None], JevSystemOneRequest],
prompt: str,
system_prompt: str | None,
models: Sequence[str],
budget: int,
truncate: Callable[[str, int], str],
unknown_limits: bool,
truncated: bool,
) -> JevBudgetedRequest:
attempts: Final = len(prompt).bit_length() + len(system_prompt or "").bit_length() + 1
components: tuple[str, str | None, bool, int] = ( # rebind-ok: bounded iterative fitting state
prompt,
system_prompt,
truncated,
attempts,
)
while True:
current_prompt, current_system, was_truncated, remaining = components
request = build_request(current_prompt, current_system)
count, method = _jev_request_count(request, models, budget)
telemetry: StandardLoggingRoutingDecision = {
"classifier_input_budget_tokens": budget,
"classifier_input_tokens": count,
"classifier_input_count_method": method,
"classifier_input_truncated": was_truncated,
"classifier_input_budget_status": "unknown_limits"
if unknown_limits
else "truncated"
if was_truncated
else "within_budget",
}
if count <= budget:
return JevBudgetedRequest(request, telemetry)
if remaining == 0 or (len(current_system or "") <= 123 and len(current_prompt) <= 123):
return JevBudgetedRequest(None, {**telemetry, "classifier_input_budget_status": "untrimmable_input"})
components = (
(current_prompt, truncate(current_system, max(120, len(current_system) // 2)), True, remaining - 1)
if current_system is not None and len(current_system) > 123
else (truncate(current_prompt, max(120, len(current_prompt) // 2)), current_system, True, remaining - 1)
)
class JevChoiceAnswer(LiteLLMBaseModel):
model_config = ConfigDict(frozen=True, allow_inf_nan=False)

View file

@ -3323,6 +3323,11 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
classifier_cost: float
classifier_failure_reason: ReadOnly[ClassifierFailureReason]
classifier_error_type: ReadOnly[str]
classifier_input_budget_tokens: ReadOnly[int]
classifier_input_tokens: ReadOnly[int]
classifier_input_count_method: ReadOnly[str]
classifier_input_truncated: ReadOnly[bool]
classifier_input_budget_status: ReadOnly[str]
classifier_probabilities: ReadOnly[Mapping[str, float]]
classifier_confidence: ReadOnly[float]
heuristic_v2_forecast: ReadOnly[StandardLoggingHeuristicV2Forecast]
@ -3381,6 +3386,11 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
"classifier_cost",
"classifier_failure_reason",
"classifier_error_type",
"classifier_input_budget_tokens",
"classifier_input_tokens",
"classifier_input_count_method",
"classifier_input_truncated",
"classifier_input_budget_status",
"classifier_probabilities",
"classifier_confidence",
"heuristic_v2_forecast",

View file

@ -207,6 +207,11 @@ def test_routing_diagnostics_reach_llm_spans_without_promoting_prompt_text(fail:
"classifier_cost": 0.001,
"classifier_failure_reason": "timeout",
"classifier_error_type": "TimeoutError",
"classifier_input_budget_tokens": 24576,
"classifier_input_tokens": 23000,
"classifier_input_count_method": "estimated_tokens",
"classifier_input_truncated": True,
"classifier_input_budget_status": "truncated",
"escalated": False,
}
payload: Final = _payload(

View file

@ -344,6 +344,11 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada
"cause": "heuristic",
"classifier_failure_reason": "timeout",
"classifier_error_type": "TimeoutError",
"classifier_input_budget_tokens": 24576,
"classifier_input_tokens": 23000,
"classifier_input_count_method": "estimated_tokens",
"classifier_input_truncated": True,
"classifier_input_budget_status": "truncated",
"classifier_cost": 0.001,
}
with provider.get_tracer(__name__).start_as_current_span("inference") as span:

View file

@ -6,14 +6,19 @@ import httpx
import pytest
from pydantic import TypeAdapter, ValidationError
import litellm
from litellm.llms.azure_ai.decisions.transformation import AzureAIDecisionsConfig
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.llms.base_llm.decisions.transformation import (
BaseDecisionsConfig,
DecisionsInputLimits,
ir_to_systemone_request,
ir_to_systemone_response,
parse_systemone_response,
systemone_request_to_ir,
)
from litellm.llms.openai.decisions.transformation import openai_request_to_ir
from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig
from litellm.types.decisions import (
MAX_DECISION_QUESTIONS,
DecisionsIRRefusal,
@ -49,6 +54,61 @@ _SYSTEMONE_REQUEST: Final[Mapping[str, object]] = {
}
@pytest.mark.parametrize("declared", (4096, 0, -1, True, None))
def test_decision_input_limits_use_only_positive_catalog_limits(
local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, declared: int | None
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"bespoke/test-input-limit",
{"litellm_provider": "bespoke", "mode": "evaluation", "max_input_tokens": declared},
)
limits: Final = BaseDecisionsConfig().get_input_limits("test-input-limit", "bespoke", {})
assert limits == DecisionsInputLimits(max_request_tokens=4096 if declared == 4096 else None)
def test_unmapped_decision_model_limits_remain_unknown(local_model_cost_map: None) -> None:
assert BaseDecisionsConfig().get_input_limits("unmapped-input-limit", "databricks", {}) == DecisionsInputLimits()
@pytest.mark.parametrize(
("headers", "long_context"),
(
({}, False),
({"TypeSafe-Danger-Zone-Long-Context": "false"}, False),
({"typesafe-danger-zone-long-context": " TRUE "}, True),
({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "true"}, True),
({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "false"}, False),
({"typesafe-danger-zone-long-context": "true,false"}, False),
),
)
def test_typesafe_limits_require_unambiguous_long_context_opt_in(
headers: Mapping[str, str], long_context: bool
) -> None:
"""Provider contract from https://docs.typesafe.ai/models, verified 2026-10-30"""
limits: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", headers)
assert limits == (
DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144)
if long_context
else DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768)
)
@pytest.mark.parametrize("model", ("Microsoft-Decision-1", "unmapped-decision-deployment"))
def test_azure_decision_deployments_use_catalog_limits_without_typesafe_modes(
local_model_cost_map: None, model: str
) -> None:
catalog_limit: Final = litellm.get_model_info("Microsoft-Decision-1", "azure_ai")["max_input_tokens"]
limits: Final = AzureAIDecisionsConfig().get_input_limits(
model, "azure_ai", {"TypeSafe-Danger-Zone-Long-Context": "true"}
)
assert limits == DecisionsInputLimits(max_request_tokens=catalog_limit)
def _openai_ir(raw: Mapping[str, object]) -> DecisionsIRRequest:
return openai_request_to_ir(_OPENAI_BODY.validate_python(raw))

View file

@ -9,6 +9,7 @@ from unittest.mock import create_autospec
import httpx
import pytest
import respx
import tiktoken
import litellm
from litellm._logging import verbose_router_logger
@ -16,12 +17,14 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter, _ClassifierCircuitBreaker
from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, JevClassifierConfig
from litellm.router_strategy.complexity_router.jev_classifier import (
DEFAULT_JEV_INSTRUCTIONS,
HttpJevClassifierClient,
JevChoiceAnswer,
JevSystemOneRequest,
JevSystemOneResponse,
JevUsage,
build_jev_request,
@ -1055,3 +1058,214 @@ 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"}
@pytest.mark.asyncio
@pytest.mark.parametrize("large_field", ["user", "system"])
@pytest.mark.parametrize("disable_counter", [False, True])
async def test_jev_bounds_complete_input_without_changing_forwarded_messages(
large_field: str, disable_counter: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_token_counter", disable_counter)
limit: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", {}).max_question_tokens
assert limit is not None
rubric: Final = "Operator routing rubric " + "document the task category " * 500
criterion: Final = ("Simple tasks " + "the classification criteria " * 15).rstrip()
large_text: Final = "BEGIN_SOURCE " + "budget 预算 function(x) return x; " * 18000 + " END_SOURCE"
ask: Final = large_text if large_field == "user" else "CURRENT_ASK Explain the example"
constraints: Final = large_text if large_field == "system" else "Use concise language"
messages: Final = [
{"role": "system", "content": constraints},
{"role": "user", "content": "PRIOR_TURN Explain the budget"},
{"role": "user", "content": ask},
]
original: Final = deepcopy(messages)
handler: Final = AsyncHTTPHandler()
handler.client = httpx.AsyncClient()
classifier: Final = ComplexityRouter(
"bounded-classifier", litellm.Router(model_list=[]),
{
"classifier_type": "oss_classifier",
"opensource_classifier_config": {"instructions": rubric},
"tier_definitions": [{"name": "SIMPLE", "description": criterion}, {"name": "REASONING"}],
"tiers": {"SIMPLE": "cheap", "REASONING": "deep"},
"fallback_tier": "SIMPLE",
"session_affinity": False,
"deployment_affinity": False,
},
jev_client=HttpJevClassifierClient("test", "https://budget.test", handler),
derive_savings_baseline=False,
)
try:
with respx.mock(assert_all_called=True) as upstream:
route: Final = upstream.post("https://budget.test/v1/systemone").respond(
json={"answers": {"tier": _answer().model_dump()}}
)
result: Final = await classifier.async_pre_routing_hook(
model="bounded-classifier", request_kwargs={"messages": messages}, messages=messages
)
sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content)
complete_input: Final = sent.model_dump_json(exclude={"model"})
measured: Final = len(tiktoken.get_encoding("o200k_base").encode(complete_input))
assert measured < limit
assert sent.questions["tier"].instructions == rubric
assert sent.questions["tier"].criteria["SIMPLE"] == criterion
assert sent.state.count("BEGIN_SOURCE") == sent.state.count("END_SOURCE") == 1
assert sent.state.count("PRIOR_TURN") == 1
assert result is not None and result.routing_decision is not None
decision: Final = result.routing_decision
assert decision["cause"] == "jev_classifier"
assert decision["classifier_input_budget_status"] == "truncated"
assert decision["classifier_input_truncated"] is True
assert measured <= decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] < limit
assert messages == original
assert result.messages is None or result.messages == original
if disable_counter:
assert len(complete_input.encode()) < limit
assert decision["classifier_input_count_method"] == "utf8_bytes"
finally:
await GLOBAL_LOGGING_WORKER.flush()
await handler.client.aclose()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("mode", "default_header", "explicit_cap", "other_member", "expected_status"),
[
("false", None, None, None, "truncated"),
("true", None, None, None, "within_budget"),
("true", "false", None, None, "truncated"),
("true", None, 8192, None, "truncated"),
("true", None, 8192, "default", "truncated"),
("true", None, None, "default", "truncated"),
("false", None, None, "unknown", "unknown_limits"),
],
)
async def test_jev_budget_uses_effective_headers_and_every_deployment(
mode: str, default_header: str | None, explicit_cap: int | None,
other_member: str | None, expected_status: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"]))
header: Final = "TypeSafe-Danger-Zone-Long-Context"
deployment: Final = {
"model_name": "budget-model",
"litellm_params": {
"model": "typesafe/jev-latest", "api_key": "test", "api_base": "https://budget.test",
"extra_headers": {header: mode},
},
"model_info": {"max_input_tokens": explicit_cap} if explicit_cap is not None else {},
}
sibling: Final = {
"model_name": "budget-model",
"litellm_params": {
"model": "databricks/unknown-budget-endpoint" if other_member == "unknown" else "typesafe/jev-latest",
"api_key": "test", "api_base": "https://budget.test",
},
}
router: Final = litellm.Router(
model_list=[deployment, sibling] if other_member else [deployment],
default_litellm_params={"extra_headers": {header: default_header}} if default_header else {},
num_retries=0,
)
classifier: Final = ComplexityRouter(
"bounded-classifier", router,
{"classifier_type": "oss_classifier", "opensource_classifier_config": {"deployment_name": "budget-model"},
"tiers": {"SIMPLE": "cheap"}},
derive_savings_baseline=False,
)
ask: Final = "Explain this code " + "word001 word002 word003 " * 12000
with respx.mock(assert_all_called=True) as upstream:
route: Final = upstream.route(method="POST").respond(
json={"answers": {"tier": _answer().model_dump()}}
)
outcome: Final = await classifier.aclassify(ask)
assert outcome.cause == "jev_classifier"
telemetry: Final = outcome.classifier_input_telemetry
assert telemetry is not None and telemetry["classifier_input_budget_status"] == expected_status
assert telemetry["classifier_input_tokens"] <= telemetry["classifier_input_budget_tokens"]
sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content)
if other_member is None:
assert route.calls[0].request.headers[header] == (default_header or mode)
if explicit_cap is not None:
assert telemetry["classifier_input_budget_tokens"] < explicit_cap
assert (ask in sent.state) is (expected_status == "within_budget")
@pytest.mark.asyncio
@pytest.mark.parametrize(
("oversized_field", "fallback"),
[("instructions", "heuristic"), ("instructions", "default_model"), ("criteria", "custom")],
)
async def test_jev_untrimmable_rubric_uses_configured_fallback_without_a_provider_call(
oversized_field: str, fallback: str, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"]))
router: Final = litellm.Router(model_list=[{
"model_name": "budget-model",
"litellm_params": {"model": "typesafe/jev-latest", "api_key": "test", "api_base": "https://budget.test"},
"model_info": {"max_input_tokens": 128 if oversized_field == "criteria" else 1024},
}])
rubric: Final = "word001 word002 word003 " * 2000
classifier: Final = ComplexityRouter(
"bounded-classifier", router,
{"classifier_type": "oss_classifier", "opensource_classifier_config": {
"deployment_name": "budget-model", "instructions": rubric if oversized_field == "instructions" else "Choose",
}, "tiers": {"SIMPLE": "cheap", "REASONING": "deep"},
**({"tier_definitions": [{"name": name, "description": "word001 " * 50}
for name in ("SIMPLE", "REASONING")], "fallback_tier": "SIMPLE"}
if fallback == "custom" else {"classifier_fallback": fallback, "default_model": "default"})},
derive_savings_baseline=False,
)
with respx.mock(assert_all_called=False) as upstream:
route: Final = upstream.route(method="POST").respond(json={"answers": {"tier": _answer().model_dump()}})
outcome: Final = await classifier.aclassify("What is two plus two?")
assert route.call_count == 0
assert outcome.cause == {
"heuristic": "heuristic_scorer", "default_model": "default_model_fallback", "custom": "classifier_fallback",
}[fallback]
assert outcome.classifier_failure_reason == "unsupported_input"
assert outcome.classifier_input_telemetry is not None
assert outcome.classifier_input_telemetry["classifier_input_budget_status"] == "untrimmable_input"
@pytest.mark.asyncio
@pytest.mark.parametrize("fallback", ["heuristic", "default_model"])
async def test_jev_open_breaker_preserves_input_budget_in_fallback_decision(fallback: str) -> None:
def clock() -> float:
return 100.0
transport: Final = create_autospec(httpx.AsyncBaseTransport, instance=True)
transport.handle_async_request.side_effect = TimeoutError("controlled classifier timeout")
handler: Final = AsyncHTTPHandler()
handler.client = httpx.AsyncClient(transport=transport)
classifier: Final = ComplexityRouter(
"breaker-budget", litellm.Router(model_list=[]),
{"classifier_type": "oss_classifier", "opensource_classifier_config": {},
"tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, "classifier_fallback": fallback,
"default_model": "default", "session_affinity": False, "deployment_affinity": False},
jev_client=HttpJevClassifierClient("test", "https://budget.test", handler),
derive_savings_baseline=False,
)
classifier._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=clock)
try:
first: Final = await classifier.async_pre_routing_hook(
model="breaker-budget", request_kwargs={}, messages=[{"role": "user", "content": "Explain budgets"}]
)
second: Final = await classifier.async_pre_routing_hook(
model="breaker-budget", request_kwargs={},
messages=[{"role": "user", "content": "Explain this code " + "word001 word002 word003 " * 12000}],
)
assert first is not None and first.routing_decision is not None
assert first.routing_decision["classifier_failure_reason"] == "timeout"
assert second is not None and second.routing_decision is not None
decision: Final = second.routing_decision
assert decision["classifier_failure_reason"] == "circuit_open"
assert decision["cause"] == ("heuristic_scorer" if fallback == "heuristic" else "default_model_fallback")
assert decision.get("classifier_input_budget_status") == "truncated"
assert decision["classifier_input_truncated"] is True
assert 0 < decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"]
transport.handle_async_request.assert_awaited_once()
finally:
await handler.client.aclose()

View file

@ -46205,6 +46205,16 @@ export interface components {
* @enum {string}
*/
classifier_failure_reason?: "timeout" | "circuit_open" | "not_configured" | "unsupported_input" | "invalid_response" | "declined" | "classifier_error";
/** Classifier Input Budget Status */
classifier_input_budget_status?: string;
/** Classifier Input Budget Tokens */
classifier_input_budget_tokens?: number;
/** Classifier Input Count Method */
classifier_input_count_method?: string;
/** Classifier Input Tokens */
classifier_input_tokens?: number;
/** Classifier Input Truncated */
classifier_input_truncated?: boolean;
/** Classifier Max Quality Gap */
classifier_max_quality_gap?: number;
/** Classifier Model */