mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
merge: support V2 alongside released capability classifier
This commit is contained in:
Tin Chi Lo
commit
52295c09fa
38 changed files with 2692 additions and 103 deletions
|
|
@ -214,26 +214,33 @@ def _message_has_cache_control(message: Mapping[str, object]) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
|
||||
last_breakpoint: Final = max(
|
||||
(index for index, msg in enumerate(messages) if _message_has_cache_control(msg)),
|
||||
default=-1,
|
||||
)
|
||||
return tuple(range(last_breakpoint + 1))
|
||||
|
||||
|
||||
def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]:
|
||||
"""
|
||||
Return indices of messages that must never be compressed:
|
||||
- All system messages
|
||||
- The last user message
|
||||
- The last assistant message
|
||||
- Any message carrying an Anthropic cache_control breakpoint
|
||||
- Every message up to and including the last one carrying an Anthropic cache_control breakpoint
|
||||
|
||||
The last user message is what the model is being asked to act on right now,
|
||||
so compressing it replaces the live instruction with a marker. Compression
|
||||
guardrails share this policy; see the Headroom guardrail. A cache_control
|
||||
breakpoint pins the provider's prompt-cache prefix to that row's exact
|
||||
bytes, so rewriting a marked row anywhere in history turns the next
|
||||
request's cache read into a cache write.
|
||||
breakpoint pins the provider's prompt-cache prefix to the exact bytes of every
|
||||
row up to it, so rewriting any row inside that prefix turns the next request's
|
||||
cache read into a cache write.
|
||||
"""
|
||||
system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system")
|
||||
last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:]
|
||||
assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant")
|
||||
cache_control_indices: Final = tuple(index for index, msg in enumerate(messages) if _message_has_cache_control(msg))
|
||||
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + cache_control_indices))
|
||||
return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages)))
|
||||
|
||||
|
||||
def _combine_scores(
|
||||
|
|
|
|||
|
|
@ -644,6 +644,24 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"""Keep ``_response_ms`` / ``litellm_overhead_time_ms`` for a result that has no ``_hidden_params``."""
|
||||
self.response_timing_metrics = dict(timing_metrics) # mutable-ok: kept deep-copyable
|
||||
|
||||
def add_dynamic_callback(self, callback: CustomLogger) -> None:
|
||||
self.dynamic_input_callbacks = self._with_dynamic_callback(self.dynamic_input_callbacks, callback)
|
||||
self.dynamic_success_callbacks = self._with_dynamic_callback(self.dynamic_success_callbacks, callback)
|
||||
self.dynamic_async_success_callbacks = self._with_dynamic_callback(
|
||||
self.dynamic_async_success_callbacks, callback
|
||||
)
|
||||
self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, callback)
|
||||
self.dynamic_async_failure_callbacks = self._with_dynamic_callback(
|
||||
self.dynamic_async_failure_callbacks, callback
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _with_dynamic_callback(
|
||||
callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger
|
||||
) -> list[str | Callable | CustomLogger]:
|
||||
existing: Final = tuple(callbacks or ())
|
||||
return [*existing, *(() if callback in existing else (callback,))]
|
||||
|
||||
def process_dynamic_callbacks(self):
|
||||
"""
|
||||
Initializes CustomLogger compatible callbacks in self.dynamic_* callbacks
|
||||
|
|
|
|||
|
|
@ -4716,6 +4716,7 @@ class JWTAuthBuilderResult(TypedDict):
|
|||
org_id: str | None
|
||||
team_membership: LiteLLM_TeamMembership | None
|
||||
jwt_claims: dict # Decoded JWT token claims (avoids re-decoding)
|
||||
agent_id: ReadOnly[str | None]
|
||||
|
||||
|
||||
class ClientSideFallbackModel(TypedDict, total=False):
|
||||
|
|
@ -4954,6 +4955,14 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase):
|
|||
user_allowed_roles: list[str] | None = None
|
||||
user_id_upsert: bool = Field(default=False, description="If user doesn't exist, upsert them into the db.")
|
||||
end_user_id_jwt_field: str | None = None
|
||||
agent_id_jwt_field: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"The field in the JWT token that identifies the calling agent (e.g. 'azp' for a Microsoft Entra ID "
|
||||
"app token). Supports dot notation. The value is matched against a registered agent's agent_id, "
|
||||
"then agent_name, and the request is rejected when it matches neither."
|
||||
),
|
||||
)
|
||||
public_key_ttl: float = 600
|
||||
public_key_stale_ttl: float = Field(
|
||||
default=DEFAULT_JWKS_STALE_TTL,
|
||||
|
|
|
|||
|
|
@ -5773,8 +5773,7 @@ async def _organization_max_budget_check(
|
|||
if org_table.litellm_budget_table is not None:
|
||||
org_max_budget = org_table.litellm_budget_table.max_budget
|
||||
|
||||
# Only check if organization has a valid max_budget set
|
||||
if org_max_budget is None or org_max_budget <= 0:
|
||||
if org_max_budget is None:
|
||||
return
|
||||
|
||||
# Read spend from cross-pod counter (Redis-first) or cached object (fallback)
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ import hashlib
|
|||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -61,6 +61,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
from .auth_checks import (
|
||||
_allowed_routes_check,
|
||||
|
|
@ -127,6 +128,26 @@ class _UserInfoResponse(Protocol):
|
|||
def json(self) -> dict[str, object]: ...
|
||||
|
||||
|
||||
class AgentLookup(Protocol):
|
||||
"""The registered-agent lookups a JWT agent claim is matched against."""
|
||||
|
||||
def get_agent_by_id(self, agent_id: str) -> AgentResponse | None:
|
||||
"""The agent registered under ``agent_id``, if any."""
|
||||
|
||||
def get_agent_by_name(self, agent_name: str) -> AgentResponse | None:
|
||||
"""The agent registered under ``agent_name``, if any."""
|
||||
|
||||
|
||||
class _NoRegisteredAgents:
|
||||
"""The lookup in force until the proxy binds its agent registry: no agent is registered, so no claim matches."""
|
||||
|
||||
def get_agent_by_id(self, agent_id: str) -> None:
|
||||
return None
|
||||
|
||||
def get_agent_by_name(self, agent_name: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _discovery_document(response: _OIDCDiscoveryResponse) -> _OIDCDiscoveryBody:
|
||||
"""Decode an OIDC discovery response body."""
|
||||
return response.json()
|
||||
|
|
@ -198,6 +219,10 @@ class JWTHandler:
|
|||
self.leeway = 0
|
||||
# Per-cache-key locks so a TTL lapse triggers one refresh instead of one per in-flight request.
|
||||
self._refresh_locks: dict[str, asyncio.Lock] = {} # mutable-ok: lock registry, keyed by JWKS url
|
||||
self.agent_lookup: AgentLookup = _NoRegisteredAgents()
|
||||
|
||||
def bind_agent_lookup(self, agent_lookup: AgentLookup) -> None:
|
||||
self.agent_lookup = agent_lookup
|
||||
|
||||
def update_environment(
|
||||
self,
|
||||
|
|
@ -623,6 +648,12 @@ class JWTHandler:
|
|||
object_id = default_value
|
||||
return object_id
|
||||
|
||||
def get_agent_claim(self, token: Mapping[str, object]) -> str | None:
|
||||
if self.litellm_jwtauth.agent_id_jwt_field is None:
|
||||
return None
|
||||
claim: Final[object] = get_nested_value(data=token, key_path=self.litellm_jwtauth.agent_id_jwt_field)
|
||||
return claim if isinstance(claim, str) and claim else None
|
||||
|
||||
def get_org_id(self, token: dict, default_value: str | None) -> str | None:
|
||||
if self._has_trusted_issuer_normalized_claim(token=token, claim=self.LITELLM_ORG_ID_CLAIM):
|
||||
return token.get(self.LITELLM_ORG_ID_CLAIM)
|
||||
|
|
@ -1380,6 +1411,7 @@ class JWTAuthManager:
|
|||
api_key: str,
|
||||
jwt_valid_token: dict | None = None,
|
||||
user_email: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
) -> JWTAuthBuilderResult | None:
|
||||
"""Check admin status and route access permissions"""
|
||||
if not jwt_handler.is_admin(scopes=scopes):
|
||||
|
|
@ -1409,8 +1441,28 @@ class JWTAuthManager:
|
|||
org_id=org_id,
|
||||
team_membership=None,
|
||||
jwt_claims=jwt_valid_token or {},
|
||||
agent_id=agent_id,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def resolve_agent_id(
|
||||
jwt_handler: JWTHandler,
|
||||
jwt_valid_token: Mapping[str, object],
|
||||
agent_registry: AgentLookup,
|
||||
) -> str | None:
|
||||
agent_claim: Final = jwt_handler.get_agent_claim(token=jwt_valid_token)
|
||||
if agent_claim is None:
|
||||
return None
|
||||
agent: Final = agent_registry.get_agent_by_id(agent_id=agent_claim) or agent_registry.get_agent_by_name(
|
||||
agent_name=agent_claim
|
||||
)
|
||||
if agent is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=f"No registered agent matches JWT claim {jwt_handler.litellm_jwtauth.agent_id_jwt_field}={agent_claim}",
|
||||
)
|
||||
return agent.agent_id
|
||||
|
||||
@staticmethod
|
||||
async def find_and_validate_specific_team_id(
|
||||
jwt_handler: JWTHandler,
|
||||
|
|
@ -2268,9 +2320,23 @@ class JWTAuthManager:
|
|||
elif rbac_role == LitellmUserRoles.INTERNAL_USER:
|
||||
user_id = object_id
|
||||
|
||||
agent_id: Final = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token=jwt_valid_token,
|
||||
agent_registry=jwt_handler.agent_lookup,
|
||||
)
|
||||
|
||||
# Check admin access
|
||||
admin_result: Final = await JWTAuthManager.check_admin_access(
|
||||
jwt_handler, scopes, route, user_id, org_id, api_key, jwt_valid_token, user_email=user_email
|
||||
jwt_handler,
|
||||
scopes,
|
||||
route,
|
||||
user_id,
|
||||
org_id,
|
||||
api_key,
|
||||
jwt_valid_token,
|
||||
user_email=user_email,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
if admin_result:
|
||||
await JWTAuthManager._attach_team_from_header_for_admin(
|
||||
|
|
@ -2514,4 +2580,5 @@ class JWTAuthManager:
|
|||
token=api_key,
|
||||
team_membership=team_membership_object,
|
||||
jwt_claims=jwt_valid_token,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -852,6 +852,7 @@ async def _auto_register_jwt_mapping(
|
|||
user_id: str | None = None,
|
||||
org_id: str | None = None,
|
||||
end_user_id: str | None = None,
|
||||
agent_id: str | None = None,
|
||||
) -> UserAPIKeyAuth | None:
|
||||
"""
|
||||
Auto-register: create a new virtual key + mapping for an unrecognised JWT
|
||||
|
|
@ -884,6 +885,7 @@ async def _auto_register_jwt_mapping(
|
|||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
organization_id=org_id,
|
||||
agent_id=agent_id,
|
||||
metadata={
|
||||
"auto_registered": True,
|
||||
"jwt_claim_field": virtual_key_claim_field,
|
||||
|
|
@ -1567,6 +1569,7 @@ async def _user_api_key_auth_builder(
|
|||
org_id: Final = result["org_id"]
|
||||
team_membership: Final[LiteLLM_TeamMembership | None] = result.get("team_membership", None)
|
||||
jwt_claims = result.get("jwt_claims", None)
|
||||
agent_id: Final[str | None] = result.get("agent_id")
|
||||
|
||||
if is_proxy_admin:
|
||||
# Proxy admins authenticate via auth_builder (full
|
||||
|
|
@ -1592,6 +1595,7 @@ async def _user_api_key_auth_builder(
|
|||
end_user_id=end_user_id,
|
||||
parent_otel_span=parent_otel_span,
|
||||
jwt_claims=jwt_claims,
|
||||
agent_id=agent_id,
|
||||
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
|
||||
)
|
||||
|
||||
|
|
@ -1612,6 +1616,7 @@ async def _user_api_key_auth_builder(
|
|||
user_rpm_limit=(user_object.rpm_limit if user_object is not None else None),
|
||||
user_model_max_budget=(user_object.model_max_budget if user_object is not None else None),
|
||||
jwt_claims=jwt_claims,
|
||||
agent_id=agent_id,
|
||||
**team_grants(team_object=team_object, team_membership=team_membership, user_id=user_id),
|
||||
)
|
||||
|
||||
|
|
@ -1635,6 +1640,7 @@ async def _user_api_key_auth_builder(
|
|||
user_id=user_id,
|
||||
org_id=org_id,
|
||||
end_user_id=end_user_id,
|
||||
agent_id=agent_id,
|
||||
)
|
||||
if auto_registered is not None:
|
||||
auto_registered.jwt_claims = jwt_claims
|
||||
|
|
|
|||
|
|
@ -1856,9 +1856,9 @@ def validate_team_org_change(
|
|||
|
||||
# Check if the team's budget is less than the org's max_budget
|
||||
if (
|
||||
team.max_budget
|
||||
and organization.litellm_budget_table
|
||||
and organization.litellm_budget_table.max_budget
|
||||
team.max_budget is not None
|
||||
and organization.litellm_budget_table is not None
|
||||
and organization.litellm_budget_table.max_budget is not None
|
||||
and team.max_budget > organization.litellm_budget_table.max_budget
|
||||
):
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -9516,6 +9516,9 @@ class ProxyStartupEvent:
|
|||
user_api_key_cache=user_api_key_cache,
|
||||
litellm_jwtauth=litellm_jwtauth,
|
||||
)
|
||||
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
|
||||
|
||||
jwt_handler.bind_agent_lookup(global_agent_registry)
|
||||
|
||||
@classmethod
|
||||
def _add_proxy_budget_to_db(cls):
|
||||
|
|
|
|||
|
|
@ -430,6 +430,14 @@ def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback:
|
|||
detail.setdefault("guardrail_mode", event_hook)
|
||||
|
||||
|
||||
def _is_client_error_exception(exc: Exception) -> bool:
|
||||
if isinstance(exc, HTTPException):
|
||||
return exc.status_code < 500
|
||||
if isinstance(exc, ProxyException):
|
||||
return not (exc.code.isdigit() and int(exc.code) >= 500)
|
||||
return False
|
||||
|
||||
|
||||
def _exception_changes_request_flow(exc: BaseException) -> bool:
|
||||
"""
|
||||
True for guardrail exceptions the proxy turns into an alternate request flow
|
||||
|
|
@ -2886,9 +2894,7 @@ class ProxyLogging:
|
|||
|
||||
### ALERTING ###
|
||||
await self.update_request_status(litellm_call_id=request_data.get("litellm_call_id", ""), status="fail")
|
||||
if AlertType.llm_exceptions in self.alert_types and not isinstance(
|
||||
original_exception, (HTTPException, ProxyException)
|
||||
):
|
||||
if AlertType.llm_exceptions in self.alert_types and not _is_client_error_exception(original_exception):
|
||||
"""
|
||||
Just alert on LLM API exceptions. Do not alert on user errors
|
||||
|
||||
|
|
|
|||
|
|
@ -1622,6 +1622,24 @@ class Router:
|
|||
return
|
||||
await selector.async_pre_call_check(deployment, parent_otel_span)
|
||||
|
||||
def _bind_override_selector_to_request(
|
||||
self, strategy: str, selector: RouterStrategySelector | None, request_kwargs: Mapping[str, object] | None
|
||||
) -> None:
|
||||
if selector is None or request_kwargs is None or strategy in self._globally_registered_strategies():
|
||||
return
|
||||
logging_obj: Final = request_kwargs.get("litellm_logging_obj")
|
||||
if isinstance(logging_obj, LiteLLMLogging):
|
||||
logging_obj.add_dynamic_callback(selector)
|
||||
|
||||
def _globally_registered_strategies(self) -> frozenset[str]:
|
||||
configured: Final = (
|
||||
self.routing_strategy,
|
||||
*(group.routing_strategy for group in self._routing_groups.values()),
|
||||
)
|
||||
return frozenset(
|
||||
normalized for normalized in map(self._normalize_strategy, configured) if normalized is not None
|
||||
)
|
||||
|
||||
def _get_routing_context(
|
||||
self, model: str, request_kwargs: dict | None = None
|
||||
) -> tuple[str | None, RouterStrategySelector | None]:
|
||||
|
|
@ -1647,7 +1665,9 @@ class Router:
|
|||
override: Final = self._get_request_routing_strategy_override(request_kwargs)
|
||||
if override is not None:
|
||||
verbose_router_logger.debug("routing_group=request-override model=%s strategy=%s", model, override)
|
||||
return override, self._get_override_strategy_selector(override)
|
||||
override_selector: Final = self._get_override_strategy_selector(override)
|
||||
self._bind_override_selector_to_request(override, override_selector, request_kwargs)
|
||||
return override, override_selector
|
||||
|
||||
group_name: Final = model if self.get_routing_group(model) is not None else self._model_to_group.get(model)
|
||||
if group_name is None:
|
||||
|
|
@ -2461,7 +2481,7 @@ class Router:
|
|||
|
||||
### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit)
|
||||
## only run if model group given, not model id
|
||||
if not self.has_model_id(model):
|
||||
if model in self.model_names or not self.has_model_id(model):
|
||||
self.routing_strategy_pre_call_checks(deployment=deployment)
|
||||
|
||||
input_kwargs: Final = {
|
||||
|
|
@ -12512,7 +12532,7 @@ class Router:
|
|||
# check if aliases set on litellm model alias map
|
||||
if specific_deployment is True:
|
||||
return model, self._get_deployment_by_litellm_model(model=model)
|
||||
elif self.has_model_id(model):
|
||||
elif model not in self.model_names and self.has_model_id(model):
|
||||
deployment: Final = self.get_deployment(model_id=model)
|
||||
if deployment is not None:
|
||||
deployment_model: Final = deployment.litellm_params.model
|
||||
|
|
|
|||
|
|
@ -68,6 +68,117 @@ still resolve to a deployment in `model_list`; this configuration does not creat
|
|||
- abc
|
||||
```
|
||||
|
||||
### Capability forecasting
|
||||
|
||||
Set `classifier_type: capability` to use
|
||||
[NVIDIA NeMo Switchyard's packaged capability classifier](https://github.com/NVIDIA-NeMo/Switchyard/blob/main/crates/libsy/src/prompts/capability-classifier/prompt.md).
|
||||
The classifier forecasts the probability that an efficient model completes
|
||||
the whole task, identifies the capability-card boundary that applies, and leaves the
|
||||
route choice to a deterministic threshold policy
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: smart-router
|
||||
litellm_params:
|
||||
model: auto_router/complexity_router
|
||||
complexity_router_config:
|
||||
classifier_type: capability
|
||||
classifier_llm_config:
|
||||
model: classifier-model
|
||||
capability_classifier_config:
|
||||
efficient_tier: SIMPLE
|
||||
capable_tier: REASONING
|
||||
base_threshold: 0.5
|
||||
threshold_step: 0.1
|
||||
tiers:
|
||||
SIMPLE:
|
||||
- efficient-model-a
|
||||
- efficient-model-b
|
||||
REASONING: capable-model
|
||||
```
|
||||
|
||||
The structured classifier verdict contains `crux`, `primary_rule`,
|
||||
`capability_boundary`, and `p_solve`. The policy computes the required solve
|
||||
probability as follows
|
||||
|
||||
- `supported`: `base_threshold`
|
||||
- `uncertain` or `unmatched`: `base_threshold + threshold_step`
|
||||
- `unsupported`: `base_threshold + 2 * threshold_step`
|
||||
|
||||
The efficient tier is selected when `p_solve` is greater than or equal to the
|
||||
adjusted threshold. Otherwise the capable tier is selected. A malformed,
|
||||
inconsistent, empty, or unavailable verdict always fails closed to the capable
|
||||
tier. `base_threshold` is required, `threshold_step` defaults to `0`, and their
|
||||
maximum adjusted threshold must not exceed `1`
|
||||
|
||||
The classifier receives the packaged Switchyard system prompt, the opening user
|
||||
task, and the latest user follow-up when present. Caller system messages,
|
||||
assistant turns, and intermediate tool results are not sent. The classifier call
|
||||
uses strict JSON Schema output and the existing classifier timeout, circuit
|
||||
breaker, attribution, redaction, reasoning-effort, and optional vision settings
|
||||
|
||||
`efficient_tier` and `capable_tier` name built-in complexity tiers with configured
|
||||
model pools. The forecast still makes one binary quality decision, while the
|
||||
ordinary tier pool may contain multiple equivalent deployments. Session affinity,
|
||||
keyword overrides, plan-mode floors, modality checks, and other post-classification
|
||||
complexity-router controls continue to apply
|
||||
|
||||
Routing decisions record the adjusted threshold and the complete valid forecast:
|
||||
`classifier_p_solve`, `classifier_capability_boundary`, `classifier_primary_rule`,
|
||||
and `classifier_crux`. Prompt redaction removes `classifier_crux` while retaining
|
||||
the derived fields needed to audit the decision
|
||||
|
||||
#### Calibrating solve probabilities
|
||||
|
||||
Supply a fitted monotone logit calibration under `capability_classifier_config`
|
||||
to transform the forecast before applying the threshold. Calibration is opt-in;
|
||||
without it the router uses the raw probability. Fit coefficients on benchmark
|
||||
outcomes from separate training repositories, select thresholds on a validation
|
||||
split, and report quality and cost on an untouched evaluation split
|
||||
|
||||
```yaml
|
||||
capability_classifier_config:
|
||||
efficient_tier: SIMPLE
|
||||
capable_tier: REASONING
|
||||
base_threshold: 0.66
|
||||
threshold_step: 0
|
||||
max_output_tokens: 512
|
||||
response_format: json_object
|
||||
calibration:
|
||||
version: your-benchmark-artifact-v1
|
||||
slope: 1.0
|
||||
intercept: 0.0
|
||||
```
|
||||
|
||||
The example coefficients are an identity mapping, not a trained calibration.
|
||||
The mapping is `sigmoid(slope * logit(clip(p_solve, 1e-6, 1-1e-6)) + intercept)`.
|
||||
The slope must be nonnegative, so calibration cannot improve ranking. It can
|
||||
make probabilities more accurate and thresholds easier to interpret. The version
|
||||
is recorded for auditing; the router does not check whether an artifact matches
|
||||
the judge, capability card, efficient solver, or agent harness. Operators must
|
||||
keep those aligned and refit when they change
|
||||
|
||||
Logs retain `classifier_p_solve` and add `classifier_calibrated_p_solve` and
|
||||
`classifier_calibration_version`. `classifier_threshold` is compared to the
|
||||
calibrated probability. Invalid verdicts still route to the capable tier
|
||||
|
||||
`response_format` defaults to `json_schema`. For endpoints that support JSON
|
||||
objects but not strict schemas, `json_object` appends the same schema to the
|
||||
unchanged capability prompt and retains strict local validation. Set
|
||||
`classifier_llm_config.timeout_ms` to cover the measured judge latency; a local
|
||||
judge may need longer than the default 3000 ms. `max_output_tokens` still defaults
|
||||
to 4096; 512 is an explicit benchmark setting for a short, non-reasoning judge
|
||||
|
||||
For a controlled whole-task benchmark, use `adaptive: false`,
|
||||
`session_affinity: true`, and a unique session ID for every task and policy arm.
|
||||
Disable keyword, plan-mode, housekeeping, and other optional overrides when
|
||||
measuring only the capability policy. When adaptive selection is enabled, it
|
||||
cannot select below the capability decision, including a capable-tier fallback
|
||||
|
||||
Configure capability forecasting through YAML or the model-management API.
|
||||
The dashboard preserves its classifier and calibration on an untouched save;
|
||||
it does not provide a capability-card editor
|
||||
|
||||
### Heuristic v2
|
||||
|
||||
Set `classifier_type: heuristic_v2` to classify with the bundled calibrated
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ from litellm.router_strategy.complexity_router.complexity_router import (
|
|||
from litellm.router_strategy.complexity_router.config import (
|
||||
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
|
||||
DEFAULT_COMPLEXITY_CONFIG,
|
||||
CapabilityCalibrationConfig,
|
||||
CapabilityClassifierConfig,
|
||||
ClassificationRubric,
|
||||
ComplexityRouterConfig,
|
||||
ComplexityTier,
|
||||
|
|
@ -28,6 +30,8 @@ from litellm.router_strategy.complexity_router.config import (
|
|||
__all__ = [
|
||||
"DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE",
|
||||
"DEFAULT_COMPLEXITY_CONFIG",
|
||||
"CapabilityCalibrationConfig",
|
||||
"CapabilityClassifierConfig",
|
||||
"ClassificationRubric",
|
||||
"ComplexityRouter",
|
||||
"ComplexityRouterConfig",
|
||||
|
|
|
|||
|
|
@ -0,0 +1,211 @@
|
|||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
"""Capability forecast contract and routing policy adapted from NVIDIA NeMo Switchyard."""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from sys import float_info
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, NamedTuple, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, StrictFloat, TypeAdapter, model_validator
|
||||
|
||||
CapabilityBoundary: TypeAlias = Literal["supported", "uncertain", "unsupported", "unmatched"]
|
||||
CapabilityRule: TypeAlias = Literal[
|
||||
"SUP-1",
|
||||
"SUP-2",
|
||||
"SUP-3",
|
||||
"SUP-4",
|
||||
"SUP-5",
|
||||
"UNC-1",
|
||||
"UNC-2",
|
||||
"LIM-1",
|
||||
"LIM-2",
|
||||
"none",
|
||||
]
|
||||
|
||||
CAPABILITY_CLASSIFIER_SYSTEM_PROMPT: Final = """You are a task-level probability forecaster for a model router. You receive the
|
||||
task's opening instruction and, when present, its latest user follow-up, plus
|
||||
the qualitative capability card below.
|
||||
|
||||
Forecast one binary event:
|
||||
|
||||
SUCCESS means that the efficient agent completes the whole task correctly on
|
||||
one fresh run under the actual harness, tools, and budget, as judged by the
|
||||
final verifier. FAILURE means any other outcome. The two outcomes are
|
||||
exhaustive.
|
||||
|
||||
Use only evidence in the instruction and the capability card. Do not assume
|
||||
hidden repository state, unmentioned tools, validators, documentation, access,
|
||||
or future work habits. Do not invent empirical counts, success rates, or base
|
||||
rates. The capability card is qualitative evidence, not a measured prior.
|
||||
|
||||
# Assessment procedure
|
||||
|
||||
1. State the crux: the hardest material requirement for whole-task success.
|
||||
2. Select the one capability rule that best describes the crux. Use
|
||||
primary_rule=none and capability_boundary=unmatched when no rule applies.
|
||||
Rule ids are opaque labels. Do not infer a boundary from an id's spelling.
|
||||
3. Privately identify the strongest instruction-visible reasons for SUCCESS
|
||||
and FAILURE, then imagine the most likely concrete failure.
|
||||
4. Privately consider material unknowns. Missing information should limit
|
||||
extreme estimates, but it is not evidence that p_solve must equal 0.50.
|
||||
5. Estimate p_solve last. It is the probability of whole-task SUCCESS, not
|
||||
confidence in this assessment, a route recommendation, or a cost judgment.
|
||||
|
||||
Interpret probabilities as natural frequencies. If p_solve is 0.70 for 100
|
||||
comparable fresh runs, about 70 should succeed and 30 should fail. Use the full
|
||||
range when justified. Reserve 0.00 and 1.00 for outcomes that are logically
|
||||
impossible or certain under the visible contract. Supported does not mean 1.00,
|
||||
and unsupported does not mean 0.00. The downstream routing threshold is not
|
||||
part of this forecast.
|
||||
|
||||
# Efficient-agent capability card
|
||||
|
||||
The route verbs in this source card are inherited qualitative descriptions.
|
||||
They do not ask you to output a route and do not assign a fixed probability to
|
||||
any boundary.
|
||||
|
||||
- SUP-1 [supported]: Route to the Efficient model when the task provides a complete output contract and a deterministic local validator that covers the material requirements.
|
||||
- SUP-2 [supported]: Route to the Efficient model when all required inputs are available, the target environment can be inspected, and correctness can be verified end-to-end without inaccessible external state.
|
||||
- SUP-3 [supported]: Route to the Efficient model when mathematical behavior, interfaces, shapes, data types, tolerances, and performance requirements are explicit and exercised by a representative harness.
|
||||
- SUP-4 [supported]: Route to the Efficient model when the required mechanism is identified, the relevant search space is bounded, and the success condition is executable. Do not infer this rule merely from the task's technical domain.
|
||||
- SUP-5 [supported]: Route to the Efficient model when reconstruction or behavioral reproduction is constrained by an executable reference, parser, format specification, or checker strong enough to distinguish correct from merely plausible output.
|
||||
- UNC-1 [uncertain]: Treat the route as uncertain when multiple reasonable interpretations of preprocessing, representation, indexing, naming, or output placement would produce different results and neither the instructions nor a validator resolve the choice.
|
||||
- UNC-2 [uncertain]: Treat the route as uncertain when success requires finding every relevant item across heterogeneous inputs or environment state, but the task does not define the search boundary or provide a completeness check.
|
||||
- LIM-1 [unsupported]: Prefer the Capable model when correctness depends primarily on extracting precise information from noisy visual, temporal, or rendered media and no machine-checkable extraction or replay mechanism is available.
|
||||
- LIM-2 [unsupported]: Prefer the Capable model when success depends on reproducing undocumented reference behavior, hidden intermediate state, or an unknown configuration, and small deviations fail despite satisfying the visible specification.
|
||||
|
||||
# Output
|
||||
|
||||
Return exactly one JSON object matching the response schema supplied with the
|
||||
request. Do not include markdown or commentary.
|
||||
|
||||
p_solve must be between 0.00 and 1.00. p_fail is exactly 1.00 - p_solve and
|
||||
must not be emitted separately. Do not output recommended_route, confidence,
|
||||
abstain, counts, task totals, empirical rates, or any other field."""
|
||||
|
||||
_BOUNDARY_STEPS: Final = MappingProxyType(
|
||||
{
|
||||
"supported": 0,
|
||||
"uncertain": 1,
|
||||
"unmatched": 1,
|
||||
"unsupported": 2,
|
||||
}
|
||||
)
|
||||
|
||||
_RULE_BOUNDARIES: Final = MappingProxyType(
|
||||
{
|
||||
"SUP-1": "supported",
|
||||
"SUP-2": "supported",
|
||||
"SUP-3": "supported",
|
||||
"SUP-4": "supported",
|
||||
"SUP-5": "supported",
|
||||
"UNC-1": "uncertain",
|
||||
"UNC-2": "uncertain",
|
||||
"LIM-1": "unsupported",
|
||||
"LIM-2": "unsupported",
|
||||
"none": "unmatched",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class CapabilityClassifierVerdict(BaseModel):
|
||||
"""Strict structured verdict returned by the capability forecaster."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
crux: str = Field(min_length=1)
|
||||
primary_rule: CapabilityRule
|
||||
capability_boundary: CapabilityBoundary
|
||||
p_solve: StrictFloat = Field(ge=0.0, le=1.0)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_rule_boundary_pair(self) -> "CapabilityClassifierVerdict":
|
||||
if not self.crux.strip():
|
||||
raise ValueError("crux must contain non-whitespace text")
|
||||
expected: Final = _RULE_BOUNDARIES[self.primary_rule]
|
||||
if self.capability_boundary != expected:
|
||||
raise ValueError(
|
||||
f"primary_rule {self.primary_rule!r} requires capability_boundary {expected!r}, "
|
||||
f"got {self.capability_boundary!r}"
|
||||
)
|
||||
return self
|
||||
|
||||
def routing_threshold(self, base_threshold: float, threshold_step: float) -> float:
|
||||
"""Required efficient-model solve probability for this boundary."""
|
||||
return base_threshold + _BOUNDARY_STEPS[self.capability_boundary] * threshold_step
|
||||
|
||||
def meets_routing_threshold(self, threshold: float) -> bool:
|
||||
"""Inclusive comparison with Switchyard's one-epsilon rounding guard."""
|
||||
return self.p_solve >= threshold or abs(threshold - self.p_solve) <= float_info.epsilon
|
||||
|
||||
|
||||
class CapabilityClassifierForecast(NamedTuple):
|
||||
verdict: CapabilityClassifierVerdict
|
||||
threshold: float
|
||||
p_solve: float
|
||||
calibration_version: str | None
|
||||
|
||||
def meets_routing_threshold(self) -> bool:
|
||||
return self.p_solve >= self.threshold or abs(self.threshold - self.p_solve) <= float_info.epsilon
|
||||
|
||||
|
||||
_CAPABILITY_CLASSIFIER_RESPONSE_FORMAT_JSON: Final = """{
|
||||
"type": "json_schema",
|
||||
"json_schema": {
|
||||
"name": "CapabilityClassifierDecision",
|
||||
"strict": true,
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"additionalProperties": false,
|
||||
"required": ["crux", "primary_rule", "capability_boundary", "p_solve"],
|
||||
"properties": {
|
||||
"crux": {"type": "string", "minLength": 1},
|
||||
"primary_rule": {
|
||||
"type": "string",
|
||||
"enum": ["SUP-1", "SUP-2", "SUP-3", "SUP-4", "SUP-5", "UNC-1", "UNC-2", "LIM-1", "LIM-2", "none"]
|
||||
},
|
||||
"capability_boundary": {
|
||||
"type": "string",
|
||||
"enum": ["supported", "uncertain", "unsupported", "unmatched"]
|
||||
},
|
||||
"p_solve": {"type": "number", "minimum": 0.0, "maximum": 1.0}
|
||||
}
|
||||
}
|
||||
}
|
||||
}"""
|
||||
|
||||
_RESPONSE_FORMAT_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def capability_classifier_response_format(
|
||||
mode: Literal["json_schema", "json_object"] = "json_schema",
|
||||
) -> Mapping[str, object]:
|
||||
"""Fresh copy of Switchyard's packaged strict JSON Schema wrapper."""
|
||||
return (
|
||||
_RESPONSE_FORMAT_ADAPTER.validate_json('{"type": "json_object"}')
|
||||
if mode == "json_object"
|
||||
else _RESPONSE_FORMAT_ADAPTER.validate_json(_CAPABILITY_CLASSIFIER_RESPONSE_FORMAT_JSON)
|
||||
)
|
||||
|
||||
|
||||
def capability_classifier_system_prompt(mode: Literal["json_schema", "json_object"]) -> str:
|
||||
if mode == "json_schema":
|
||||
return CAPABILITY_CLASSIFIER_SYSTEM_PROMPT
|
||||
wrapper: Final = _RESPONSE_FORMAT_ADAPTER.validate_python(capability_classifier_response_format()["json_schema"])
|
||||
return (
|
||||
CAPABILITY_CLASSIFIER_SYSTEM_PROMPT
|
||||
+ "\n\nReturn exactly one JSON object matching this JSON Schema:\n"
|
||||
+ json.dumps(wrapper["schema"], indent=2, sort_keys=True)
|
||||
)
|
||||
|
||||
|
||||
def parse_capability_classifier_verdict(content: str) -> CapabilityClassifierVerdict:
|
||||
"""Parse raw JSON or the fenced JSON shape tolerated by Switchyard."""
|
||||
text: Final = content.strip()
|
||||
if not text.startswith("```"):
|
||||
return CapabilityClassifierVerdict.model_validate_json(text)
|
||||
unfenced: Final = text.removeprefix("```").removeprefix("json").lstrip("\n\r")
|
||||
return CapabilityClassifierVerdict.model_validate_json(unfenced.removesuffix("```").strip())
|
||||
|
|
@ -5,8 +5,9 @@ A rule-based routing strategy that uses weighted scoring across multiple dimensi
|
|||
to classify requests by complexity and route them to appropriate models.
|
||||
|
||||
By default, scoring is local (regex/keyword-based) with no external API calls and <1ms
|
||||
latency. Optionally, classifier_type="llm" routes classification through a configured
|
||||
model instead, trading that latency/cost guarantee for potentially better accuracy.
|
||||
latency. Optionally, classifier_type="llm" selects a tier through a configured model,
|
||||
while classifier_type="capability" forecasts efficient-model success and applies a
|
||||
Switchyard-compatible threshold policy.
|
||||
keyword_tier_rules (lexical or, with semantic_keyword_matching, embedding-based) are
|
||||
evaluated before either classification strategy and force a tier outright when matched.
|
||||
|
||||
|
|
@ -75,6 +76,12 @@ from litellm.types.utils import (
|
|||
StandardLoggingRoutingDecisionTierBoundaries,
|
||||
)
|
||||
|
||||
from .capability_classifier import (
|
||||
CapabilityClassifierForecast,
|
||||
capability_classifier_response_format,
|
||||
capability_classifier_system_prompt,
|
||||
parse_capability_classifier_verdict,
|
||||
)
|
||||
from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section
|
||||
from .config import (
|
||||
CALIBRATION_EXAMPLES_HEADING,
|
||||
|
|
@ -997,6 +1004,7 @@ class ClassificationOutcome(NamedTuple):
|
|||
"heuristic_v2",
|
||||
"reasoning_override",
|
||||
"llm_classifier",
|
||||
"capability_classifier",
|
||||
"llm_v2_classifier",
|
||||
"llm_v2_fallback",
|
||||
"heuristic_first_short_circuit",
|
||||
|
|
@ -1004,15 +1012,43 @@ class ClassificationOutcome(NamedTuple):
|
|||
"housekeeping",
|
||||
"classifier_plugin",
|
||||
"classifier_fallback",
|
||||
"capability_classifier_fallback",
|
||||
"default_model_fallback",
|
||||
]
|
||||
classifier_cost: float | None = None
|
||||
capability_forecast: CapabilityClassifierForecast | None = None
|
||||
|
||||
|
||||
def _with_signal(outcome: ClassificationOutcome, signal: str | None) -> ClassificationOutcome:
|
||||
return outcome if signal is None else outcome._replace(signals=(*outcome.signals, signal))
|
||||
|
||||
|
||||
def _with_capability_forecast(
|
||||
decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome
|
||||
) -> StandardLoggingRoutingDecision:
|
||||
"""Attach the validated capability verdict and applied threshold to its decision record."""
|
||||
forecast: Final = outcome.capability_forecast
|
||||
if forecast is None:
|
||||
return decision
|
||||
verdict: Final = forecast.verdict
|
||||
enriched: Final[StandardLoggingRoutingDecision] = { # mutable-ok: routing decisions are JSON TypedDict records
|
||||
**decision,
|
||||
"classifier_crux": verdict.crux,
|
||||
"classifier_primary_rule": verdict.primary_rule,
|
||||
"classifier_capability_boundary": verdict.capability_boundary,
|
||||
"classifier_p_solve": verdict.p_solve,
|
||||
"classifier_threshold": forecast.threshold,
|
||||
}
|
||||
if forecast.calibration_version is None:
|
||||
return enriched
|
||||
calibrated: Final[StandardLoggingRoutingDecision] = {
|
||||
**enriched,
|
||||
"classifier_calibrated_p_solve": forecast.p_solve,
|
||||
"classifier_calibration_version": forecast.calibration_version,
|
||||
}
|
||||
return calibrated
|
||||
|
||||
|
||||
class _ClassifierCircuitBreaker:
|
||||
"""Process-local timeout breaker for one complexity-router classifier.
|
||||
|
||||
|
|
@ -1281,10 +1317,15 @@ class ComplexityRouter(CustomLogger):
|
|||
self._classifier_system_prompt: str | None = (
|
||||
self._build_classifier_system_prompt() if llm_classifier_configured else None
|
||||
)
|
||||
capability_config: Final = self.config.capability_classifier_config
|
||||
self._classifier_response_format: Mapping[str, object] | None = (
|
||||
(
|
||||
llm_v2_response_format(self.config.llm_v2_config.response_format)
|
||||
if self.config.llm_v2_config is not None
|
||||
else capability_classifier_response_format(
|
||||
capability_config.response_format if capability_config is not None else "json_schema"
|
||||
)
|
||||
if self.config.classifier_type == "capability"
|
||||
else type_to_response_format_param(_tier_classification_model(self.config.classifier_wire_labels()))
|
||||
)
|
||||
if llm_classifier_configured
|
||||
|
|
@ -1316,6 +1357,11 @@ class ComplexityRouter(CustomLogger):
|
|||
if v2 is not None:
|
||||
pools: Final = self._tier_pools()
|
||||
return v2.system_prompt(pools[v2.efficient_tier][0], pools[v2.capable_tier][0])
|
||||
if self.config.classifier_type == "capability":
|
||||
capability: Final = self.config.capability_classifier_config
|
||||
return capability_classifier_system_prompt(
|
||||
capability.response_format if capability is not None else "json_schema"
|
||||
)
|
||||
definitions: Final = self.config.tier_definitions
|
||||
if definitions is not None:
|
||||
return custom_tier_classification_prompt(
|
||||
|
|
@ -1733,6 +1779,8 @@ class ComplexityRouter(CustomLogger):
|
|||
return await self._classify_heuristic_first(prompt, system_prompt, request_kwargs, messages)
|
||||
if self.config.classifier_type == "hybrid" and self.config.classifier_llm_config is not None:
|
||||
return await self._classify_hybrid(prompt, system_prompt, request_kwargs, messages)
|
||||
if self.config.classifier_type == "capability" and self.config.classifier_llm_config is not None:
|
||||
return await self._capability_classifier_outcome(prompt, request_kwargs, messages)
|
||||
if self.config.classifier_type not in ("llm", "llm_v2") or self.config.classifier_llm_config is None:
|
||||
tier, score, signals, cause = self._score_and_classify(prompt, system_prompt)
|
||||
return ClassificationOutcome(tier=tier, score=score, signals=signals, cause=cause)
|
||||
|
|
@ -1844,6 +1892,66 @@ class ComplexityRouter(CustomLogger):
|
|||
)
|
||||
)
|
||||
|
||||
async def _capability_classifier_outcome(
|
||||
self,
|
||||
prompt: str,
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
) -> ClassificationOutcome:
|
||||
"""Forecast efficient-tier success, then apply the deterministic boundary policy."""
|
||||
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._capability_classifier_failure_outcome(
|
||||
"capability classifier circuit is open", signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL
|
||||
)
|
||||
try:
|
||||
tier, classifier_cost, forecast = await self._classify_with_capability_llm(prompt, request_kwargs, messages)
|
||||
if breaker is not None and permit is not None:
|
||||
breaker.record_success(permit)
|
||||
return ClassificationOutcome(
|
||||
tier=tier,
|
||||
score=None,
|
||||
signals=(
|
||||
f"capability-boundary:{forecast.verdict.capability_boundary}",
|
||||
f"capability-rule:{forecast.verdict.primary_rule}",
|
||||
),
|
||||
cause="capability_classifier",
|
||||
classifier_cost=classifier_cost,
|
||||
capability_forecast=forecast,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
if breaker is not None and permit is not None:
|
||||
breaker.record_failure(permit, is_timeout=False)
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 -- every unavailable or invalid judge verdict must fail closed
|
||||
if breaker is not None and permit is not None:
|
||||
breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e))
|
||||
return self._capability_classifier_failure_outcome(f"capability classifier failed ({e})")
|
||||
|
||||
def _capability_classifier_failure_outcome(self, reason: str, signal: str | None = None) -> ClassificationOutcome:
|
||||
"""Fail closed to the configured capable tier without consulting another taxonomy."""
|
||||
capability: Final = self.config.capability_classifier_config
|
||||
if capability is None:
|
||||
raise ValueError("capability_classifier_config is not set")
|
||||
verbose_router_logger.warning(
|
||||
"ComplexityRouter: %s, routing to capable_tier %s", reason, capability.capable_tier
|
||||
)
|
||||
signals: Final = (
|
||||
("capability-classifier-fallback",)
|
||||
if signal is None
|
||||
else (
|
||||
"capability-classifier-fallback",
|
||||
signal,
|
||||
)
|
||||
)
|
||||
return ClassificationOutcome(
|
||||
tier=ComplexityTier(capability.capable_tier),
|
||||
score=None,
|
||||
signals=signals,
|
||||
cause="capability_classifier_fallback",
|
||||
)
|
||||
|
||||
async def _llm_classifier_outcome(
|
||||
self,
|
||||
prompt: str,
|
||||
|
|
@ -2172,6 +2280,73 @@ class ComplexityRouter(CustomLogger):
|
|||
classifier_cost=classifier_cost,
|
||||
)
|
||||
|
||||
async def _classify_with_capability_llm(
|
||||
self,
|
||||
prompt: str,
|
||||
request_kwargs: Mapping[str, object] | None,
|
||||
messages: Sequence[Mapping[str, object]] | None,
|
||||
) -> tuple[ComplexityTier, float | None, CapabilityClassifierForecast]:
|
||||
"""Call the packaged capability forecaster and apply its two-tier policy."""
|
||||
capability: Final = self.config.capability_classifier_config
|
||||
classifier_system_prompt: Final = self._classifier_system_prompt
|
||||
if capability is None or classifier_system_prompt is None:
|
||||
raise ValueError("capability classifier is not configured")
|
||||
|
||||
markers: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING)
|
||||
encrypted_task: Final = _encrypted_classifier_task(request_kwargs, markers)
|
||||
asks_newest_first: Final = (
|
||||
() if encrypted_task is not None else tuple(_iter_human_asks_newest_first(messages or (), markers))
|
||||
)
|
||||
opening_task: Final = (
|
||||
"The delegated task in the following agent_message."
|
||||
if encrypted_task is not None
|
||||
else asks_newest_first[-1]
|
||||
if asks_newest_first
|
||||
else prompt
|
||||
)
|
||||
latest_follow_up: Final = asks_newest_first[0] if len(asks_newest_first) > 1 else None
|
||||
task_messages: list[AllMessageValues] = [ # mutable-ok: the latest message gains optional image parts below
|
||||
{"role": "user", "content": opening_task}, # mutable-ok: SDK messages are dict-shaped
|
||||
]
|
||||
if latest_follow_up is not None:
|
||||
task_messages.append( # mutable-ok: the provider SDK requires a concrete message list
|
||||
{"role": "user", "content": latest_follow_up} # mutable-ok: SDK messages are dict-shaped
|
||||
)
|
||||
|
||||
image_parts: Final = self._classifier_image_parts(messages)
|
||||
if image_parts:
|
||||
latest_text: Final = latest_follow_up or opening_task
|
||||
task_messages[-1] = { # mutable-ok: SDK messages are dict-shaped
|
||||
"role": "user",
|
||||
"content": [ # mutable-ok: multimodal SDK content is a JSON array
|
||||
{"type": "text", "text": latest_text}, # mutable-ok: SDK content parts are dict-shaped
|
||||
*image_parts,
|
||||
],
|
||||
}
|
||||
messages_for_call: Final[list[AllMessageValues]] = [ # mutable-ok: provider SDK requires a concrete list
|
||||
{"role": "system", "content": classifier_system_prompt}, # mutable-ok: SDK messages are dict-shaped
|
||||
*task_messages,
|
||||
]
|
||||
content, classifier_cost = await self._call_classifier_model(
|
||||
messages_for_call,
|
||||
request_kwargs,
|
||||
max_output_tokens=capability.max_output_tokens,
|
||||
encrypted_task=encrypted_task,
|
||||
)
|
||||
verdict: Final = parse_capability_classifier_verdict(content)
|
||||
threshold: Final = verdict.routing_threshold(capability.base_threshold, capability.threshold_step)
|
||||
calibration: Final = capability.calibration
|
||||
forecast: Final = CapabilityClassifierForecast(
|
||||
verdict=verdict,
|
||||
threshold=threshold,
|
||||
p_solve=calibration.calibrate(verdict.p_solve) if calibration is not None else verdict.p_solve,
|
||||
calibration_version=calibration.version if calibration is not None else None,
|
||||
)
|
||||
selected_tier: Final = (
|
||||
capability.efficient_tier if forecast.meets_routing_threshold() else capability.capable_tier
|
||||
)
|
||||
return ComplexityTier(selected_tier), classifier_cost, forecast
|
||||
|
||||
async def _call_classifier_model(
|
||||
self,
|
||||
messages_for_call: list[AllMessageValues], # mutable-ok: SDKs require a list
|
||||
|
|
@ -4197,7 +4372,12 @@ class ComplexityRouter(CustomLogger):
|
|||
housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None
|
||||
# A context-escalated tier becomes the hard floor: a floor the bandit can slide
|
||||
# under is not a floor.
|
||||
adaptive_floor: Final = tier if context_original_tier is not None else plan_floor
|
||||
adaptive_floor: Final = (
|
||||
tier
|
||||
if context_original_tier is not None
|
||||
or outcome.cause in ("capability_classifier", "capability_classifier_fallback")
|
||||
else plan_floor
|
||||
)
|
||||
adaptive_fit: Final = context_placement.holdable_models if context_placement is not None else None
|
||||
sampled_model: Final = self._soft_floor_pick(
|
||||
tier,
|
||||
|
|
@ -4247,7 +4427,7 @@ class ComplexityRouter(CustomLogger):
|
|||
tier_litellm_params: Final = self._litellm_params_for_model(tier, routed_model)
|
||||
classifier_model: Final = (
|
||||
self.config.classifier_llm_config.model
|
||||
if outcome.cause in ("llm_classifier", "llm_v2_classifier", "llm_v2_fallback")
|
||||
if outcome.cause in ("llm_classifier", "capability_classifier", "llm_v2_classifier", "llm_v2_fallback")
|
||||
and self.config.classifier_llm_config is not None
|
||||
else None
|
||||
)
|
||||
|
|
@ -4271,23 +4451,24 @@ class ComplexityRouter(CustomLogger):
|
|||
decision_keyword: Final = (
|
||||
plan_mode_sentinel if plan_floored else (housekeeping_sentinel if outcome.cause == "housekeeping" else None)
|
||||
)
|
||||
routing_decision: Final = self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
conversation_continuing=conversation_continuing,
|
||||
cause=decision_cause,
|
||||
tier=classified_pool_tier,
|
||||
score=score,
|
||||
signals=decision_signals,
|
||||
matched_keyword=decision_keyword,
|
||||
escalation_keyword=escalation_keyword,
|
||||
escalated=escalated,
|
||||
classifier_model=classifier_model,
|
||||
classifier_cost=outcome.classifier_cost,
|
||||
tier_litellm_params=tier_litellm_params,
|
||||
context_escalation_original_tier=context_original_tier,
|
||||
)
|
||||
return PreRoutingHookResponse(
|
||||
model=routed_model,
|
||||
messages=messages if has_original_messages else None,
|
||||
litellm_params=tier_litellm_params,
|
||||
routing_decision=self._build_routing_decision(
|
||||
routed_model=routed_model,
|
||||
conversation_continuing=conversation_continuing,
|
||||
cause=decision_cause,
|
||||
tier=classified_pool_tier,
|
||||
score=score,
|
||||
signals=decision_signals,
|
||||
matched_keyword=decision_keyword,
|
||||
escalation_keyword=escalation_keyword,
|
||||
escalated=escalated,
|
||||
classifier_model=classifier_model,
|
||||
classifier_cost=outcome.classifier_cost,
|
||||
tier_litellm_params=tier_litellm_params,
|
||||
context_escalation_original_tier=context_original_tier,
|
||||
),
|
||||
routing_decision=_with_capability_forecast(routing_decision, outcome),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -13,7 +13,16 @@ from enum import Enum
|
|||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, Literal, NamedTuple
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, SkipValidation, field_serializer, field_validator, model_validator
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
SkipValidation,
|
||||
StrictFloat,
|
||||
field_serializer,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore", DeprecationWarning)
|
||||
|
|
@ -54,7 +63,7 @@ DEFAULT_CLASSIFICATION_RUBRIC: Final[ClassificationRubric] = ClassificationRubri
|
|||
# The classifier_type values that can call classifier_llm_config.model. Every consumer asking
|
||||
# "is the classifier model a real dependency of this router" resolves it here, including the ones
|
||||
# that only hold the raw config mapping and cannot reach ComplexityRouterConfig.uses_llm_classifier.
|
||||
LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "llm_v2", "heuristic_first", "hybrid"})
|
||||
LLM_CLASSIFIER_TYPES: Final[frozenset[str]] = frozenset({"llm", "capability", "llm_v2", "heuristic_first", "hybrid"})
|
||||
|
||||
|
||||
TIER_SEVERITY_ORDER: Final[tuple[ComplexityTier, ...]] = (
|
||||
|
|
@ -592,6 +601,78 @@ class ClassifierLLMConfig(BaseModel):
|
|||
return self
|
||||
|
||||
|
||||
class CapabilityCalibrationConfig(BaseModel):
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
version: str = Field(min_length=1, max_length=128, pattern=r"^\S(?:.*\S)?$")
|
||||
slope: StrictFloat = Field(ge=0.0, le=20.0, allow_inf_nan=False)
|
||||
intercept: StrictFloat = Field(ge=-20.0, le=20.0, allow_inf_nan=False)
|
||||
|
||||
def calibrate(self, p_solve: float) -> float:
|
||||
clipped: Final = min(max(p_solve, 1e-6), 1.0 - 1e-6)
|
||||
log_odds: Final = self.slope * (math.log(clipped) - math.log1p(-clipped)) + self.intercept
|
||||
return 1.0 / (1.0 + math.exp(-log_odds))
|
||||
|
||||
|
||||
class CapabilityClassifierConfig(BaseModel):
|
||||
"""Switchyard-compatible probability threshold policy for two model tiers."""
|
||||
|
||||
model_config = ConfigDict(extra="forbid", frozen=True)
|
||||
|
||||
efficient_tier: str = Field(
|
||||
description="Tier used when the efficient model's forecasted solve probability meets the adjusted threshold",
|
||||
)
|
||||
capable_tier: str = Field(
|
||||
description=(
|
||||
"Higher, fail-closed tier used below the adjusted threshold or when the classifier verdict is unavailable"
|
||||
),
|
||||
)
|
||||
base_threshold: StrictFloat = Field(
|
||||
ge=0.0,
|
||||
le=1.0,
|
||||
description="Lowest p_solve that routes a supported task to efficient_tier",
|
||||
)
|
||||
threshold_step: StrictFloat = Field(
|
||||
default=0.0,
|
||||
ge=0.0,
|
||||
description=("Amount added once for uncertain or unmatched verdicts and twice for unsupported verdicts"),
|
||||
)
|
||||
max_output_tokens: int = Field(
|
||||
default=4096,
|
||||
ge=1,
|
||||
description="Maximum completion tokens available to the capability classifier verdict",
|
||||
)
|
||||
calibration: CapabilityCalibrationConfig | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Optional versioned sigmoid calibration fitted for this judge, capability card, efficient model, "
|
||||
"and execution setup. Applies sigmoid(slope * logit(clip(p_solve, 1e-6, 1-1e-6)) + intercept) "
|
||||
"before the threshold policy. Omit to route on the raw forecast."
|
||||
),
|
||||
)
|
||||
response_format: Literal["json_schema", "json_object"] = Field(
|
||||
default="json_schema",
|
||||
description=(
|
||||
"Use json_object for judges without strict JSON Schema support. This appends the verdict schema "
|
||||
"to the packaged system prompt; both modes validate the returned verdict identically."
|
||||
),
|
||||
)
|
||||
|
||||
@field_validator("efficient_tier", "capable_tier")
|
||||
@classmethod
|
||||
def _normalize_tier(cls, value: str) -> str:
|
||||
normalized: Final = value.strip()
|
||||
if not normalized:
|
||||
raise ValueError("tier must be non-empty")
|
||||
return normalized
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_threshold_range(self) -> "CapabilityClassifierConfig":
|
||||
if self.base_threshold + 2 * self.threshold_step > 1.0:
|
||||
raise ValueError("base_threshold + 2 * threshold_step must be at most 1")
|
||||
return self
|
||||
|
||||
|
||||
MAX_CUSTOM_PATTERN_REPEAT: Final[int] = 64
|
||||
MAX_CUSTOM_PATTERN_WORK: Final[int] = 2048
|
||||
MAX_CUSTOM_DIMENSIONS_WORK: Final[int] = 8192
|
||||
|
|
@ -883,16 +964,18 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
|
||||
# Classifier strategy
|
||||
classifier_type: Literal["heuristic", "heuristic_v2", "llm", "llm_v2", "custom", "heuristic_first", "hybrid"] = (
|
||||
Field(
|
||||
default="heuristic",
|
||||
description=(
|
||||
"Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, "
|
||||
"an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays "
|
||||
"for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', "
|
||||
"which trusts the local scorer everywhere except when its score lands near a tier boundary"
|
||||
),
|
||||
)
|
||||
classifier_type: Literal[
|
||||
"heuristic", "heuristic_v2", "llm", "capability", "llm_v2", "custom", "heuristic_first", "hybrid"
|
||||
] = Field(
|
||||
default="heuristic",
|
||||
description=(
|
||||
"Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, "
|
||||
"an LLM tier-selection call, a Switchyard-compatible capability forecast, a joint task-demand and "
|
||||
"capability forecast, a custom classifier plugin, 'heuristic_first', which scores locally and only "
|
||||
"pays for the LLM classifier when the "
|
||||
"local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer "
|
||||
"everywhere except when its score lands near a tier boundary"
|
||||
),
|
||||
)
|
||||
llm_v2_config: LLMV2Config | None = Field(
|
||||
default=None,
|
||||
|
|
@ -909,7 +992,15 @@ class ComplexityRouterConfig(BaseModel):
|
|||
default=None,
|
||||
description=(
|
||||
"Configuration for the LLM classifier; required when classifier_type is 'llm', "
|
||||
"'heuristic_first' or 'hybrid'"
|
||||
"'capability', 'heuristic_first' or 'hybrid'"
|
||||
),
|
||||
)
|
||||
capability_classifier_config: CapabilityClassifierConfig | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"Probability threshold policy required when classifier_type is 'capability'. The classifier "
|
||||
"forecasts p_solve for efficient_tier, adjusts base_threshold using the capability-card boundary, "
|
||||
"and otherwise routes to capable_tier"
|
||||
),
|
||||
)
|
||||
heuristic_first_max_tier: str | None = Field(
|
||||
|
|
@ -1434,6 +1525,66 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_capability_classifier_config(self) -> "ComplexityRouterConfig":
|
||||
capability: Final = self.capability_classifier_config
|
||||
if self.classifier_type != "capability":
|
||||
if capability is not None:
|
||||
raise ValueError(
|
||||
"capability_classifier_config requires classifier_type 'capability'; otherwise it has no effect"
|
||||
)
|
||||
return self
|
||||
if capability is None:
|
||||
raise ValueError("capability_classifier_config is required when classifier_type is 'capability'")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_capability_classifier_tiers(self) -> "ComplexityRouterConfig":
|
||||
capability: Final = self.capability_classifier_config
|
||||
if self.classifier_type != "capability" or capability is None:
|
||||
return self
|
||||
if self.tier_definitions is not None:
|
||||
raise ValueError(
|
||||
"classifier_type 'capability' uses the built-in tier map and cannot be combined with tier_definitions"
|
||||
)
|
||||
for field, tier in (
|
||||
("efficient_tier", capability.efficient_tier),
|
||||
("capable_tier", capability.capable_tier),
|
||||
):
|
||||
if tier not in self.tier_names():
|
||||
raise ValueError(
|
||||
f"{field} {tier!r} is not an active tier: it must name one of {', '.join(self.tier_names())}"
|
||||
)
|
||||
if not self.tiers.get(tier):
|
||||
raise ValueError(f"{field} {tier!r} has no model configured in tiers")
|
||||
names: Final = self.tier_names()
|
||||
if names.index(capability.capable_tier) <= names.index(capability.efficient_tier):
|
||||
raise ValueError("capable_tier must be a higher tier than efficient_tier")
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_capability_classifier_prompt_policy(self) -> "ComplexityRouterConfig":
|
||||
if self.classifier_type != "capability":
|
||||
return self
|
||||
llm_config: Final = self.classifier_llm_config
|
||||
if llm_config is not None and (
|
||||
llm_config.system_prompt is not None or llm_config.classification_rubric is not None
|
||||
):
|
||||
raise ValueError(
|
||||
"classifier_type 'capability' uses the packaged capability card; classifier_llm_config.system_prompt "
|
||||
"and classification_rubric are not supported"
|
||||
)
|
||||
if self.classification_prompt is not None or self.classification_examples is not None:
|
||||
raise ValueError(
|
||||
"classifier_type 'capability' uses the packaged capability card; classification_prompt and "
|
||||
"classification_examples are not supported"
|
||||
)
|
||||
if self.classifier_fallback != "heuristic":
|
||||
raise ValueError(
|
||||
"classifier_type 'capability' always fails closed to capable_tier; classifier_fallback cannot override it"
|
||||
)
|
||||
return self
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _validate_llm_v2(self) -> "ComplexityRouterConfig":
|
||||
v2: Final = self.llm_v2_config
|
||||
|
|
@ -1731,7 +1882,7 @@ class ComplexityRouterConfig(BaseModel):
|
|||
)
|
||||
if duplicated:
|
||||
raise ValueError(f"tier_definitions names must be unique (case-insensitive): {', '.join(duplicated)}")
|
||||
if self.classifier_type in ("heuristic", "heuristic_v2", "heuristic_first", "hybrid"):
|
||||
if self.classifier_type in ("heuristic", "heuristic_v2", "capability", "heuristic_first", "hybrid"):
|
||||
raise ValueError(
|
||||
"tier_definitions requires classifier_type 'llm' or 'custom': the heuristic scorer only "
|
||||
"produces the built-in tiers from SIMPLE up, as does heuristic_v2"
|
||||
|
|
|
|||
|
|
@ -113,7 +113,7 @@ def strategy_router_dependencies(
|
|||
"""The model names a strategy-router deployment must reach, in no particular order.
|
||||
|
||||
A field is a dependency only under the condition the runtime itself reads it: the
|
||||
classifier model needs `classifier_type: llm`, and the complexity embedding model needs
|
||||
classifier model needs an LLM-backed classifier type, and the complexity embedding model needs
|
||||
`semantic_keyword_matching`. Listing one the router never calls reds a working deployment.
|
||||
|
||||
The two default-model spellings are not symmetric. A quality router falls back to its
|
||||
|
|
|
|||
|
|
@ -2891,6 +2891,7 @@ RoutingDecisionCause = Literal[
|
|||
# meant anything that filtered `signals` silently changed what the row claimed.
|
||||
"reasoning_override",
|
||||
"llm_classifier",
|
||||
"capability_classifier",
|
||||
"llm_v2_classifier",
|
||||
"llm_v2_fallback",
|
||||
# classifier_type 'heuristic_first': the local scorer produced at least one signal and landed at
|
||||
|
|
@ -2905,6 +2906,9 @@ RoutingDecisionCause = Literal[
|
|||
# The LLM classifier or classifier plugin failed on a router with an operator-defined
|
||||
# tier set, so the request routed to the configured fallback_tier without being classified.
|
||||
"classifier_fallback",
|
||||
# The capability judge failed or returned an invalid verdict, so its fail-closed policy
|
||||
# routed to capable_tier without consulting the unrelated complexity heuristic.
|
||||
"capability_classifier_fallback",
|
||||
# The LLM classifier or classifier plugin failed and classifier_fallback is
|
||||
# 'default_model', so the request went to default_model without being classified.
|
||||
# Distinct from "default_fallback",
|
||||
|
|
@ -2980,6 +2984,13 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
|
|||
escalation_keyword: str
|
||||
classifier_model: str
|
||||
classifier_cost: float
|
||||
classifier_crux: str # writable-ok: added only when a capability verdict is available
|
||||
classifier_primary_rule: str # writable-ok: added only when a capability verdict is available
|
||||
classifier_capability_boundary: str # writable-ok: added only when a capability verdict is available
|
||||
classifier_p_solve: float # writable-ok: added only when a capability verdict is available
|
||||
classifier_calibrated_p_solve: ReadOnly[float]
|
||||
classifier_calibration_version: ReadOnly[str]
|
||||
classifier_threshold: float # writable-ok: added only when a capability verdict is available
|
||||
escalated: bool
|
||||
context_escalated: bool # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
context_escalation_original_tier: str # writable-ok: Pydantic warns on ReadOnly TypedDict fields
|
||||
|
|
@ -2995,7 +3006,9 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
|
|||
# logging off. Every other field aggregates the prompt without reproducing it and is kept,
|
||||
# so a redacted row stays explainable. `test_every_routing_decision_field_is_classified`
|
||||
# fails if a field is added to the record without being placed in one set or the other.
|
||||
PROMPT_QUOTING_ROUTING_DECISION_FIELDS: frozenset[str] = frozenset({"signals", "matched_keyword", "escalation_keyword"})
|
||||
PROMPT_QUOTING_ROUTING_DECISION_FIELDS: frozenset[str] = frozenset(
|
||||
{"signals", "matched_keyword", "escalation_keyword", "classifier_crux"}
|
||||
)
|
||||
DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
"router_model_name",
|
||||
|
|
@ -3008,6 +3021,12 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset(
|
|||
"score",
|
||||
"classifier_model",
|
||||
"classifier_cost",
|
||||
"classifier_primary_rule",
|
||||
"classifier_capability_boundary",
|
||||
"classifier_p_solve",
|
||||
"classifier_calibrated_p_solve",
|
||||
"classifier_calibration_version",
|
||||
"classifier_threshold",
|
||||
"escalated",
|
||||
"context_escalated",
|
||||
"context_escalation_original_tier",
|
||||
|
|
|
|||
|
|
@ -2202,15 +2202,20 @@ def _is_streaming_request(
|
|||
|
||||
def _select_tokenizer(model: str, custom_tokenizer: CustomHuggingfaceTokenizer | None = None):
|
||||
if custom_tokenizer is not None:
|
||||
_tokenizer: Final = create_pretrained_tokenizer(
|
||||
return _select_custom_tokenizer_helper(
|
||||
identifier=custom_tokenizer["identifier"],
|
||||
revision=custom_tokenizer["revision"],
|
||||
auth_token=custom_tokenizer["auth_token"],
|
||||
)
|
||||
return _tokenizer
|
||||
return _select_tokenizer_helper(model=model)
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _select_custom_tokenizer_helper(identifier: str, revision: str, auth_token: str | None) -> SelectTokenizerResponse:
|
||||
verbose_logger.debug("Loading custom HuggingFace tokenizer %s (revision %s)", identifier, revision)
|
||||
return create_pretrained_tokenizer(identifier=identifier, revision=revision, auth_token=auth_token)
|
||||
|
||||
|
||||
@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
|
||||
def _select_tokenizer_helper(model: str) -> SelectTokenizerResponse:
|
||||
if litellm.disable_hf_tokenizer_download is True:
|
||||
|
|
|
|||
|
|
@ -56,11 +56,87 @@ def test_no_user_or_assistant_rows():
|
|||
assert get_protected_indices([]) == ()
|
||||
|
||||
|
||||
def test_rows_before_last_cache_control_breakpoint_are_protected():
|
||||
messages = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "old question"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "old answer",
|
||||
"tool_calls": [{"id": "t1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "large file body"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "ack",
|
||||
"tool_calls": [{"id": "t2", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "t2", "content": "later tool output"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
protected = sorted(get_protected_indices(messages))
|
||||
|
||||
assert protected == [0, 1, 2, 3, 4, 5, 7]
|
||||
assert 6 not in protected
|
||||
|
||||
|
||||
def test_cache_control_directly_on_message_protects_prefix():
|
||||
messages = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "tool", "tool_call_id": "before", "content": "large file body"},
|
||||
{"role": "user", "content": "old question"},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "marked",
|
||||
"content": "cached tool",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "after", "content": "later tool output"},
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
protected = sorted(get_protected_indices(messages))
|
||||
|
||||
assert 1 in protected
|
||||
assert 3 in protected
|
||||
assert 4 not in protected
|
||||
|
||||
|
||||
def test_no_cache_control_leaves_history_compressible():
|
||||
messages = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": "old question"},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "large file body"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
assert sorted(get_protected_indices(messages)) == [0, 2, 4]
|
||||
|
||||
|
||||
def test_non_mapping_content_parts_are_not_cache_control():
|
||||
messages = [
|
||||
{"role": "system", "content": "sys"},
|
||||
{"role": "user", "content": ["not", "a", "dict"]},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
{"role": "tool", "tool_call_id": "t1", "content": "plain string"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
protected = sorted(get_protected_indices(messages))
|
||||
|
||||
assert protected == [0, 2, 4]
|
||||
assert 1 not in protected
|
||||
assert 3 not in protected
|
||||
|
||||
|
||||
def test_mid_history_cache_control_part_is_protected():
|
||||
# A large cached tool result from a few turns back, not the last user or
|
||||
# last assistant row -- exactly the row a provider prompt-cache pins to
|
||||
# exact bytes. Rewriting it (even leaving the marker on) changes those
|
||||
# bytes and turns the next request's cache read into a cache write.
|
||||
messages = [
|
||||
{"role": "user", "content": "old question"},
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
|
|
@ -74,9 +150,7 @@ def test_mid_history_cache_control_part_is_protected():
|
|||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
# index 3 = last assistant, index 4 = last user (both protected by role
|
||||
# regardless), index 2 = the cache_control-marked row itself.
|
||||
assert sorted(get_protected_indices(messages)) == [2, 3, 4]
|
||||
assert sorted(get_protected_indices(messages)) == [0, 1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_cache_control_directly_on_message_is_protected():
|
||||
|
|
@ -116,8 +190,6 @@ def test_content_that_is_not_a_list_of_mappings_is_not_treated_as_cache_control(
|
|||
|
||||
|
||||
def test_compress_keeps_part_level_cache_control_row_verbatim():
|
||||
# compress() scores text-only copies of the rows, where a part-level marker
|
||||
# is gone; protection has to read the original rows or the pinned row is stubbed.
|
||||
stale_log = {"role": "user", "content": [{"type": "text", "text": "stale log line " * 2000}]}
|
||||
pinned = {
|
||||
"role": "user",
|
||||
|
|
@ -126,9 +198,9 @@ def test_compress_keeps_part_level_cache_control_row_verbatim():
|
|||
],
|
||||
}
|
||||
messages = [
|
||||
stale_log,
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
pinned,
|
||||
{"role": "assistant", "content": "old answer"},
|
||||
stale_log,
|
||||
{"role": "assistant", "content": "ack"},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
|
@ -142,6 +214,6 @@ def test_compress_keeps_part_level_cache_control_row_verbatim():
|
|||
)
|
||||
|
||||
assert len(result["messages"]) == len(messages)
|
||||
assert result["messages"][2] == pinned
|
||||
assert result["messages"][0] != stale_log
|
||||
assert result["messages"][0] == pinned
|
||||
assert result["messages"][2] != stale_log
|
||||
assert len(result["cache"]) >= 1
|
||||
|
|
|
|||
|
|
@ -7155,3 +7155,21 @@ def test_get_additional_headers_survives_a_thread_growing_headers_mid_copy():
|
|||
assert copied["llm_provider-x-custom-1999"] == "1999"
|
||||
|
||||
_run_while_a_thread_grows(headers, read, reads=300)
|
||||
|
||||
|
||||
def test_add_dynamic_callback_registers_once_per_list_without_touching_the_callers_list(logging_obj: LitellmLogging):
|
||||
callback: Final = CustomLogger()
|
||||
caller_owned: Final = ["langfuse"]
|
||||
logging_obj.dynamic_success_callbacks = caller_owned
|
||||
|
||||
logging_obj.add_dynamic_callback(callback)
|
||||
logging_obj.add_dynamic_callback(callback)
|
||||
|
||||
assert caller_owned == ["langfuse"]
|
||||
assert logging_obj.dynamic_success_callbacks == ["langfuse", callback]
|
||||
assert logging_obj.dynamic_input_callbacks == [callback]
|
||||
assert logging_obj.dynamic_async_success_callbacks == [callback]
|
||||
assert logging_obj.dynamic_failure_callbacks == [callback]
|
||||
assert logging_obj.dynamic_async_failure_callbacks == [callback]
|
||||
assert LitellmLogging._with_dynamic_callback(None, callback) == [callback]
|
||||
assert LitellmLogging._with_dynamic_callback((callback,), callback) == [callback]
|
||||
|
|
|
|||
|
|
@ -5855,6 +5855,71 @@ async def test_organization_budget_check_carries_org_state_on_the_token():
|
|||
assert token.org_budget_snapshot == OrgBudgetSnapshot(spend=12.5, max_budget=100.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"max_budget, spend, expect_blocked",
|
||||
[
|
||||
(0.0, 0.0, True), # explicit zero budget blocks even a fresh org with no spend
|
||||
(0.0, 7.4e-06, True), # any spend at all against a zero budget blocks
|
||||
(None, 999.0, False), # unlimited (None) never blocks, regardless of spend
|
||||
(5.0, 4.99, False), # a positive budget under its cap still passes
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_organization_zero_max_budget_is_enforced(max_budget, spend, expect_blocked):
|
||||
"""An explicit organization max_budget of 0 must mean zero allowance, matching
|
||||
key/team/user semantics, not unlimited.
|
||||
|
||||
Regression for LIT-7797: `_organization_max_budget_check` returned early
|
||||
whenever `org_max_budget <= 0`, so an org configured with max_budget=0 could
|
||||
spend without limit.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_OrganizationTable
|
||||
from litellm.proxy.auth.auth_checks import _organization_max_budget_check
|
||||
|
||||
org_table = LiteLLM_OrganizationTable(
|
||||
organization_id="o1",
|
||||
organization_alias="zero-budget-org",
|
||||
budget_id="b1",
|
||||
created_by="admin",
|
||||
updated_by="admin",
|
||||
spend=spend,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=max_budget) if max_budget is not None else None,
|
||||
)
|
||||
token = UserAPIKeyAuth(token="k1", org_id="o1")
|
||||
user_api_key_cache = UserApiKeyCache()
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="org_id:o1:with_budget", value=org_table, model_type=LiteLLM_OrganizationTable
|
||||
)
|
||||
|
||||
async def _spend(counter_key, fallback_spend, max_budget=None, **kwargs):
|
||||
return spend
|
||||
|
||||
proxy_logging_obj = MagicMock()
|
||||
proxy_logging_obj.budget_alerts = AsyncMock()
|
||||
|
||||
with patch( # test-quality-ok: _organization_max_budget_check imports get_current_spend locally
|
||||
"litellm.proxy.proxy_server.get_current_spend", _spend
|
||||
):
|
||||
if expect_blocked:
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _organization_max_budget_check(
|
||||
valid_token=token,
|
||||
team_object=None,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
assert exc_info.value.max_budget == max_budget
|
||||
else:
|
||||
await _organization_max_budget_check(
|
||||
valid_token=token,
|
||||
team_object=None,
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("route", ["/health", "/health/services", "/health/test_connection"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_capable_non_llm_routes_still_enforce_budget(route):
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from litellm.proxy._types import (
|
|||
ProxyException,
|
||||
)
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
|
||||
from litellm.proxy.auth.handle_jwt import (
|
||||
JWKS_FETCH_ATTEMPTS,
|
||||
STALE_CACHE_KEY_PREFIX,
|
||||
|
|
@ -32,6 +33,7 @@ from litellm.proxy.auth.handle_jwt import (
|
|||
JWTHandler,
|
||||
NoMatchingJWTPublicKeyError,
|
||||
)
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -6786,3 +6788,180 @@ async def test_sync_user_role_and_teams_singular_claim_only_recognized_under_fla
|
|||
}
|
||||
assert mock_patch.call_args.kwargs["teams_ids_to_add_user_to"] == []
|
||||
assert user.teams == []
|
||||
|
||||
|
||||
def _entra_agent_registry() -> AgentRegistry:
|
||||
registry = AgentRegistry()
|
||||
registry.register_agent(
|
||||
AgentResponse(
|
||||
agent_id="canonical-agent-id",
|
||||
agent_name="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
|
||||
agent_card_params={"name": "research-agent", "url": "http://localhost:9999/a2a", "version": "1.0.0"},
|
||||
litellm_params={"require_trace_id_on_calls_by_agent": True},
|
||||
)
|
||||
)
|
||||
return registry
|
||||
|
||||
|
||||
def _entra_agent_jwt_handler(agent_id_jwt_field: str | None) -> JWTHandler:
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=DualCache(),
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(user_id_jwt_field="sub", agent_id_jwt_field=agent_id_jwt_field),
|
||||
)
|
||||
return jwt_handler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"claim_value",
|
||||
["canonical-agent-id", "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"],
|
||||
ids=["matches_agent_id", "matches_agent_name"],
|
||||
)
|
||||
def test_resolve_agent_id_returns_canonical_agent_id(claim_value: str):
|
||||
"""An Entra app token's azp claim binds to the registered agent by id or by name and yields its canonical id."""
|
||||
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
|
||||
|
||||
resolved = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token={"sub": "sp-object-id-1234", "azp": claim_value},
|
||||
agent_registry=_entra_agent_registry(),
|
||||
)
|
||||
|
||||
assert resolved == "canonical-agent-id"
|
||||
|
||||
|
||||
def test_resolve_agent_id_reads_nested_claim():
|
||||
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="entra.client_id")
|
||||
|
||||
resolved = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token={"sub": "sp-object-id-1234", "entra": {"client_id": "canonical-agent-id"}},
|
||||
agent_registry=_entra_agent_registry(),
|
||||
)
|
||||
|
||||
assert resolved == "canonical-agent-id"
|
||||
|
||||
|
||||
def test_resolve_agent_id_rejects_claim_for_unregistered_agent():
|
||||
"""A configured agent claim naming no registered agent fails closed with 403 instead of falling back to an unbound identity."""
|
||||
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"},
|
||||
agent_registry=_entra_agent_registry(),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"token",
|
||||
[
|
||||
{"sub": "sp-object-id-1234"},
|
||||
{"sub": "sp-object-id-1234", "azp": ""},
|
||||
{"sub": "sp-object-id-1234", "azp": ["canonical-agent-id"]},
|
||||
],
|
||||
ids=["claim_absent", "claim_empty", "claim_not_a_string"],
|
||||
)
|
||||
def test_resolve_agent_id_returns_none_when_claim_unusable(token: dict):
|
||||
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field="azp")
|
||||
|
||||
assert (
|
||||
JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler, jwt_valid_token=token, agent_registry=_entra_agent_registry()
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_agent_id_ignores_claim_when_field_not_configured():
|
||||
"""Without agent_id_jwt_field an azp claim (even an unknown one) leaves JWT auth behaviour unchanged."""
|
||||
jwt_handler = _entra_agent_jwt_handler(agent_id_jwt_field=None)
|
||||
|
||||
resolved = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=jwt_handler,
|
||||
jwt_valid_token={"sub": "sp-object-id-1234", "azp": "00000000-0000-0000-0000-000000000000"},
|
||||
agent_registry=_entra_agent_registry(),
|
||||
)
|
||||
|
||||
assert resolved is None
|
||||
|
||||
|
||||
def _entra_signed_app_token(monkeypatch, azp: str, scope: str) -> tuple[JWTHandler, str]:
|
||||
"""A JWTHandler that verifies RS256 tokens against a pre-cached JWKS, plus a signed Entra-style app token."""
|
||||
jwks_url = "https://login.microsoftonline.test/discovery/v2.0/keys"
|
||||
monkeypatch.setenv("JWT_PUBLIC_KEY_URL", jwks_url)
|
||||
monkeypatch.delenv("JWT_AUDIENCE", raising=False)
|
||||
private_key, jwk = _get_rsa_key_and_jwk(kid="entra-kid")
|
||||
cache = DualCache()
|
||||
cache.set_cache(key=f"litellm_jwt_auth_keys_{jwks_url}", value=[jwk])
|
||||
jwt_handler = JWTHandler()
|
||||
jwt_handler.update_environment(
|
||||
prisma_client=None,
|
||||
user_api_key_cache=cache,
|
||||
litellm_jwtauth=LiteLLM_JWTAuth(agent_id_jwt_field="azp"),
|
||||
)
|
||||
token = _encode_rsa_jwt(
|
||||
private_key,
|
||||
issuer="https://login.microsoftonline.test/lit7664-tenant/v2.0",
|
||||
audience="api://litellm",
|
||||
kid="entra-kid",
|
||||
extra_claims={"sub": "sp-object-id-1234", "azp": azp, "scope": scope},
|
||||
)
|
||||
return jwt_handler, token
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_admin_token", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
|
||||
async def test_auth_builder_propagates_agent_id_from_jwt_claim(monkeypatch, is_admin_token: bool):
|
||||
"""auth_builder carries the resolved agent id into JWTAuthBuilderResult on both the admin and standard paths."""
|
||||
jwt_handler, token = _entra_signed_app_token(
|
||||
monkeypatch,
|
||||
azp="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
|
||||
scope=LiteLLM_JWTAuth().admin_jwt_scope if is_admin_token else "",
|
||||
)
|
||||
jwt_handler.bind_agent_lookup(_entra_agent_registry())
|
||||
|
||||
result = await JWTAuthManager.auth_builder(
|
||||
api_key=token,
|
||||
jwt_handler=jwt_handler,
|
||||
request_data={"model": "gpt-5.6"},
|
||||
general_settings={"enforce_rbac": False},
|
||||
route="/key/info" if is_admin_token else "/chat/completions",
|
||||
prisma_client=None,
|
||||
user_api_key_cache=None,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
assert result["is_proxy_admin"] is is_admin_token
|
||||
assert result["agent_id"] == "canonical-agent-id"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_builder_denies_jwt_naming_unregistered_agent_before_admin_check(monkeypatch):
|
||||
"""An unknown agent claim is rejected even when the token would otherwise be a proxy admin."""
|
||||
jwt_handler, token = _entra_signed_app_token(
|
||||
monkeypatch,
|
||||
azp="00000000-0000-0000-0000-000000000000",
|
||||
scope=LiteLLM_JWTAuth().admin_jwt_scope,
|
||||
)
|
||||
jwt_handler.bind_agent_lookup(_entra_agent_registry())
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await JWTAuthManager.auth_builder(
|
||||
api_key=token,
|
||||
jwt_handler=jwt_handler,
|
||||
request_data={"model": "gpt-5.6"},
|
||||
general_settings={"enforce_rbac": False},
|
||||
route="/key/info",
|
||||
prisma_client=None,
|
||||
user_api_key_cache=None,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
|
|
|
|||
|
|
@ -1937,6 +1937,76 @@ async def test_standard_jwt_auth_propagates_user_email():
|
|||
assert result.api_key is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_proxy_admin", [False, True], ids=["standard_jwt", "proxy_admin_jwt"])
|
||||
async def test_jwt_auth_propagates_agent_id_to_user_api_key_auth(is_proxy_admin: bool):
|
||||
"""The agent id resolved by auth_builder must land on UserAPIKeyAuth.agent_id so
|
||||
agent-scoped checks (trace id requirement, MCP server/tool restrictions, spend
|
||||
attribution) apply to JWT callers the same way they apply to agent-bound keys."""
|
||||
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
|
||||
general_settings = {"enable_jwt_auth": True}
|
||||
user_api_key_cache = DualCache()
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.is_jwt.return_value = True
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(agent_id_jwt_field="azp")
|
||||
|
||||
user_object = LiteLLM_UserTable(user_id="sp-object-id-1234", user_role="internal_user")
|
||||
mock_jwt_result = {
|
||||
"is_proxy_admin": is_proxy_admin,
|
||||
"team_object": None,
|
||||
"user_object": user_object,
|
||||
"end_user_object": None,
|
||||
"org_object": None,
|
||||
"token": jwt_token,
|
||||
"team_id": None,
|
||||
"user_id": "sp-object-id-1234",
|
||||
"user_email": None,
|
||||
"end_user_id": None,
|
||||
"org_id": None,
|
||||
"team_membership": None,
|
||||
"jwt_claims": {"sub": "sp-object-id-1234", "azp": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"},
|
||||
"agent_id": "canonical-agent-id",
|
||||
}
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.method = "POST"
|
||||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
mock_request.state = SimpleNamespace()
|
||||
|
||||
with (
|
||||
patch.multiple( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
|
||||
"litellm.proxy.proxy_server",
|
||||
general_settings=general_settings,
|
||||
premium_user=True,
|
||||
master_key="sk-master",
|
||||
prisma_client=None,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
jwt_handler=jwt_handler,
|
||||
),
|
||||
patch( # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_jwt_result,
|
||||
),
|
||||
):
|
||||
result = await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=jwt_token,
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={"model": "gpt-5.6"},
|
||||
)
|
||||
|
||||
assert result.agent_id == "canonical-agent-id"
|
||||
assert result.user_id == "sp-object-id-1234"
|
||||
assert result.api_key is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_binds_api_key_to_token_hash():
|
||||
"""
|
||||
|
|
@ -2106,6 +2176,222 @@ async def test_auto_register_first_request_propagates_user_email():
|
|||
assert result.api_key == "hashed-auto-key"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auto_register_stamps_new_key_with_jwt_agent_id():
|
||||
"""The virtual key AUTO_REGISTER creates must carry the agent id auth_builder bound
|
||||
from the JWT claim, and the first request's principal must carry it too, or the
|
||||
mapped-key path would drop the agent policies on that request and every later one."""
|
||||
from litellm.proxy.auth.auth_method import AuthMethod
|
||||
from litellm.proxy.auth.resolvers.models import CredentialRef
|
||||
from litellm.proxy.auth.resolvers.store import IdentityStore
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
from litellm.proxy.proxy_server import hash_token
|
||||
|
||||
plaintext = "sk-auto-registered-agent"
|
||||
token_hash = hash_token(plaintext)
|
||||
persisted_principal = IdentityStore._principal_from_key(
|
||||
UserAPIKeyAuth(token=token_hash, user_id="validated-user", team_id="validated-team", agent_id="canonical-agent-id"),
|
||||
auth_method=AuthMethod.API_KEY,
|
||||
credential_ref=CredentialRef(token_id=token_hash),
|
||||
)
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock()
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
|
||||
generate_key = AsyncMock(return_value={"token": plaintext})
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
generate_key,
|
||||
),
|
||||
patch( # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists
|
||||
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
|
||||
new_callable=AsyncMock,
|
||||
return_value=persisted_principal,
|
||||
),
|
||||
):
|
||||
result = await _auto_register_jwt_mapping(
|
||||
virtual_key_claim_field="appid",
|
||||
claim_value="2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
cache_key="jwt_key_mapping:appid:2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21",
|
||||
team_id="validated-team",
|
||||
user_id="validated-user",
|
||||
agent_id="canonical-agent-id",
|
||||
)
|
||||
|
||||
assert generate_key.await_args is not None
|
||||
assert generate_key.await_args.kwargs["agent_id"] == "canonical-agent-id"
|
||||
assert result is not None
|
||||
assert result.agent_id == "canonical-agent-id"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("losing_agent_id", ["other-agent", None], ids=["different_agent", "no_agent_claim"])
|
||||
async def test_auto_register_race_loser_keeps_winners_agent_id(losing_agent_id: str | None):
|
||||
"""When two requests race to AUTO_REGISTER the same mapping claim, the loser must run as
|
||||
the persisted key, agent binding included. Every later request on that mapping uses the
|
||||
winner's key, so stamping the loser's own (or missing) agent id on it would give one request
|
||||
different agent policies and spend attribution than all the others."""
|
||||
from litellm.proxy.auth.auth_method import AuthMethod
|
||||
from litellm.proxy.auth.resolvers.models import CredentialRef
|
||||
from litellm.proxy.auth.resolvers.store import IdentityStore
|
||||
from litellm.proxy.auth.user_api_key_auth import _auto_register_jwt_mapping
|
||||
|
||||
winner_hash = "winner-key-hash"
|
||||
winner_principal = IdentityStore._principal_from_key(
|
||||
UserAPIKeyAuth(token=winner_hash, user_id="validated-user", team_id="validated-team", agent_id="winner-agent"),
|
||||
auth_method=AuthMethod.API_KEY,
|
||||
credential_ref=CredentialRef(token_id=winner_hash),
|
||||
)
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_jwtkeymapping.create = AsyncMock(
|
||||
side_effect=Exception("Unique constraint failed on the fields: (`jwt_claim_name`,`jwt_claim_value`)")
|
||||
)
|
||||
prisma_client.db.litellm_verificationtoken.delete = AsyncMock()
|
||||
user_api_key_cache = MagicMock()
|
||||
user_api_key_cache.async_set_cache = AsyncMock()
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_mapping_cache_ttl=300)
|
||||
|
||||
with (
|
||||
patch( # test-quality-ok: key creation is an inline import inside the helper; no dependency injection seam exists
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"token": "sk-orphaned-loser-key"},
|
||||
),
|
||||
patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
|
||||
"litellm.proxy.auth.user_api_key_auth.get_jwt_key_mapping_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=winner_hash,
|
||||
),
|
||||
patch( # test-quality-ok: the helper constructs IdentityStore itself; no dependency injection seam exists
|
||||
"litellm.proxy.auth.resolvers.store.IdentityStore.resolve",
|
||||
new_callable=AsyncMock,
|
||||
return_value=winner_principal,
|
||||
),
|
||||
):
|
||||
result = await _auto_register_jwt_mapping(
|
||||
virtual_key_claim_field="tid",
|
||||
claim_value="shared-tenant",
|
||||
jwt_handler=jwt_handler,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
cache_key="jwt_key_mapping:tid:shared-tenant",
|
||||
team_id="validated-team",
|
||||
user_id="validated-user",
|
||||
agent_id=losing_agent_id,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.token == winner_hash
|
||||
assert result.agent_id == "winner-agent"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_jwt_auto_register_forwards_bound_agent_id():
|
||||
"""When a JWT under AUTO_REGISTER also carries the configured agent claim, the agent
|
||||
id auth_builder resolved must reach the key creation, not be dropped when
|
||||
valid_token is swapped for the freshly registered key."""
|
||||
jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.signature"
|
||||
user_api_key_cache = DualCache()
|
||||
jwt_handler = MagicMock()
|
||||
jwt_handler.is_jwt.return_value = True
|
||||
jwt_handler.auth_jwt = AsyncMock(return_value={"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"})
|
||||
jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(
|
||||
virtual_key_claim_field="sub",
|
||||
virtual_key_mapping_cache_ttl=300,
|
||||
agent_id_jwt_field="appid",
|
||||
)
|
||||
user_object = LiteLLM_UserTable(user_id="validated-user", user_role="internal_user")
|
||||
mock_jwt_result = {
|
||||
"is_proxy_admin": False,
|
||||
"team_object": None,
|
||||
"user_object": user_object,
|
||||
"end_user_object": None,
|
||||
"org_object": None,
|
||||
"token": jwt_token,
|
||||
"team_id": "validated-team",
|
||||
"user_id": "validated-user",
|
||||
"user_email": None,
|
||||
"end_user_id": None,
|
||||
"org_id": None,
|
||||
"team_membership": None,
|
||||
"jwt_claims": {"sub": "user1", "appid": "2f5c9b1e-6a4d-4c8e-9f0b-7d1a3e5c9b21"},
|
||||
"agent_id": "canonical-agent-id",
|
||||
}
|
||||
auto_register = AsyncMock(
|
||||
return_value=UserAPIKeyAuth(
|
||||
token="hashed-auto-key",
|
||||
api_key="hashed-auto-key",
|
||||
team_id="validated-team",
|
||||
user_id="validated-user",
|
||||
agent_id="canonical-agent-id",
|
||||
)
|
||||
)
|
||||
|
||||
mock_request = MagicMock()
|
||||
mock_request.url.path = "/v1/chat/completions"
|
||||
mock_request.method = "POST"
|
||||
mock_request.headers = {"authorization": f"Bearer {jwt_token}"}
|
||||
mock_request.query_params = {}
|
||||
mock_request.state = SimpleNamespace()
|
||||
|
||||
with (
|
||||
patch.multiple( # test-quality-ok: production auth reads these module globals; no dependency injection seam exists
|
||||
"litellm.proxy.proxy_server",
|
||||
general_settings={"enable_jwt_auth": True},
|
||||
premium_user=True,
|
||||
master_key="sk-master",
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
jwt_handler=jwt_handler,
|
||||
),
|
||||
patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
|
||||
"litellm.proxy.auth.user_api_key_auth._resolve_jwt_to_virtual_key",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_PendingAutoRegister(
|
||||
claim_field="sub",
|
||||
claim_value="user1",
|
||||
cache_key="jwt_key_mapping:sub:user1",
|
||||
),
|
||||
),
|
||||
patch( # test-quality-ok: the builder calls this static method directly; no dependency injection seam exists
|
||||
"litellm.proxy.auth.user_api_key_auth.JWTAuthManager.auth_builder",
|
||||
new_callable=AsyncMock,
|
||||
return_value=mock_jwt_result,
|
||||
),
|
||||
patch( # test-quality-ok: module-level helper called by the builder; no dependency injection seam exists
|
||||
"litellm.proxy.auth.user_api_key_auth._auto_register_jwt_mapping",
|
||||
auto_register,
|
||||
),
|
||||
):
|
||||
result = await _user_api_key_auth_builder(
|
||||
request=mock_request,
|
||||
api_key=jwt_token,
|
||||
azure_api_key_header="",
|
||||
anthropic_api_key_header=None,
|
||||
google_ai_studio_api_key_header=None,
|
||||
azure_apim_header=None,
|
||||
request_data={"model": "gpt-5.6"},
|
||||
)
|
||||
|
||||
assert auto_register.await_args is not None
|
||||
assert auto_register.await_args.kwargs["agent_id"] == "canonical-agent-id"
|
||||
assert result.agent_id == "canonical-agent-id"
|
||||
assert result.api_key == "hashed-auto-key"
|
||||
|
||||
|
||||
class TestJWTOAuth2Coexistence:
|
||||
"""
|
||||
Test that JWT and OAuth2 auth can coexist on the same instance.
|
||||
|
|
|
|||
|
|
@ -900,7 +900,7 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr
|
|||
)
|
||||
|
||||
assert has_headroom_retrieve_tool(result.get("tools") or [])
|
||||
(issued, _expiry), = guardrail._issued_hashes_by_call_id.values()
|
||||
((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values()
|
||||
assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"})
|
||||
|
||||
|
||||
|
|
@ -953,7 +953,6 @@ async def test_anthropic_assistant_history_never_reaches_compression_service(gua
|
|||
assert result["messages"][1]["content"] == [{"type": "text", "text": table}]
|
||||
|
||||
|
||||
|
||||
def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape():
|
||||
"""By the time an Anthropic Messages API response reaches the agentic-loop
|
||||
gate, the OpenAI-shaped tool this guardrail injects (type: "function")
|
||||
|
|
@ -2342,9 +2341,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end(
|
|||
)
|
||||
assert streamed_text == final_answer
|
||||
assert not any("function_call" in str(getattr(event, "type", "")) for event in events)
|
||||
assert not any(
|
||||
getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
|
||||
)
|
||||
assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events)
|
||||
mock_get.assert_called_once()
|
||||
assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0])
|
||||
|
||||
|
|
@ -2399,9 +2396,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end(
|
|||
getattr(event, "delta", "") for event in events if getattr(event, "type", None) == "response.output_text.delta"
|
||||
)
|
||||
assert streamed_text == final_answer
|
||||
assert not any(
|
||||
getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events
|
||||
)
|
||||
assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events)
|
||||
mock_get.assert_called_once()
|
||||
assert len(upstream.calls) == 2
|
||||
assert not json.loads(upstream.calls[1].request.content).get("stream")
|
||||
|
|
@ -2514,6 +2509,38 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail):
|
|||
assert messages[3] == compressed_history[1]
|
||||
|
||||
|
||||
CACHED_PREFIX_MESSAGES = [
|
||||
{"role": "system", "content": "You are Claude Code. " + "S" * 5000},
|
||||
{"role": "user", "content": "old question " + "Q" * 5000},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Reading the file now.",
|
||||
"tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "old_1", "content": "large file body " + "F" * 5000},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "cached turn", "cache_control": {"type": "ephemeral"}}],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Listing now.",
|
||||
"tool_calls": [{"id": "new_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "new_1", "content": "volatile tail output " + "T" * 5000},
|
||||
{"role": "assistant", "content": "Finished listing."},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rows_before_last_cache_control_breakpoint_are_never_sent(guardrail: HeadroomGuardrail):
|
||||
wire, result = await _wire_and_result(guardrail, CACHED_PREFIX_MESSAGES)
|
||||
|
||||
assert [row.get("tool_call_id") for row in wire] == ["new_1"]
|
||||
assert result["structured_messages"][:5] == CACHED_PREFIX_MESSAGES[:5]
|
||||
|
||||
|
||||
CACHE_MARKED_HISTORY_MESSAGES = [
|
||||
{"role": "system", "content": "You are Claude Code. " + "S" * 5000},
|
||||
{"role": "user", "content": "old question " + "Q" * 5000},
|
||||
|
|
@ -2529,6 +2556,7 @@ CACHE_MARKED_HISTORY_MESSAGES = [
|
|||
"cache_control": {"type": "ephemeral"},
|
||||
},
|
||||
{"role": "assistant", "content": "Summarized the file for you."},
|
||||
{"role": "tool", "tool_call_id": "tail", "content": "volatile tail output " + "T" * 5000},
|
||||
{"role": "user", "content": "live instruction"},
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -338,6 +338,15 @@ def test_a_request_must_carry_exactly_one_usable_conversation(body: dict):
|
|||
"config_overrides",
|
||||
[
|
||||
{"classifier_type": "llm", "classifier_llm_config": {"model": "classifier-model"}},
|
||||
{
|
||||
"classifier_type": "capability",
|
||||
"classifier_llm_config": {"model": "classifier-model"},
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.5,
|
||||
},
|
||||
},
|
||||
{
|
||||
"semantic_keyword_matching": True,
|
||||
"embedding_model": "classifier-model",
|
||||
|
|
|
|||
|
|
@ -245,6 +245,55 @@ async def test_validate_team_org_change_same_org_id():
|
|||
mock_access_check.assert_not_called() # Ensure access check wasn't called
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"org_max_budget, team_max_budget, expect_blocked",
|
||||
[
|
||||
(0.0, 100.0, True), # explicit zero org budget must still cap the team's budget
|
||||
(0.0, None, False), # team has no budget of its own, nothing to compare
|
||||
(None, 100.0, False), # unlimited (None) org budget never blocks
|
||||
(50.0, 100.0, True), # a positive org budget is still enforced normally
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_team_org_change_zero_org_budget_is_enforced(
|
||||
org_max_budget, team_max_budget, expect_blocked
|
||||
):
|
||||
"""An organization with an explicit max_budget of 0 must still block moving in a
|
||||
team with a larger budget, matching key/team/user zero-budget semantics.
|
||||
|
||||
Regression for LIT-7797: the truthy check `organization.litellm_budget_table.max_budget`
|
||||
treated an explicit 0 the same as no budget table at all, silently skipping this guard.
|
||||
"""
|
||||
org_id = "team-org-123"
|
||||
new_org_id = "new-org-456"
|
||||
|
||||
team = MagicMock(spec=LiteLLM_TeamTable)
|
||||
team.organization_id = org_id
|
||||
team.models = []
|
||||
team.max_budget = team_max_budget
|
||||
team.tpm_limit = None
|
||||
team.rpm_limit = None
|
||||
team.members_with_roles = []
|
||||
|
||||
organization = MagicMock(spec=LiteLLM_OrganizationTableWithMembers)
|
||||
organization.organization_id = new_org_id
|
||||
organization.models = []
|
||||
organization.litellm_budget_table = (
|
||||
LiteLLM_BudgetTable(max_budget=org_max_budget) if org_max_budget is not None else None
|
||||
)
|
||||
organization.members = []
|
||||
|
||||
mock_router = MagicMock(spec=Router)
|
||||
|
||||
if expect_blocked:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
validate_team_org_change(team=team, organization=organization, llm_router=mock_router)
|
||||
assert exc_info.value.status_code == 403
|
||||
else:
|
||||
result = validate_team_org_change(team=team, organization=organization, llm_router=mock_router)
|
||||
assert result is None or result is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_team_org_change_members_in_org():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -3764,6 +3764,55 @@ async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("agents_source", ["config", "db", "api"])
|
||||
async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_registry(
|
||||
clean_agent_registry, agents_source
|
||||
):
|
||||
"""A JWT agent claim must resolve against every agent the proxy knows, including ones created after startup."""
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import LiteLLM_JWTAuth
|
||||
from litellm.proxy.auth.handle_jwt import JWTAuthManager
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
original_lookup = proxy_server.jwt_handler.agent_lookup
|
||||
try:
|
||||
proxy_server.ProxyStartupEvent._initialize_jwt_auth(
|
||||
general_settings={"litellm_jwtauth": {"agent_id_jwt_field": "appid"}},
|
||||
prisma_client=None,
|
||||
user_api_key_cache=UserApiKeyCache(),
|
||||
)
|
||||
if agents_source == "config":
|
||||
await ProxyConfig()._init_non_llm_configs(
|
||||
config={"agents": [_config_agent("loaded-agent")]},
|
||||
config_file_path=None,
|
||||
)
|
||||
elif agents_source == "db":
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_agentstable.find_many = AsyncMock(
|
||||
return_value=[_FakeAgentRow("db-id", "loaded-agent")]
|
||||
)
|
||||
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
|
||||
else:
|
||||
clean_agent_registry.register_agent(
|
||||
agent_config=AgentResponse(agent_id="api-id", **_config_agent("loaded-agent"))
|
||||
)
|
||||
|
||||
resolved = JWTAuthManager.resolve_agent_id(
|
||||
jwt_handler=proxy_server.jwt_handler,
|
||||
jwt_valid_token={"appid": "loaded-agent"},
|
||||
agent_registry=proxy_server.jwt_handler.agent_lookup,
|
||||
)
|
||||
finally:
|
||||
proxy_server.jwt_handler.bind_agent_lookup(original_lookup)
|
||||
proxy_server.jwt_handler.update_environment(
|
||||
prisma_client=None, user_api_key_cache=UserApiKeyCache(), litellm_jwtauth=LiteLLM_JWTAuth()
|
||||
)
|
||||
|
||||
assert resolved == clean_agent_registry.get_agent_by_name(agent_name="loaded-agent").agent_id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"config, expected_agent_names",
|
||||
|
|
|
|||
|
|
@ -13521,3 +13521,54 @@ async def test_token_counter_loads_a_custom_tokenizer_off_the_event_loop(monkeyp
|
|||
assert response.tokenizer_type == "huggingface_tokenizer"
|
||||
assert response.total_tokens > 0
|
||||
assert_loop_stayed_free(took, lags)
|
||||
|
||||
|
||||
async def test_token_counter_loads_a_custom_tokenizer_once_per_identifier_revision_and_token(monkeypatch):
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
from litellm import Router
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
|
||||
claude_tokenizer: Final[Tokenizer] = litellm.utils._select_tokenizer("claude-fable-5")["tokenizer"]
|
||||
from_pretrained: Final = MagicMock(return_value=claude_tokenizer)
|
||||
|
||||
def deployment(model_name: str, revision: str, auth_token: str | None) -> DeploymentTypedDict:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": "openai/self-hosted-model", "api_base": "http://localhost:8080/v1"},
|
||||
"model_info": {
|
||||
"custom_tokenizer": {"identifier": "my-org/tokenizer", "revision": revision, "auth_token": auth_token}
|
||||
},
|
||||
}
|
||||
|
||||
monkeypatch.setattr(litellm.utils, "Tokenizer", MagicMock(from_pretrained=from_pretrained))
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.proxy_server.llm_router",
|
||||
Router(
|
||||
model_list=[
|
||||
deployment("self-hosted", "main", None),
|
||||
deployment("self-hosted-pinned", "v2", None),
|
||||
deployment("self-hosted-private", "main", "hf_test_token"),
|
||||
]
|
||||
),
|
||||
)
|
||||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
try:
|
||||
responses: Final = [
|
||||
await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted", prompt="count me once"))
|
||||
for _ in range(3)
|
||||
]
|
||||
assert from_pretrained.call_args_list == [mock.call("my-org/tokenizer", revision="main", token=None)]
|
||||
assert all(response.tokenizer_type == "huggingface_tokenizer" for response in responses)
|
||||
assert len({response.total_tokens for response in responses}) == 1
|
||||
assert responses[0].total_tokens > 0
|
||||
|
||||
await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-pinned", prompt="count me once"))
|
||||
await proxy_server_module.token_counter(TokenCountRequest(model="self-hosted-private", prompt="count me once"))
|
||||
assert from_pretrained.call_args_list == [
|
||||
mock.call("my-org/tokenizer", revision="main", token=None),
|
||||
mock.call("my-org/tokenizer", revision="v2", token=None),
|
||||
mock.call("my-org/tokenizer", revision="main", token="hf_test_token"),
|
||||
]
|
||||
finally:
|
||||
litellm.utils._select_custom_tokenizer_helper.cache_clear()
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging
|
|||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from litellm.proxy.utils import get_custom_url, join_paths
|
||||
|
||||
|
|
@ -1303,12 +1303,10 @@ class TestPostCallFailureHookLLMExceptionAlerting:
|
|||
"""The llm_exceptions alert is for infra / LLM-API failures, not user
|
||||
errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized
|
||||
client errors must be excluded so a guardrail content-policy block never
|
||||
pages on-call. ProxyException is such an error; before LIT-3751 only
|
||||
HTTPException was excluded, so AIM blocks paged as if the LLM API failed."""
|
||||
pages on-call. 5xx proxy errors still alert."""
|
||||
|
||||
async def _alerted(self, exc) -> bool:
|
||||
async def _alerted(self, exc: Exception) -> AsyncMock:
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.proxy._types import AlertType, UserAPIKeyAuth
|
||||
|
||||
|
|
@ -1325,7 +1323,7 @@ class TestPostCallFailureHookLLMExceptionAlerting:
|
|||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
await asyncio.sleep(0) # let the fire-and-forget alert task run
|
||||
return alerting_handler.called
|
||||
return alerting_handler
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_exception_does_not_alert(self):
|
||||
|
|
@ -1338,15 +1336,49 @@ class TestPostCallFailureHookLLMExceptionAlerting:
|
|||
code=400,
|
||||
openai_code="content_policy_violation",
|
||||
)
|
||||
assert await self._alerted(exc) is False
|
||||
assert (await self._alerted(exc)).called is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_exception_does_not_alert(self):
|
||||
assert await self._alerted(HTTPException(status_code=400, detail="blocked")) is False
|
||||
assert (await self._alerted(HTTPException(status_code=400, detail="blocked"))).called is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_genuine_llm_api_error_still_alerts(self):
|
||||
assert await self._alerted(Exception("upstream 503")) is True
|
||||
assert (await self._alerted(Exception("upstream 503"))).called is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_exception_5xx_alerts(self):
|
||||
alerting_handler = await self._alerted(
|
||||
HTTPException(
|
||||
status_code=502,
|
||||
detail={
|
||||
"error": "Headroom compression service returned an error",
|
||||
"status_code": 503,
|
||||
"guardrail_name": "headroom-compression-global",
|
||||
},
|
||||
)
|
||||
)
|
||||
assert alerting_handler.called is True
|
||||
assert "headroom-compression-global" in alerting_handler.call_args.kwargs["message"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_exception_5xx_alerts(self):
|
||||
from litellm.proxy._types import ProxyException
|
||||
|
||||
alerting_handler = await self._alerted(
|
||||
ProxyException(
|
||||
message="guardrail backend down",
|
||||
type="internal_server_error",
|
||||
param=None,
|
||||
code=503,
|
||||
)
|
||||
)
|
||||
assert alerting_handler.called is True
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_exception_429_does_not_alert(self):
|
||||
alerting_handler = await self._alerted(HTTPException(status_code=429, detail="rate limited"))
|
||||
assert alerting_handler.called is False
|
||||
|
||||
|
||||
class TestPostCallFailureHookProxyExceptionLogging:
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ Tests the rule-based complexity scoring and tier assignment logic.
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
|
|
@ -52,7 +53,13 @@ from litellm.router_strategy.complexity_router.complexity_router import (
|
|||
classification_system_prompt,
|
||||
custom_tier_classification_prompt,
|
||||
)
|
||||
from litellm.router_strategy.complexity_router.capability_classifier import (
|
||||
CAPABILITY_CLASSIFIER_SYSTEM_PROMPT,
|
||||
CapabilityClassifierVerdict,
|
||||
)
|
||||
from litellm.router_strategy.complexity_router.config import (
|
||||
CapabilityCalibrationConfig,
|
||||
CapabilityClassifierConfig,
|
||||
DEFAULT_CLASSIFICATION_RUBRIC,
|
||||
DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE,
|
||||
DEFAULT_COMPLEXITY_CONFIG,
|
||||
|
|
@ -2363,6 +2370,515 @@ class TestLLMClassifierConfig:
|
|||
)
|
||||
|
||||
|
||||
CAPABILITY_TIERS: Dict[str, str] = {
|
||||
"SIMPLE": "efficient-model",
|
||||
"REASONING": "capable-model",
|
||||
}
|
||||
|
||||
|
||||
def _capability_router_config(**overrides):
|
||||
return {
|
||||
"tiers": dict(CAPABILITY_TIERS),
|
||||
"classifier_type": "capability",
|
||||
"classifier_llm_config": {"model": "judge-model", "timeout_ms": 400},
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.5,
|
||||
"threshold_step": 0.1,
|
||||
},
|
||||
**overrides,
|
||||
}
|
||||
|
||||
|
||||
def _capability_reply(
|
||||
*,
|
||||
p_solve: float,
|
||||
primary_rule: str = "SUP-1",
|
||||
capability_boundary: str = "supported",
|
||||
crux: str = "complete the requested change",
|
||||
) -> str:
|
||||
return json.dumps(
|
||||
{
|
||||
"crux": crux,
|
||||
"primary_rule": primary_rule,
|
||||
"capability_boundary": capability_boundary,
|
||||
"p_solve": p_solve,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class TestCapabilityClassifierConfig:
|
||||
@pytest.mark.parametrize(
|
||||
"calibration",
|
||||
(
|
||||
{"version": "v1", "slope": -1.0, "intercept": 0.0},
|
||||
{"version": "v1", "slope": float("nan"), "intercept": 0.0},
|
||||
{"version": "v1", "slope": 1.0, "intercept": float("inf")},
|
||||
{"version": "v1", "slope": True, "intercept": 0.0},
|
||||
{"version": " ", "slope": 1.0, "intercept": 0.0},
|
||||
{"version": "v1", "slope": 1.0, "intercept": 0.0, "typo": 1},
|
||||
),
|
||||
)
|
||||
def test_rejects_invalid_calibration(self, calibration: dict[str, object]) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
CapabilityCalibrationConfig.model_validate(calibration)
|
||||
|
||||
def test_calibration_round_trip_and_probability_endpoints(self) -> None:
|
||||
calibration: Final = CapabilityCalibrationConfig(version="held-out-v1", slope=0.0, intercept=0.0)
|
||||
config: Final = CapabilityClassifierConfig(
|
||||
efficient_tier="SIMPLE", capable_tier="REASONING", base_threshold=0.6, calibration=calibration
|
||||
)
|
||||
restored: Final = CapabilityClassifierConfig.model_validate_json(config.model_dump_json())
|
||||
assert restored.calibration == calibration
|
||||
assert tuple(calibration.calibrate(p) for p in (0.0, 0.5, 1.0)) == (0.5, 0.5, 0.5)
|
||||
steep: Final = CapabilityCalibrationConfig(version="endpoints", slope=20.0, intercept=-20.0)
|
||||
values: Final = tuple(steep.calibrate(p) for p in (0.0, 0.5, 1.0))
|
||||
assert all(math.isfinite(p) and 0.0 <= p <= 1.0 for p in values)
|
||||
assert values[0] < values[1] < values[2]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"patch,error_match",
|
||||
[
|
||||
({"classifier_llm_config": None}, "classifier_llm_config is required"),
|
||||
({"capability_classifier_config": None}, "capability_classifier_config is required"),
|
||||
(
|
||||
{
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "SIMPLE",
|
||||
"base_threshold": 0.5,
|
||||
}
|
||||
},
|
||||
"must be a higher tier",
|
||||
),
|
||||
(
|
||||
{
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "REASONING",
|
||||
"capable_tier": "SIMPLE",
|
||||
"base_threshold": 0.5,
|
||||
}
|
||||
},
|
||||
"must be a higher tier",
|
||||
),
|
||||
(
|
||||
{
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "MEDIUM",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.5,
|
||||
}
|
||||
},
|
||||
"has no model configured",
|
||||
),
|
||||
(
|
||||
{
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.9,
|
||||
"threshold_step": 0.1,
|
||||
}
|
||||
},
|
||||
r"base_threshold \+ 2 \* threshold_step must be at most 1",
|
||||
),
|
||||
({"classifier_fallback": "default_model", "default_model": "fallback"}, "always fails closed"),
|
||||
(
|
||||
{"classifier_llm_config": {"model": "judge-model", "system_prompt": "pick one"}},
|
||||
"uses the packaged capability card",
|
||||
),
|
||||
({"classification_examples": "example"}, "uses the packaged capability card"),
|
||||
],
|
||||
)
|
||||
def test_rejects_incoherent_configuration(self, patch, error_match):
|
||||
with pytest.raises(ValidationError, match=error_match):
|
||||
ComplexityRouterConfig(**{**_capability_router_config(), **patch})
|
||||
|
||||
def test_capability_config_is_rejected_on_other_classifier_types(self):
|
||||
config = _capability_router_config(classifier_type="llm")
|
||||
with pytest.raises(ValidationError, match="requires classifier_type 'capability'"):
|
||||
ComplexityRouterConfig(**config)
|
||||
|
||||
def test_rejects_misspelled_optional_policy_instead_of_using_defaults(self) -> None:
|
||||
with pytest.raises(ValidationError, match="threshold_steps"):
|
||||
CapabilityClassifierConfig.model_validate(
|
||||
{
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.5,
|
||||
"threshold_steps": 0.2,
|
||||
}
|
||||
)
|
||||
|
||||
def test_threshold_defaults_match_switchyard(self):
|
||||
config = CapabilityClassifierConfig(efficient_tier=" SIMPLE ", capable_tier=" REASONING ", base_threshold=0.5)
|
||||
assert config.efficient_tier == "SIMPLE"
|
||||
assert config.capable_tier == "REASONING"
|
||||
assert config.threshold_step == 0.0
|
||||
assert config.max_output_tokens == 4096
|
||||
|
||||
def test_classifier_model_is_registered_as_a_dependency(self):
|
||||
assert ComplexityRouterConfig(**_capability_router_config()).uses_llm_classifier is True
|
||||
|
||||
|
||||
class TestCapabilityClassifierVerdict:
|
||||
@pytest.mark.parametrize(
|
||||
"primary_rule,capability_boundary",
|
||||
[
|
||||
*((f"SUP-{index}", "supported") for index in range(1, 6)),
|
||||
*((f"UNC-{index}", "uncertain") for index in range(1, 3)),
|
||||
*((f"LIM-{index}", "unsupported") for index in range(1, 3)),
|
||||
("none", "unmatched"),
|
||||
],
|
||||
)
|
||||
def test_accepts_every_valid_rule_boundary_pair(self, primary_rule, capability_boundary):
|
||||
verdict = CapabilityClassifierVerdict(
|
||||
crux="the hard part",
|
||||
primary_rule=primary_rule,
|
||||
capability_boundary=capability_boundary,
|
||||
p_solve=0.5,
|
||||
)
|
||||
assert verdict.primary_rule == primary_rule
|
||||
assert verdict.capability_boundary == capability_boundary
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"payload,error_match",
|
||||
[
|
||||
(
|
||||
{
|
||||
"crux": "x",
|
||||
"primary_rule": "SUP-1",
|
||||
"capability_boundary": "unsupported",
|
||||
"p_solve": 0.5,
|
||||
},
|
||||
"requires capability_boundary",
|
||||
),
|
||||
(
|
||||
{"crux": " ", "primary_rule": "none", "capability_boundary": "unmatched", "p_solve": 0.5},
|
||||
"non-whitespace",
|
||||
),
|
||||
(
|
||||
{
|
||||
"crux": "x",
|
||||
"primary_rule": "none",
|
||||
"capability_boundary": "unmatched",
|
||||
"p_solve": 0.5,
|
||||
"recommended_route": "efficient",
|
||||
},
|
||||
"Extra inputs are not permitted",
|
||||
),
|
||||
(
|
||||
{"crux": "x", "primary_rule": "none", "capability_boundary": "unmatched", "p_solve": True},
|
||||
"valid number",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_rejects_invalid_or_inconsistent_verdicts(self, payload, error_match):
|
||||
with pytest.raises(ValidationError, match=error_match):
|
||||
CapabilityClassifierVerdict.model_validate(payload)
|
||||
|
||||
|
||||
class TestCapabilityClassifier:
|
||||
@staticmethod
|
||||
def _router(mock_router_instance, **overrides):
|
||||
return ComplexityRouter(
|
||||
model_name="capability-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config=_capability_router_config(**overrides),
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_encrypted_task_is_not_replaced_by_plaintext_envelope(self, mock_router_instance: MagicMock) -> None:
|
||||
mock_router_instance.aresponses = AsyncMock(
|
||||
return_value=_native_classifier_response(_capability_reply(p_solve=0.8))
|
||||
)
|
||||
router: Final = self._router(mock_router_instance)
|
||||
task: Final = _encrypted_agent_task()
|
||||
request: Final = {"input": [task]}
|
||||
original: Final = deepcopy(request)
|
||||
result: Final = await router.async_pre_routing_hook(model="capability-router", request_kwargs=request)
|
||||
assert result is not None and result.model == "efficient-model"
|
||||
assert result.routing_decision is not None
|
||||
assert result.routing_decision["cause"] == "capability_classifier"
|
||||
mock_router_instance.aresponses.assert_awaited_once()
|
||||
call: Final = mock_router_instance.aresponses.call_args.kwargs
|
||||
assert call["input"][-1] == task
|
||||
plaintext: Final = json.dumps(call["input"][:-1])
|
||||
assert "The delegated task in the following agent_message." in plaintext
|
||||
assert "Message Type: NEW_TASK" not in plaintext
|
||||
assert "opaque-provider-task" not in plaintext
|
||||
assert request == original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("custom_markers", (False, True))
|
||||
async def test_task_forecast_uses_request_scoped_codex_markers(
|
||||
self, mock_router_instance: MagicMock, custom_markers: bool
|
||||
) -> None:
|
||||
completion: Final = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8)))
|
||||
mock_router_instance.acompletion = completion
|
||||
router: Final = self._router(
|
||||
mock_router_instance,
|
||||
escalation_keywords=[],
|
||||
**({"reminder_markers": [{"open": "<custom>", "close": "</custom>"}]} if custom_markers else {}),
|
||||
)
|
||||
envelope: Final = "\n".join(_CODEX_ENVELOPES)
|
||||
opening: Final = f"{envelope}\nFix nested behavior"
|
||||
messages: Final = [
|
||||
{"role": "user", "content": opening},
|
||||
{"role": "user", "content": "Preserve empty inputs"},
|
||||
{"role": "user", "content": envelope},
|
||||
]
|
||||
original: Final = deepcopy(messages)
|
||||
for user_agent in ("codex-tui", "curl/8.7.1", "codex_cli_rs/0.62.0"):
|
||||
result: Final = await router.async_pre_routing_hook(
|
||||
model="capability-router", messages=messages, request_kwargs={"metadata": {"user_agent": user_agent}}
|
||||
)
|
||||
assert result is not None and result.model == "efficient-model"
|
||||
sent: Final = completion.call_args.kwargs["messages"]
|
||||
if user_agent.startswith("codex") and not custom_markers:
|
||||
assert [message["content"] for message in sent[1:]] == ["Fix nested behavior", "Preserve empty inputs"]
|
||||
else:
|
||||
assert [message["content"] for message in sent[1:]] == [opening, envelope]
|
||||
assert result.messages == original
|
||||
assert completion.await_count == 3
|
||||
assert messages == original
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("p_solve,expected_model", ((0.95, "capable-model"), (0.98, "efficient-model")))
|
||||
async def test_fitted_probability_controls_routing_and_preserves_raw_score(
|
||||
self, mock_router_instance: MagicMock, p_solve: float, expected_model: str
|
||||
) -> None:
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=p_solve)))
|
||||
router: Final = self._router(
|
||||
mock_router_instance,
|
||||
capability_classifier_config={
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.66,
|
||||
"threshold_step": 0.1,
|
||||
"calibration": {
|
||||
"version": "qwen3-haiku45-mini-swe-v1",
|
||||
"slope": 0.1482462649948327,
|
||||
"intercept": 0.1895438369492216,
|
||||
},
|
||||
},
|
||||
)
|
||||
result: Final = await router.async_pre_routing_hook(
|
||||
model="capability-router", request_kwargs={}, messages=[{"role": "user", "content": "Fix the issue"}]
|
||||
)
|
||||
assert result is not None and result.model == expected_model
|
||||
decision: Final = result.routing_decision
|
||||
assert decision is not None
|
||||
assert decision["classifier_p_solve"] == p_solve
|
||||
assert decision["classifier_threshold"] == 0.66
|
||||
assert decision["classifier_calibration_version"] == "qwen3-haiku45-mini-swe-v1"
|
||||
assert 0.65 < decision["classifier_calibrated_p_solve"] < 0.69
|
||||
assert (decision["classifier_calibrated_p_solve"] >= 0.66) == (expected_model == "efficient-model")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("mode", ("json_schema", "json_object"))
|
||||
async def test_response_modes_preserve_the_card_and_validate_the_same_verdict(
|
||||
self, mock_router_instance: MagicMock, mode: str
|
||||
) -> None:
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8)))
|
||||
router: Final = self._router(
|
||||
mock_router_instance,
|
||||
capability_classifier_config={
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.5,
|
||||
"response_format": mode,
|
||||
},
|
||||
)
|
||||
outcome: Final = await router.aclassify("Fix the issue")
|
||||
assert outcome.tier == ComplexityTier.SIMPLE
|
||||
call: Final = mock_router_instance.acompletion.call_args.kwargs
|
||||
system_prompt: Final = call["messages"][0]["content"]
|
||||
assert call["response_format"]["type"] == mode
|
||||
if mode == "json_object":
|
||||
marker: Final = "\n\nReturn exactly one JSON object matching this JSON Schema:\n"
|
||||
assert system_prompt.startswith(CAPABILITY_CLASSIFIER_SYSTEM_PROMPT + marker)
|
||||
schema: Final = json.loads(system_prompt.split(marker)[1])
|
||||
assert schema["required"] == ["crux", "primary_rule", "capability_boundary", "p_solve"]
|
||||
assert schema["additionalProperties"] is False
|
||||
else:
|
||||
assert system_prompt == CAPABILITY_CLASSIFIER_SYSTEM_PROMPT
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response("invalid JSON"))
|
||||
assert (await router.aclassify("Fix another issue")).tier == ComplexityTier.REASONING
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("reply", ("invalid JSON", _capability_reply(p_solve=0.0)))
|
||||
async def test_adaptive_selection_cannot_undo_a_capable_verdict(
|
||||
self, mock_router_instance: MagicMock, reply: str
|
||||
) -> None:
|
||||
from litellm.router_strategy.adaptive_router.bandit import BanditCell
|
||||
from litellm.types.router import RequestType
|
||||
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply))
|
||||
mock_router_instance.model_list = [
|
||||
{"model_name": "efficient-model", "litellm_params": {"input_cost_per_token": 0.000001}},
|
||||
{"model_name": "capable-model", "litellm_params": {"input_cost_per_token": 0.00001}},
|
||||
]
|
||||
mock_router_instance.model_name_to_deployment_indices = {"efficient-model": [0], "capable-model": [1]}
|
||||
router: Final = self._router(
|
||||
mock_router_instance,
|
||||
adaptive=True,
|
||||
adaptive_eligible="all",
|
||||
adaptive_weights={"quality": 0.0, "cost": 1.0},
|
||||
tier_distance_penalty=0.0,
|
||||
tiers={"SIMPLE": ["efficient-model"], "REASONING": ["capable-model"]},
|
||||
)
|
||||
adaptive: Final = router._ensure_adaptive_router()
|
||||
assert adaptive is not None
|
||||
for model in ("efficient-model", "capable-model"):
|
||||
adaptive._cells[(RequestType.GENERAL, model)] = BanditCell(alpha=20.0, beta=1.0)
|
||||
assert router._soft_floor_pick(ComplexityTier.REASONING, "Fix the issue") == "efficient-model"
|
||||
result: Final = await router.async_pre_routing_hook(
|
||||
model="capability-router", request_kwargs={}, messages=[{"role": "user", "content": "Fix the issue"}]
|
||||
)
|
||||
assert result is not None and result.model == "capable-model"
|
||||
assert result.routing_decision is not None
|
||||
assert result.routing_decision["tier"] == "REASONING"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"p_solve,primary_rule,boundary,expected_tier,expected_threshold",
|
||||
[
|
||||
(0.5, "SUP-1", "supported", ComplexityTier.SIMPLE, 0.5),
|
||||
(0.59, "UNC-1", "uncertain", ComplexityTier.REASONING, 0.6),
|
||||
(0.6, "UNC-1", "uncertain", ComplexityTier.SIMPLE, 0.6),
|
||||
(0.59, "none", "unmatched", ComplexityTier.REASONING, 0.6),
|
||||
(0.69, "LIM-1", "unsupported", ComplexityTier.REASONING, 0.7),
|
||||
(0.7, "LIM-1", "unsupported", ComplexityTier.SIMPLE, 0.7),
|
||||
],
|
||||
)
|
||||
async def test_boundary_adjusted_threshold_is_inclusive(
|
||||
self, mock_router_instance, p_solve, primary_rule, boundary, expected_tier, expected_threshold
|
||||
):
|
||||
mock_router_instance.acompletion = AsyncMock(
|
||||
return_value=_llm_response(
|
||||
_capability_reply(p_solve=p_solve, primary_rule=primary_rule, capability_boundary=boundary)
|
||||
)
|
||||
)
|
||||
outcome = await self._router(mock_router_instance).aclassify("do the task")
|
||||
assert outcome.tier == expected_tier
|
||||
assert outcome.cause == "capability_classifier"
|
||||
assert outcome.capability_forecast is not None
|
||||
assert outcome.capability_forecast.threshold == pytest.approx(expected_threshold)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fenced_json_verdict_is_accepted(self, mock_router_instance):
|
||||
reply = _capability_reply(p_solve=0.8)
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(f"```json\n{reply}\n```"))
|
||||
outcome = await self._router(mock_router_instance).aclassify("do the task")
|
||||
assert outcome.tier == ComplexityTier.SIMPLE
|
||||
assert outcome.cause == "capability_classifier"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_decimal_rounding_does_not_break_inclusive_threshold(self, mock_router_instance):
|
||||
config = _capability_router_config(
|
||||
capability_classifier_config={
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.1,
|
||||
"threshold_step": 0.1,
|
||||
}
|
||||
)
|
||||
mock_router_instance.acompletion = AsyncMock(
|
||||
return_value=_llm_response(
|
||||
_capability_reply(p_solve=0.3, primary_rule="LIM-1", capability_boundary="unsupported")
|
||||
)
|
||||
)
|
||||
router = ComplexityRouter(
|
||||
model_name="capability-router",
|
||||
litellm_router_instance=mock_router_instance,
|
||||
complexity_router_config=config,
|
||||
)
|
||||
outcome = await router.aclassify("do the task")
|
||||
assert outcome.capability_forecast is not None
|
||||
assert outcome.capability_forecast.threshold == 0.30000000000000004
|
||||
assert outcome.tier == ComplexityTier.SIMPLE
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_uses_packaged_prompt_schema_and_opening_plus_latest_user_task(self, mock_router_instance):
|
||||
mock_router_instance.acompletion = AsyncMock(
|
||||
return_value=_llm_response(_capability_reply(p_solve=0.8), response_cost=0.002)
|
||||
)
|
||||
router = self._router(mock_router_instance)
|
||||
messages = [
|
||||
{"role": "system", "content": "Never expose this caller instruction to the judge"},
|
||||
{"role": "user", "content": "Build the feature"},
|
||||
{"role": "assistant", "content": "I need more information"},
|
||||
{"role": "user", "content": "Use the existing API"},
|
||||
]
|
||||
|
||||
response = await router.async_pre_routing_hook(model="capability-router", request_kwargs={}, messages=messages)
|
||||
|
||||
assert response.model == "efficient-model"
|
||||
call = mock_router_instance.acompletion.call_args.kwargs
|
||||
assert call["messages"] == [
|
||||
{"role": "system", "content": CAPABILITY_CLASSIFIER_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": "Build the feature"},
|
||||
{"role": "user", "content": "Use the existing API"},
|
||||
]
|
||||
schema = call["response_format"]["json_schema"]["schema"]
|
||||
assert call["response_format"]["json_schema"]["name"] == "CapabilityClassifierDecision"
|
||||
assert call["response_format"]["json_schema"]["strict"] is True
|
||||
assert schema["additionalProperties"] is False
|
||||
assert set(schema["required"]) == {"crux", "primary_rule", "capability_boundary", "p_solve"}
|
||||
assert schema["properties"]["primary_rule"]["enum"] == [
|
||||
"SUP-1",
|
||||
"SUP-2",
|
||||
"SUP-3",
|
||||
"SUP-4",
|
||||
"SUP-5",
|
||||
"UNC-1",
|
||||
"UNC-2",
|
||||
"LIM-1",
|
||||
"LIM-2",
|
||||
"none",
|
||||
]
|
||||
assert call["max_tokens"] == 4096
|
||||
decision = response.routing_decision
|
||||
assert decision["cause"] == "capability_classifier"
|
||||
assert decision["classifier_model"] == "judge-model"
|
||||
assert decision["classifier_cost"] == 0.002
|
||||
assert decision["classifier_crux"] == "complete the requested change"
|
||||
assert decision["classifier_primary_rule"] == "SUP-1"
|
||||
assert decision["classifier_capability_boundary"] == "supported"
|
||||
assert decision["classifier_p_solve"] == 0.8
|
||||
assert decision["classifier_threshold"] == 0.5
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"reply",
|
||||
[
|
||||
"not json",
|
||||
_capability_reply(p_solve=0.9, primary_rule="SUP-1", capability_boundary="unsupported"),
|
||||
'{"crux":"x","primary_rule":"SUP-1","capability_boundary":"supported","p_solve":0.9,"route":"efficient"}',
|
||||
],
|
||||
ids=["malformed", "inconsistent-pair", "extra-field"],
|
||||
)
|
||||
async def test_invalid_verdict_fails_closed_to_capable_tier(self, mock_router_instance, reply):
|
||||
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(reply))
|
||||
outcome = await self._router(mock_router_instance).aclassify("do the task")
|
||||
assert outcome.tier == ComplexityTier.REASONING
|
||||
assert outcome.cause == "capability_classifier_fallback"
|
||||
assert outcome.signals == ("capability-classifier-fallback",)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_classifier_call_failure_fails_closed_to_capable_model(self, mock_router_instance):
|
||||
mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("judge unavailable"))
|
||||
response = await self._router(mock_router_instance).async_pre_routing_hook(
|
||||
model="capability-router",
|
||||
request_kwargs={},
|
||||
messages=[{"role": "user", "content": "do the task"}],
|
||||
)
|
||||
assert response.model == "capable-model"
|
||||
assert response.routing_decision["cause"] == "capability_classifier_fallback"
|
||||
|
||||
|
||||
CUSTOM_TIER_LABELS: Dict[str, str] = {
|
||||
"SIMPLE": "Cheap",
|
||||
"MEDIUM": "Standard",
|
||||
|
|
@ -8342,6 +8858,13 @@ class TestRedactedLoggingDropsPromptText:
|
|||
"score": 0.8,
|
||||
"tier_boundaries": {"simple_medium": 0.15, "medium_complex": 0.35, "complex_reasoning": 0.6},
|
||||
"classifier_model": "claude-haiku",
|
||||
"classifier_crux": "deploy the requested service to k8s",
|
||||
"classifier_primary_rule": "SUP-2",
|
||||
"classifier_capability_boundary": "supported",
|
||||
"classifier_p_solve": 0.8,
|
||||
"classifier_calibrated_p_solve": 0.65,
|
||||
"classifier_calibration_version": "fitted-v1",
|
||||
"classifier_threshold": 0.5,
|
||||
"escalated": True,
|
||||
"tier_litellm_params": {"reasoning_effort": "xhigh"},
|
||||
"signals": ["code (python)"],
|
||||
|
|
@ -8349,7 +8872,15 @@ class TestRedactedLoggingDropsPromptText:
|
|||
"escalation_keyword": "LITELLM ESCALATE",
|
||||
}
|
||||
kept = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=full)
|
||||
assert set(full) - set(kept) == {"signals", "matched_keyword", "escalation_keyword"}
|
||||
assert set(full) - set(kept) == {
|
||||
"signals",
|
||||
"matched_keyword",
|
||||
"escalation_keyword",
|
||||
"classifier_crux",
|
||||
}
|
||||
assert kept["classifier_p_solve"] == 0.8
|
||||
assert kept["classifier_calibrated_p_solve"] == 0.65
|
||||
assert kept["classifier_calibration_version"] == "fitted-v1"
|
||||
assert kept["tier_litellm_params"] == {"reasoning_effort": "xhigh"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -8463,8 +8994,10 @@ class TestContextAwareClassifier:
|
|||
assert messages == original_messages
|
||||
assert (claude_kwargs, compared_kwargs) == original_kwargs
|
||||
calls: Final = tuple(call.kwargs["messages"] for call in dependency.acompletion.await_args_list)
|
||||
assert calls[0][0]["content"] == calls[1][0]["content"] == classification_system_prompt(
|
||||
router.config.classifier_context_window_size
|
||||
assert (
|
||||
calls[0][0]["content"]
|
||||
== calls[1][0]["content"]
|
||||
== classification_system_prompt(router.config.classifier_context_window_size)
|
||||
)
|
||||
payloads: Final = (calls[0][1]["content"], calls[1][1]["content"])
|
||||
for payload, expected_system in zip(payloads, (False, forwards_system)):
|
||||
|
|
@ -13653,11 +14186,7 @@ class TestHealthFallbackDispatch:
|
|||
"api_key": "test-only",
|
||||
"api_base": f"https://{name}.test{base_suffix}",
|
||||
**({"tags": [name]} if tagged else {}),
|
||||
**(
|
||||
{"max_budget": 1.0, "budget_duration": "1d"}
|
||||
if budgeted and name == "primary"
|
||||
else {}
|
||||
),
|
||||
**({"max_budget": 1.0, "budget_duration": "1d"} if budgeted and name == "primary" else {}),
|
||||
},
|
||||
"model_info": {"id": f"{name}-id"},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -5,15 +5,20 @@ the implicit `"default"` group driven by the router's top-level
|
|||
`routing_strategy` / `routing_strategy_args`.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.router import RoutingGroup, RoutingStrategy
|
||||
from litellm.utils import Rules, function_setup
|
||||
|
||||
|
||||
def _model_list():
|
||||
|
|
@ -954,6 +959,223 @@ def test_sync_pass_through_specific_deployment_runs_the_override_pre_call_check(
|
|||
assert plain["model_info"]["id"] == "deploy-3"
|
||||
|
||||
|
||||
def _two_deployment_model_list(**d1_params: object) -> list[dict[str, object]]:
|
||||
return [
|
||||
{
|
||||
"model_name": "grp",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test-1", "mock_response": "ok", **d1_params},
|
||||
"model_info": {"id": "d1"},
|
||||
},
|
||||
{
|
||||
"model_name": "grp",
|
||||
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test-2", "mock_response": "ok"},
|
||||
"model_info": {"id": "d2"},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _proxy_shaped_request(**data: object) -> dict[str, object]:
|
||||
"""The proxy builds the request's `Logging` object before it hands the call to the router."""
|
||||
logging_obj, kwargs = function_setup(
|
||||
"acompletion",
|
||||
Rules(),
|
||||
datetime.datetime.now(),
|
||||
litellm_call_id=str(uuid.uuid4()),
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
**data,
|
||||
)
|
||||
return {**kwargs, "litellm_logging_obj": logging_obj}
|
||||
|
||||
|
||||
async def _async_override_pick(router: Router, strategy: str) -> str:
|
||||
deployment = await router.async_get_available_deployment(
|
||||
"grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy)
|
||||
)
|
||||
return deployment["model_info"]["id"]
|
||||
|
||||
|
||||
def _sync_override_pick(router: Router, strategy: str) -> str:
|
||||
deployment = router.get_available_deployment(
|
||||
"grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy)
|
||||
)
|
||||
return deployment["model_info"]["id"]
|
||||
|
||||
|
||||
def _in_flight(router: Router, deployment_id: str) -> int | None:
|
||||
return router.cache.get_cache(f"grp_request_count:{deployment_id}")
|
||||
|
||||
|
||||
async def _async_wait_until(predicate: Callable[[], bool]) -> None:
|
||||
for _ in range(100):
|
||||
if predicate():
|
||||
return
|
||||
await asyncio.sleep(0.02)
|
||||
raise AssertionError("lifecycle callback never reached the override selector")
|
||||
|
||||
|
||||
def _sync_wait_until(predicate: Callable[[], bool]) -> None:
|
||||
for _ in range(100):
|
||||
if predicate():
|
||||
return
|
||||
time.sleep(0.02)
|
||||
raise AssertionError("lifecycle callback never reached the override selector")
|
||||
|
||||
|
||||
def _selector_is_not_global(selector: CustomLogger) -> bool:
|
||||
global_lists = (
|
||||
litellm.callbacks,
|
||||
litellm.input_callback,
|
||||
litellm.success_callback,
|
||||
litellm.failure_callback,
|
||||
litellm._async_success_callback,
|
||||
litellm._async_failure_callback,
|
||||
)
|
||||
return not any(cb is selector for cbs in global_lists for cb in cbs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_least_busy_override_sees_the_overriding_request_in_flight():
|
||||
router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0)
|
||||
|
||||
stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True))
|
||||
busy = stream._hidden_params["model_id"]
|
||||
idle = "d2" if busy == "d1" else "d1"
|
||||
assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle]
|
||||
|
||||
async for _ in stream:
|
||||
pass
|
||||
await _async_wait_until(lambda: _in_flight(router, busy) == 0)
|
||||
assert await _async_override_pick(router, "least-busy") == "d1"
|
||||
assert _selector_is_not_global(router._override_selectors["least-busy"])
|
||||
|
||||
|
||||
def test_sync_least_busy_override_sees_the_overriding_request_in_flight():
|
||||
router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0)
|
||||
|
||||
stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True))
|
||||
busy = stream._hidden_params["model_id"]
|
||||
idle = "d2" if busy == "d1" else "d1"
|
||||
assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle]
|
||||
|
||||
for _ in stream:
|
||||
pass
|
||||
_sync_wait_until(lambda: _in_flight(router, busy) == 0)
|
||||
assert _sync_override_pick(router, "least-busy") == "d1"
|
||||
assert _selector_is_not_global(router._override_selectors["least-busy"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_least_busy_override_releases_the_slot_when_the_overriding_request_fails():
|
||||
router = Router(
|
||||
model_list=_two_deployment_model_list(mock_response="litellm.InternalServerError"),
|
||||
routing_strategy="simple-shuffle",
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.InternalServerError):
|
||||
await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy"))
|
||||
|
||||
await _async_wait_until(lambda: _in_flight(router, "d1") == 0)
|
||||
assert await _async_override_pick(router, "least-busy") == "d1"
|
||||
assert _selector_is_not_global(router._override_selectors["least-busy"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_latency_based_override_learns_from_the_overriding_requests():
|
||||
router = Router(
|
||||
model_list=_two_deployment_model_list(mock_delay=0.05), routing_strategy="simple-shuffle", num_retries=0
|
||||
)
|
||||
|
||||
def samples(deployment_id: str) -> list[float]:
|
||||
recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency", [])
|
||||
return [latency for latency in recorded if latency > 0]
|
||||
|
||||
async def overriding_call() -> str:
|
||||
sampled_before = {"d1": len(samples("d1")), "d2": len(samples("d2"))}
|
||||
response = await router.acompletion(
|
||||
**_proxy_shaped_request(model="grp", routing_strategy="latency-based-routing")
|
||||
)
|
||||
deployment_id = response._hidden_params["model_id"]
|
||||
await _async_wait_until(lambda: len(samples(deployment_id)) > sampled_before[deployment_id])
|
||||
return deployment_id
|
||||
|
||||
served = [await overriding_call() for _ in range(6)]
|
||||
|
||||
assert "d1" in served
|
||||
assert served[2:] == ["d2"] * 4
|
||||
assert _selector_is_not_global(router._override_selectors["latency-based-routing"])
|
||||
|
||||
|
||||
def test_override_selector_is_bound_only_to_the_request_that_asked_for_it():
|
||||
router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle")
|
||||
overriding = _proxy_shaped_request(model="grp", routing_strategy="least-busy")
|
||||
plain = _proxy_shaped_request(model="grp")
|
||||
|
||||
router.get_available_deployment("grp", request_kwargs=overriding)
|
||||
router.get_available_deployment("grp", request_kwargs=overriding)
|
||||
router.get_available_deployment("grp", request_kwargs=plain)
|
||||
|
||||
selector = router._override_selectors["least-busy"]
|
||||
bound = overriding["litellm_logging_obj"]
|
||||
for callbacks in (
|
||||
bound.dynamic_input_callbacks,
|
||||
bound.dynamic_success_callbacks,
|
||||
bound.dynamic_async_success_callbacks,
|
||||
bound.dynamic_failure_callbacks,
|
||||
bound.dynamic_async_failure_callbacks,
|
||||
):
|
||||
assert callbacks == [selector]
|
||||
unbound = plain["litellm_logging_obj"]
|
||||
assert unbound.dynamic_input_callbacks is None and unbound.dynamic_success_callbacks is None
|
||||
assert unbound.dynamic_failure_callbacks is None and unbound.dynamic_async_failure_callbacks is None
|
||||
|
||||
|
||||
def test_override_matching_the_router_strategy_is_not_bound_twice():
|
||||
router = Router(model_list=_two_deployment_model_list(), routing_strategy="least-busy")
|
||||
request = _proxy_shaped_request(model="grp", routing_strategy="least-busy")
|
||||
|
||||
router.get_available_deployment("grp", request_kwargs=request)
|
||||
|
||||
assert request["litellm_logging_obj"].dynamic_input_callbacks is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_override_matching_a_routing_group_strategy_records_each_request_once():
|
||||
router = Router(
|
||||
model_list=_two_deployment_model_list(),
|
||||
routing_strategy="simple-shuffle",
|
||||
routing_groups=[RoutingGroup(group_name="lat", models=["grp"], routing_strategy="latency-based-routing")],
|
||||
num_retries=0,
|
||||
)
|
||||
request = _proxy_shaped_request(model="grp", routing_strategy="latency-based-routing")
|
||||
assert router._globally_registered_strategies() == {"simple-shuffle", "latency-based-routing"}
|
||||
|
||||
response = await router.acompletion(**request)
|
||||
deployment_id = response._hidden_params["model_id"]
|
||||
await _async_wait_until(lambda: (router.cache.get_cache("grp_map") or {}).get(deployment_id) is not None)
|
||||
|
||||
assert len(router.cache.get_cache("grp_map")[deployment_id]["latency"]) == 1
|
||||
assert request["litellm_logging_obj"].dynamic_success_callbacks is None
|
||||
|
||||
|
||||
def test_bind_override_selector_to_request_binds_once_and_ignores_requests_without_logging():
|
||||
router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle")
|
||||
selector = router._get_override_strategy_selector("least-busy")
|
||||
request = _proxy_shaped_request(model="grp", routing_strategy="least-busy")
|
||||
request["litellm_logging_obj"].dynamic_success_callbacks = ["langfuse"]
|
||||
|
||||
router._bind_override_selector_to_request("least-busy", selector, request)
|
||||
router._bind_override_selector_to_request("least-busy", selector, request)
|
||||
router._bind_override_selector_to_request("least-busy", selector, None)
|
||||
router._bind_override_selector_to_request("least-busy", selector, {"model": "grp"})
|
||||
|
||||
logging_obj = request["litellm_logging_obj"]
|
||||
assert logging_obj.dynamic_success_callbacks == ["langfuse", selector]
|
||||
assert logging_obj.dynamic_input_callbacks == [selector]
|
||||
assert logging_obj.dynamic_async_failure_callbacks == [selector]
|
||||
assert _selector_is_not_global(selector)
|
||||
|
||||
|
||||
def _quality_group(strategy="latency-based-routing"):
|
||||
return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}]
|
||||
|
||||
|
|
|
|||
|
|
@ -214,6 +214,22 @@ def test_config_check_ignores_the_model_entirely():
|
|||
},
|
||||
(("a", "tier"), ("clf", "classifier")),
|
||||
),
|
||||
(
|
||||
{
|
||||
"model": "auto_router/complexity_router",
|
||||
"complexity_router_config": {
|
||||
"tiers": {"SIMPLE": "a", "REASONING": "b"},
|
||||
"classifier_type": "capability",
|
||||
"classifier_llm_config": {"model": "clf"},
|
||||
"capability_classifier_config": {
|
||||
"efficient_tier": "SIMPLE",
|
||||
"capable_tier": "REASONING",
|
||||
"base_threshold": 0.5,
|
||||
},
|
||||
},
|
||||
},
|
||||
(("a", "tier"), ("b", "tier"), ("clf", "classifier")),
|
||||
),
|
||||
(
|
||||
{
|
||||
"model": "auto_router/complexity_router",
|
||||
|
|
|
|||
|
|
@ -15975,6 +15975,60 @@ async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warni
|
|||
assert any("circuit breaker is open" in record.getMessage() for record in caplog.records)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the_group():
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {"model": "openai/gpt-5-nano", "api_key": "k", "weight": 0, "mock_response": "A"},
|
||||
"model_info": {"id": "gpt-5-nano"},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k", "weight": 1, "mock_response": "B"},
|
||||
"model_info": {"id": "gpt-5-mini-dep"},
|
||||
},
|
||||
],
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
by_group = await router.acompletion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}])
|
||||
by_id = await router.acompletion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}])
|
||||
|
||||
assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep"
|
||||
assert by_group.choices[0].message.content == "B"
|
||||
assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep"
|
||||
|
||||
|
||||
def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_deployment_id():
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {"model": "openai/gpt-5-nano", "api_key": "k", "weight": 0, "mock_response": "A"},
|
||||
"model_info": {"id": "gpt-5-nano"},
|
||||
},
|
||||
{
|
||||
"model_name": "gpt-5-nano",
|
||||
"litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k", "weight": 1, "mock_response": "B"},
|
||||
"model_info": {"id": "gpt-5-mini-dep"},
|
||||
},
|
||||
],
|
||||
routing_strategy="simple-shuffle",
|
||||
)
|
||||
|
||||
with patch.object(router, "routing_strategy_pre_call_checks") as pre_call_checks:
|
||||
by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}])
|
||||
assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep"
|
||||
pre_call_checks.assert_called_once()
|
||||
assert pre_call_checks.call_args.kwargs["deployment"]["model_info"]["id"] == "gpt-5-mini-dep"
|
||||
|
||||
by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}])
|
||||
assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep"
|
||||
pre_call_checks.assert_called_once()
|
||||
|
||||
|
||||
class TestMemberAutoRouterInference:
|
||||
@pytest.fixture(autouse=True)
|
||||
def runtime(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
|
|
|||
|
|
@ -433,6 +433,15 @@ const ClassificationMethodConfig: React.FC<ClassificationMethodConfigProps> = ({
|
|||
});
|
||||
};
|
||||
|
||||
if (classifierType === "capability") {
|
||||
return (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
This router uses capability forecasting. Configure its classifier, threshold, and calibration through YAML or
|
||||
the API. Saving preserves those settings
|
||||
</p>
|
||||
);
|
||||
}
|
||||
|
||||
if (classifierType === "llm_v2") {
|
||||
return (
|
||||
<div className="rounded-md border p-4 text-sm">
|
||||
|
|
|
|||
|
|
@ -143,7 +143,14 @@ export interface ClassifierLLMConfig {
|
|||
system_prompt?: string;
|
||||
}
|
||||
|
||||
export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "llm_v2" | "heuristic_first" | "hybrid";
|
||||
export type ClassifierType =
|
||||
| "heuristic"
|
||||
| "heuristic_v2"
|
||||
| "llm"
|
||||
| "llm_v2"
|
||||
| "heuristic_first"
|
||||
| "hybrid"
|
||||
| "capability";
|
||||
|
||||
/**
|
||||
* Whether this router can call classifier_llm_config.model. Mirrors the backend's
|
||||
|
|
@ -151,7 +158,7 @@ export type ClassifierType = "heuristic" | "heuristic_v2" | "llm" | "llm_v2" | "
|
|||
* control and payload key, so a new chaining type cannot strip knobs the operator set.
|
||||
*/
|
||||
export const usesLlmClassifier = (classifierType: ClassifierType): boolean =>
|
||||
["llm", "llm_v2", "heuristic_first", "hybrid"].includes(classifierType);
|
||||
(["llm", "llm_v2", "heuristic_first", "hybrid", "capability"] as const).some((type) => type === classifierType);
|
||||
|
||||
export type ClassifierFallback = "heuristic" | "default_model";
|
||||
|
||||
|
|
@ -176,7 +183,8 @@ export const heuristicScoringRoleFor = (
|
|||
classifierType: ClassifierType,
|
||||
classifierFallback: ClassifierFallback | undefined,
|
||||
): HeuristicScoringRole => {
|
||||
if (classifierType === "heuristic_v2" || classifierType === "llm_v2") return "never";
|
||||
if (classifierType === "heuristic_v2" || classifierType === "llm_v2" || classifierType === "capability")
|
||||
return "never";
|
||||
if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid")
|
||||
return "decides";
|
||||
return (classifierFallback ?? DEFAULT_CLASSIFIER_FALLBACK) === "heuristic" ? "fallback_only" : "never";
|
||||
|
|
|
|||
|
|
@ -470,7 +470,10 @@ const classifierWireFields = (
|
|||
>,
|
||||
): Partial<ComplexityRouterConfigPayload> => ({
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierLlmConfig && { classifier_llm_config: normalizeClassifierLlmConfig(classifierLlmConfig) }),
|
||||
classifierLlmConfig && {
|
||||
classifier_llm_config:
|
||||
effectiveType === "capability" ? classifierLlmConfig : normalizeClassifierLlmConfig(classifierLlmConfig),
|
||||
}),
|
||||
...(usesLlmClassifier(effectiveType) &&
|
||||
classifierFallback !== undefined && { classifier_fallback: classifierFallback }),
|
||||
...(effectiveType === "heuristic_first" &&
|
||||
|
|
|
|||
|
|
@ -164,6 +164,32 @@ const STORED_LLM = {
|
|||
classifier_context_per_turn_chars: 300,
|
||||
};
|
||||
|
||||
describe("capability classifier configuration", () => {
|
||||
it("preserves the judge and calibrated policy through an untouched dashboard edit", () => {
|
||||
const stored = {
|
||||
tiers: { SIMPLE: ["efficient-model"], REASONING: ["capable-model"] },
|
||||
classifier_type: "capability" as const,
|
||||
classifier_llm_config: { model: "judge", timeout_ms: 30000, temperature: 0 },
|
||||
capability_classifier_config: {
|
||||
efficient_tier: "SIMPLE",
|
||||
capable_tier: "REASONING",
|
||||
base_threshold: 0.66,
|
||||
max_output_tokens: 512,
|
||||
response_format: "json_object",
|
||||
calibration: { version: "fitted-v1", slope: 0.15, intercept: 0.19 },
|
||||
},
|
||||
};
|
||||
const hydrated = hydrateComplexityRouterConfig(stored, null);
|
||||
const saved = buildUpdatedComplexityRouterConfig(stored, hydrated);
|
||||
|
||||
expect(saved.classifier_type).toBe("capability");
|
||||
expect(saved.classifier_llm_config).toEqual(stored.classifier_llm_config);
|
||||
expect(saved.capability_classifier_config).toEqual(stored.capability_classifier_config);
|
||||
expect(saved).not.toHaveProperty("classification_prompt");
|
||||
expect(saved).not.toHaveProperty("custom_dimensions");
|
||||
});
|
||||
});
|
||||
|
||||
describe("buildUpdatedComplexityRouterConfig classifier context window", () => {
|
||||
it("round-trips an untouched edit without changing the classifier context values", () => {
|
||||
const formValue = {
|
||||
|
|
|
|||
75
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
75
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -25328,6 +25328,57 @@ export interface components {
|
|||
*/
|
||||
status: "cancelled";
|
||||
};
|
||||
/** CapabilityCalibrationConfig */
|
||||
CapabilityCalibrationConfig: {
|
||||
/** Intercept */
|
||||
intercept: number;
|
||||
/** Slope */
|
||||
slope: number;
|
||||
/** Version */
|
||||
version: string;
|
||||
};
|
||||
/**
|
||||
* CapabilityClassifierConfig
|
||||
* @description Switchyard-compatible probability threshold policy for two model tiers.
|
||||
*/
|
||||
CapabilityClassifierConfig: {
|
||||
/**
|
||||
* Base Threshold
|
||||
* @description Lowest p_solve that routes a supported task to efficient_tier
|
||||
*/
|
||||
base_threshold: number;
|
||||
/** @description Optional versioned sigmoid calibration fitted for this judge, capability card, efficient model, and execution setup. Applies sigmoid(slope * logit(clip(p_solve, 1e-6, 1-1e-6)) + intercept) before the threshold policy. Omit to route on the raw forecast. */
|
||||
calibration?: components["schemas"]["CapabilityCalibrationConfig"] | null;
|
||||
/**
|
||||
* Capable Tier
|
||||
* @description Higher, fail-closed tier used below the adjusted threshold or when the classifier verdict is unavailable
|
||||
*/
|
||||
capable_tier: string;
|
||||
/**
|
||||
* Efficient Tier
|
||||
* @description Tier used when the efficient model's forecasted solve probability meets the adjusted threshold
|
||||
*/
|
||||
efficient_tier: string;
|
||||
/**
|
||||
* Max Output Tokens
|
||||
* @description Maximum completion tokens available to the capability classifier verdict
|
||||
* @default 4096
|
||||
*/
|
||||
max_output_tokens: number;
|
||||
/**
|
||||
* Response Format
|
||||
* @description Use json_object for judges without strict JSON Schema support. This appends the verdict schema to the packaged system prompt; both modes validate the returned verdict identically.
|
||||
* @default json_schema
|
||||
* @enum {string}
|
||||
*/
|
||||
response_format: "json_schema" | "json_object";
|
||||
/**
|
||||
* Threshold Step
|
||||
* @description Amount added once for uncertain or unmatched verdicts and twice for unsupported verdicts
|
||||
* @default 0
|
||||
*/
|
||||
threshold_step: number;
|
||||
};
|
||||
/** ChatCompletionAnnotation */
|
||||
ChatCompletionAnnotation: {
|
||||
/**
|
||||
|
|
@ -35615,6 +35666,8 @@ export interface components {
|
|||
adaptive_eligible: "all" | "classified_tier";
|
||||
/** @description Quality vs cost weights for adaptive selection (used when adaptive=True) */
|
||||
adaptive_weights?: components["schemas"]["AdaptiveRouterWeights"];
|
||||
/** @description Probability threshold policy required when classifier_type is 'capability'. The classifier forecasts p_solve for efficient_tier, adjusts base_threshold using the capability-card boundary, and otherwise routes to capable_tier */
|
||||
capability_classifier_config?: components["schemas"]["CapabilityClassifierConfig"] | null;
|
||||
/**
|
||||
* Classification Examples
|
||||
* @description Replaces the calibration examples of the LLM classifier rubric, and nothing else. Written as example lines only: the router renders the 'Calibration examples:' heading above them, after the per-tier bullets. Requires an LLM classifier and cannot be combined with classifier_llm_config.system_prompt. With built-in tiers the rubric preset still supplies the tier criteria and, unless classification_prompt replaces them, the classification instructions; a custom tier set ships no examples of its own, so the section renders only when this is set.
|
||||
|
|
@ -35662,7 +35715,7 @@ export interface components {
|
|||
* @enum {string}
|
||||
*/
|
||||
classifier_fallback: "heuristic" | "default_model";
|
||||
/** @description Configuration for the LLM classifier; required when classifier_type is 'llm', 'heuristic_first' or 'hybrid' */
|
||||
/** @description Configuration for the LLM classifier; required when classifier_type is 'llm', 'capability', 'heuristic_first' or 'hybrid' */
|
||||
classifier_llm_config?: components["schemas"]["ClassifierLLMConfig"] | null;
|
||||
/**
|
||||
* Classifier Plugin
|
||||
|
|
@ -35677,11 +35730,11 @@ export interface components {
|
|||
classifier_plugin_timeout_ms: number;
|
||||
/**
|
||||
* Classifier Type
|
||||
* @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM call, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary
|
||||
* @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM tier-selection call, a Switchyard-compatible capability forecast, a joint task-demand and capability forecast, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary
|
||||
* @default heuristic
|
||||
* @enum {string}
|
||||
*/
|
||||
classifier_type: "heuristic" | "heuristic_v2" | "llm" | "llm_v2" | "custom" | "heuristic_first" | "hybrid";
|
||||
classifier_type: "heuristic" | "heuristic_v2" | "llm" | "capability" | "llm_v2" | "custom" | "heuristic_first" | "hybrid";
|
||||
/**
|
||||
* Code Keywords
|
||||
* @description Keywords indicating code-related content
|
||||
|
|
@ -37081,11 +37134,25 @@ export interface components {
|
|||
* Cause
|
||||
* @enum {string}
|
||||
*/
|
||||
cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "llm_v2_classifier" | "llm_v2_fallback" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "modality_pin_override" | "health_failover" | "health_default_fallback" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit";
|
||||
cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "capability_classifier" | "llm_v2_classifier" | "llm_v2_fallback" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "capability_classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "modality_pin_override" | "health_failover" | "health_default_fallback" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit";
|
||||
/** Classifier Calibrated P Solve */
|
||||
classifier_calibrated_p_solve?: number;
|
||||
/** Classifier Calibration Version */
|
||||
classifier_calibration_version?: string;
|
||||
/** Classifier Capability Boundary */
|
||||
classifier_capability_boundary?: string;
|
||||
/** Classifier Cost */
|
||||
classifier_cost?: number;
|
||||
/** Classifier Crux */
|
||||
classifier_crux?: string;
|
||||
/** Classifier Model */
|
||||
classifier_model?: string;
|
||||
/** Classifier P Solve */
|
||||
classifier_p_solve?: number;
|
||||
/** Classifier Primary Rule */
|
||||
classifier_primary_rule?: string;
|
||||
/** Classifier Threshold */
|
||||
classifier_threshold?: number;
|
||||
/** Context Escalated */
|
||||
context_escalated?: boolean;
|
||||
/** Context Escalation Original Tier */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue