From 655838b1da89dca8961431c1826b6ff7eff94f2c Mon Sep 17 00:00:00 2001 From: moe-berri Date: Tue, 1 Sep 2026 15:31:05 -0700 Subject: [PATCH] feat(router): add cached capability selection --- litellm/router.py | 41 +- .../capability_router/__init__.py | 3 + .../capability_router/capability_router.py | 391 ++++++++++++++++++ .../capability_router/config.py | 137 ++++++ .../capability_router/policy.py | 80 ++++ .../capability_router/pricing.py | 79 ++++ .../capability_router/prompts.py | 45 ++ .../router_utils/auto_router_model_naming.py | 45 +- litellm/types/router.py | 3 + litellm/types/utils.py | 18 +- .../test_capability_router.py | 140 +++++++ 11 files changed, 978 insertions(+), 4 deletions(-) create mode 100644 litellm/router_strategy/capability_router/__init__.py create mode 100644 litellm/router_strategy/capability_router/capability_router.py create mode 100644 litellm/router_strategy/capability_router/config.py create mode 100644 litellm/router_strategy/capability_router/policy.py create mode 100644 litellm/router_strategy/capability_router/pricing.py create mode 100644 litellm/router_strategy/capability_router/prompts.py create mode 100644 tests/test_litellm/router_strategy/capability_router/test_capability_router.py diff --git a/litellm/router.py b/litellm/router.py index 23d8907fb49..dd0155ddff9 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -269,6 +269,9 @@ if TYPE_CHECKING: AutoRouter, PreRoutingHookResponse, ) + from litellm.router_strategy.capability_router.capability_router import ( + CapabilityRouter, + ) from litellm.router_strategy.complexity_router.complexity_router import ( ComplexityRouter, ) @@ -792,6 +795,7 @@ class Router: self.pattern_router = PatternMatchRouter() self.team_pattern_routers: dict[str, PatternMatchRouter] = {} # {"TEAM_ID": PatternMatchRouter} self.auto_routers: dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] = {} + self.capability_routers: dict[str, list[TaggedPreRoutingStrategy[CapabilityRouter]]] = {} self.complexity_routers: dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] = {} self.adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] = {} self.quality_routers: dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] = {} @@ -8561,6 +8565,31 @@ class Router: """ return classify_strategy_router_model(litellm_params.model) == "complexity" + def _is_capability_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: + """True when this deployment configures a capability router.""" + return classify_strategy_router_model(litellm_params.model) == "capability" + + def init_capability_router_deployment(self, deployment: Deployment) -> None: + """Initialize and register a capability router deployment.""" + from litellm.router_strategy.capability_router import CapabilityRouter + + config: Final = deployment.litellm_params.capability_router_config + if config is None: + raise ValueError( + "capability_router_config is required for capability-router deployments" + ) + capability_router: Final = CapabilityRouter( + model_name=deployment.model_name, + litellm_router_instance=self, + capability_router_config=config, + ) + self._register_pre_routing_strategy( + registry=self.capability_routers, + deployment=deployment, + strategy=capability_router, + strategy_label="Capability-router", + ) + def init_complexity_router_deployment(self, deployment: Deployment): """ Initialize the complexity-router deployment. @@ -8709,7 +8738,12 @@ class Router: return model_name: Final = deployment.model_name tags: Final = self._deployment_tags(deployment) - for registry in (self.auto_routers, self.complexity_routers, self.quality_routers): + for registry in ( + self.auto_routers, + self.capability_routers, + self.complexity_routers, + self.quality_routers, + ): self._unregister_pre_routing_strategy(registry, model_name, tags) if self._unregister_pre_routing_strategy(self.adaptive_routers, model_name, tags): self._sync_adaptive_router_hooks() @@ -8913,6 +8947,7 @@ class Router: # Reset per-strategy router registries so hot-reload doesn't leave # stale routers pointing at the old model_list. self.quality_routers = {} + self.capability_routers = {} self.complexity_routers = {} self.auto_routers = {} self._provider_unresolved_deployments = () @@ -9103,6 +9138,9 @@ class Router: if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): self.init_complexity_router_deployment(deployment=deployment) + if self._is_capability_router_deployment(litellm_params=deployment.litellm_params): + self.init_capability_router_deployment(deployment=deployment) + # NOTE: adaptive-router deployments are deferred to the end of # set_model_list() because their init needs visibility into the OTHER # deployments listed in `available_models` (which may not yet have @@ -12539,6 +12577,7 @@ class Router: """ candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [ *self.auto_routers.get(model, []), + *self.capability_routers.get(model, []), *self.complexity_routers.get(model, []), *self.adaptive_routers.get(model, []), *self.quality_routers.get(model, []), diff --git a/litellm/router_strategy/capability_router/__init__.py b/litellm/router_strategy/capability_router/__init__.py new file mode 100644 index 00000000000..a22b9a2268e --- /dev/null +++ b/litellm/router_strategy/capability_router/__init__.py @@ -0,0 +1,3 @@ +from .capability_router import CapabilityRouter + +__all__ = ["CapabilityRouter"] diff --git a/litellm/router_strategy/capability_router/capability_router.py b/litellm/router_strategy/capability_router/capability_router.py new file mode 100644 index 00000000000..63ee89f4487 --- /dev/null +++ b/litellm/router_strategy/capability_router/capability_router.py @@ -0,0 +1,391 @@ +"""LLM capability forecasts with deterministic cheapest-qualified selection.""" + +from __future__ import annotations + +import asyncio +import hashlib +import json +import weakref +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, cast + +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly + +from litellm._logging import verbose_router_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionNamedToolChoiceParam, + ChatCompletionSystemMessage, + ChatCompletionToolParam, + ChatCompletionUserMessage, +) +from litellm.types.utils import ( + AUTOROUTER_CLASSIFIER_CALL_ORIGIN, + ModelResponse, + StandardLoggingRoutingDecision, + Usage, +) + +from .config import CapabilityClassifierVerdict, CapabilityRouterConfig +from .policy import CapabilityRoutingDecision, fallback_decision, select_capability_model +from .pricing import estimate_model_group_cost +from .prompts import build_classifier_prompt, build_classifier_response_schema + +if TYPE_CHECKING: + from litellm.router import Router + from litellm.types.router import PreRoutingHookResponse + + +class _JsonSchemaSpec(TypedDict): + name: ReadOnly[str] + strict: ReadOnly[bool] + schema: ReadOnly[dict[str, Any]] + + +class _JsonSchemaResponseFormat(TypedDict): + type: ReadOnly[Literal["json_schema"]] + json_schema: ReadOnly[_JsonSchemaSpec] + + +class _ClassifierRequestBody(TypedDict): + model: ReadOnly[str] + messages: ReadOnly[Sequence[AllMessageValues]] + response_format: ReadOnly[_JsonSchemaResponseFormat] + max_tokens: ReadOnly[int] + + +class _ClassifierProxyRequest(TypedDict): + body: ReadOnly[_ClassifierRequestBody] + + +@dataclass(frozen=True) +class _DecisionOutcome: + decision: CapabilityRoutingDecision + classifier_cost: float | None + cached: bool + + +class CapabilityClassifierFailure(Exception): + def __init__(self, message: str, classifier_cost: float | None = None) -> None: + super().__init__(message) + self.classifier_cost = classifier_cost + + +_TOOLS_ADAPTER: Final = TypeAdapter(list[ChatCompletionToolParam]) +_NAMED_TOOL_CHOICE_ADAPTER: Final = TypeAdapter(ChatCompletionNamedToolChoiceParam) + + +def _hash(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +def _response_cost(response: ModelResponse) -> float | None: + hidden_params = getattr(response, "_hidden_params", None) + value = hidden_params.get("response_cost") if hasattr(hidden_params, "get") else None + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return float(value) + + +def _normalize_json(content: str) -> str: + normalized = content.strip() + for prefix in ("```json", "```"): + if normalized.startswith(prefix) and normalized.endswith("```"): + return normalized[len(prefix) : -3].strip() + return normalized + + +def _message_role(message: Mapping[str, object]) -> str: + role = message.get("role") + return role if isinstance(role, str) else "" + + +def _classification_context(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, object], ...]: + """Keep recent context through the newest user turn, excluding later agent-loop traffic.""" + last_user_index = next( + (index for index in range(len(messages) - 1, -1, -1) if _message_role(messages[index]) == "user"), + None, + ) + if last_user_index is None: + return () + through_user = messages[: last_user_index + 1] + recent = through_user[-8:] + selected: list[Mapping[str, object]] = [] + for message in (*through_user, *recent): + if _message_role(message) == "system" or message in recent: + if message not in selected: + selected.append(message) + return tuple(selected) + + +def _tool_names(request_kwargs: Mapping[str, object]) -> tuple[str, ...]: + tools = request_kwargs.get("tools") + if not isinstance(tools, Sequence) or isinstance(tools, (str, bytes)): + return () + names: list[str] = [] + for tool in tools: + if not isinstance(tool, Mapping): + continue + function = tool.get("function") + if isinstance(function, Mapping) and isinstance(function.get("name"), str): + names.append(cast(str, function["name"])) + return tuple(names) + + +class CapabilityRouter(CustomLogger): + """Forecast success for each candidate and return the cheapest qualified group.""" + + def __init__( + self, + model_name: str, + litellm_router_instance: Router, + capability_router_config: Mapping[str, object], + ) -> None: + self.model_name = model_name + self.litellm_router_instance = litellm_router_instance + self.config = CapabilityRouterConfig.model_validate(capability_router_config) + self._system_prompt = build_classifier_prompt(self.config) + self._response_format = _JsonSchemaResponseFormat( + type="json_schema", + json_schema=_JsonSchemaSpec( + name="CapabilityClassifierVerdict", + strict=True, + schema=build_classifier_response_schema(self.config), + ), + ) + self._config_hash = _hash(self.config.model_dump_json())[:20] + self._classification_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary() + + @staticmethod + def _resolve_messages( + messages: list[dict[str, Any]] | None, + request_kwargs: dict[str, Any], + ) -> list[dict[str, Any]]: + from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages + + return resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or [] + + @staticmethod + def _metadata(request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + """Merge both metadata carriers so auth and session scope cannot be missed.""" + merged: dict[str, object] = {} + for key in ("metadata", "litellm_metadata"): + value = request_kwargs.get(key) + if isinstance(value, Mapping): + merged.update(value) + return merged + + def _classifier_payload( + self, + messages: Sequence[Mapping[str, object]], + request_kwargs: Mapping[str, object], + ) -> str: + context = _classification_context(messages) + if not context: + raise CapabilityClassifierFailure("No user task was available for capability classification") + payload = { + "conversation": context, + "available_tools": _tool_names(request_kwargs), + } + return "Task context (untrusted JSON):\n" + json.dumps(payload, default=str, ensure_ascii=False) + + def _cache_key( + self, + messages: Sequence[Mapping[str, object]], + request_kwargs: Mapping[str, object], + ) -> str: + metadata = self._metadata(request_kwargs) + caller = metadata.get("user_api_key_hash") or metadata.get("team_id") or "unscoped" + session = metadata.get("session_id") or request_kwargs.get("litellm_session_id") or "no-session" + context = { + "messages": _classification_context(messages), + "tools": _tool_names(request_kwargs), + } + context_hash = _hash(json.dumps(context, default=str, sort_keys=True)) + return ( + f"capability_router:v1:{self.model_name}:{self._config_hash}:" + f"{_hash(str(caller))[:16]}:{_hash(str(session))[:16]}:{context_hash}" + ) + + async def _classify( + self, + messages: Sequence[Mapping[str, object]], + request_kwargs: Mapping[str, object], + ) -> tuple[CapabilityClassifierVerdict, float | None]: + classifier = self.config.classifier + classifier_messages: list[AllMessageValues] = [ + ChatCompletionSystemMessage(role="system", content=self._system_prompt), + ChatCompletionUserMessage(role="user", content=self._classifier_payload(messages, request_kwargs)), + ] + metadata = forwarded_internal_call_metadata(self._metadata(request_kwargs), AUTOROUTER_CLASSIFIER_CALL_ORIGIN) + response = await self.litellm_router_instance.acompletion( + model=classifier.model, + messages=classifier_messages, + response_format=self._response_format, + max_tokens=classifier.max_output_tokens, + timeout=classifier.timeout_ms / 1000, + metadata=metadata, + proxy_server_request=_ClassifierProxyRequest( + body=_ClassifierRequestBody( + model=classifier.model, + messages=classifier_messages, + response_format=self._response_format, + max_tokens=classifier.max_output_tokens, + ) + ), + ) + classifier_cost = _response_cost(response) + content = response.choices[0].message.content + if not isinstance(content, str) or not content.strip(): + raise CapabilityClassifierFailure("Capability classifier returned empty content", classifier_cost) + try: + return CapabilityClassifierVerdict.model_validate_json(_normalize_json(content)), classifier_cost + except ValidationError as exc: + raise CapabilityClassifierFailure("Capability classifier returned invalid JSON", classifier_cost) from exc + + def _estimated_usage( + self, + messages: Sequence[Mapping[str, object]], + request_kwargs: Mapping[str, object], + ) -> Usage | None: + import litellm + + try: + tools_value = request_kwargs.get("tools") + tools = _TOOLS_ADAPTER.validate_python(tools_value) if tools_value is not None else None + tool_choice_value = request_kwargs.get("tool_choice") + if tool_choice_value in ("none", "auto", "required", None): + tool_choice = tool_choice_value + else: + tool_choice = _NAMED_TOOL_CHOICE_ADAPTER.validate_python(tool_choice_value) + input_tokens = litellm.token_counter( + messages=list(messages), + tools=tools, + tool_choice=tool_choice, + use_default_image_token_count=True, + ) + except Exception as exc: # noqa: BLE001 - unpriceable requests use the explicit fallback + verbose_router_logger.warning("CapabilityRouter: token estimate failed (%s)", exc) + return None + + requested_limits = tuple( + value + for field in ("max_completion_tokens", "max_tokens", "max_output_tokens") + if isinstance((value := request_kwargs.get(field)), int) and not isinstance(value, bool) and value > 0 + ) + output_tokens = min((self.config.estimated_output_tokens, *requested_limits)) + return Usage( + prompt_tokens=input_tokens, + completion_tokens=output_tokens, + total_tokens=input_tokens + output_tokens, + ) + + async def _new_decision( + self, + messages: Sequence[Mapping[str, object]], + request_kwargs: Mapping[str, object], + ) -> tuple[CapabilityRoutingDecision, float | None]: + try: + verdict, classifier_cost = await self._classify(messages, request_kwargs) + usage = self._estimated_usage(messages, request_kwargs) + costs = { + candidate.model: ( + estimate_model_group_cost(self.litellm_router_instance, candidate.model, usage) + if usage is not None + else None + ) + for candidate in self.config.candidates + } + return select_capability_model(self.config, verdict, costs), classifier_cost + except CapabilityClassifierFailure as exc: + verbose_router_logger.warning("CapabilityRouter: classifier failed; using fallback") + return fallback_decision(self.config, "classifier_error"), exc.classifier_cost + except Exception as exc: # noqa: BLE001 - routing must fail safe + verbose_router_logger.warning("CapabilityRouter: selection failed (%s); using fallback", exc) + return fallback_decision(self.config, "classifier_error"), None + + def _cached_decision(self, value: object) -> CapabilityRoutingDecision | None: + try: + decision = CapabilityRoutingDecision.model_validate(value) + except ValidationError: + return None + configured = {candidate.model for candidate in self.config.candidates} + return decision if decision.selected_model in configured else None + + async def _decision( + self, + cache_key: str, + messages: Sequence[Mapping[str, object]], + request_kwargs: Mapping[str, object], + ) -> _DecisionOutcome: + cached = self._cached_decision(await self.litellm_router_instance.cache.async_get_cache(key=cache_key)) + if cached is not None: + return _DecisionOutcome(cached, None, True) + + lock = self._classification_locks.setdefault(cache_key, asyncio.Lock()) + async with lock: + cached = self._cached_decision(await self.litellm_router_instance.cache.async_get_cache(key=cache_key)) + if cached is not None: + return _DecisionOutcome(cached, None, True) + decision, classifier_cost = await self._new_decision(messages, request_kwargs) + await self.litellm_router_instance.cache.async_set_cache( + key=cache_key, + value=decision.model_dump(mode="json"), + ttl=self.config.cache_ttl_seconds, + ) + return _DecisionOutcome(decision, classifier_cost, False) + + def _routing_record(self, outcome: _DecisionOutcome) -> StandardLoggingRoutingDecision: + decision = outcome.decision + record = StandardLoggingRoutingDecision( + router_model_name=self.model_name, + router_type="capability", + routed_model=decision.selected_model, + cause=( + "capability_cache" + if outcome.cached + else "capability_classifier" + if decision.reason == "cheapest_qualified" + else "capability_fallback" + ), + classifier_model=self.config.classifier.model, + probability_threshold=self.config.probability_threshold, + candidate_probabilities={candidate.model: candidate.p_solve for candidate in decision.candidates}, + candidate_costs={ + candidate.model: candidate.estimated_cost + for candidate in decision.candidates + if candidate.estimated_cost is not None + }, + qualified_models=[candidate.model for candidate in decision.candidates if candidate.qualified], + fallback_reason=(decision.reason if decision.reason != "cheapest_qualified" else None), + cached=outcome.cached, + ) + if outcome.classifier_cost is not None: + record["classifier_cost"] = outcome.classifier_cost + return record + + async def async_pre_routing_hook( + self, + model: str, + request_kwargs: dict, + messages: list[dict[str, Any]] | None = None, + input: str | list | None = None, + specific_deployment: bool | None = False, + ) -> PreRoutingHookResponse: + from litellm.types.router import PreRoutingHookResponse + + resolved_messages = self._resolve_messages(messages, request_kwargs) + outcome = await self._decision( + self._cache_key(resolved_messages, request_kwargs), + resolved_messages, + request_kwargs, + ) + return PreRoutingHookResponse( + model=outcome.decision.selected_model, + messages=messages, + routing_decision=self._routing_record(outcome), + ) diff --git a/litellm/router_strategy/capability_router/config.py b/litellm/router_strategy/capability_router/config.py new file mode 100644 index 00000000000..ab4a553d260 --- /dev/null +++ b/litellm/router_strategy/capability_router/config.py @@ -0,0 +1,137 @@ +"""Configuration and classifier output for capability routing.""" + +import math +from typing import Literal + +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from typing_extensions import Self + + +def _nonblank(value: str, field_name: str) -> str: + normalized = value.strip() + if not normalized: + raise ValueError(f"{field_name} must not be blank") + return normalized + + +def _provider_model(value: str, field_name: str) -> str: + normalized = _nonblank(value, field_name) + if normalized.startswith("auto_router/"): + raise ValueError(f"{field_name} must name a provider model group") + return normalized + + +class CapabilityRouterCandidate(BaseModel): + """A model group and the operator's description of when it succeeds.""" + + model: str + description: str + + model_config = ConfigDict(extra="forbid") + + @field_validator("model") + @classmethod + def validate_model(cls, value: str) -> str: + return _provider_model(value, "candidate model") + + @field_validator("description") + @classmethod + def validate_description(cls, value: str) -> str: + return _nonblank(value, "candidate description") + + +class CapabilityClassifierConfig(BaseModel): + """The model used to estimate each candidate's probability of success.""" + + model: str + timeout_ms: int = Field(default=3000, ge=1) + max_output_tokens: int = Field(default=1024, ge=1) + + model_config = ConfigDict(extra="forbid") + + @field_validator("model") + @classmethod + def validate_model(cls, value: str) -> str: + return _provider_model(value, "classifier model") + + +class CapabilityRouterConfig(BaseModel): + """Operator configuration for cheapest-qualified model selection.""" + + candidates: tuple[CapabilityRouterCandidate, ...] = Field(min_length=2) + classifier: CapabilityClassifierConfig + probability_threshold: float = Field(default=0.7, ge=0.0, le=1.0) + fallback_model: str + estimated_output_tokens: int = Field(default=1000, ge=1) + cache_ttl_seconds: int = Field(default=3600, ge=1) + + model_config = ConfigDict(extra="forbid") + + @field_validator("fallback_model") + @classmethod + def validate_fallback_model(cls, value: str) -> str: + return _provider_model(value, "fallback_model") + + @field_validator("probability_threshold", mode="before") + @classmethod + def validate_finite_threshold(cls, value: object) -> object: + if isinstance(value, (int, float)) and not math.isfinite(value): + raise ValueError("probability_threshold must be finite") + return value + + @model_validator(mode="after") + def validate_candidates(self) -> Self: + models = tuple(candidate.model for candidate in self.candidates) + if len(models) != len(set(models)): + raise ValueError("candidate model names must be unique") + if self.fallback_model not in models: + raise ValueError("fallback_model must be one of the candidate models") + return self + + +class CapabilityCandidateScore(BaseModel): + """One classifier estimate for the current task.""" + + model: str + p_solve: float = Field(ge=0.0, le=1.0) + reason: str + + model_config = ConfigDict(extra="forbid") + + @field_validator("model", "reason") + @classmethod + def validate_nonblank(cls, value: str, info) -> str: + return _nonblank(value, info.field_name) + + @field_validator("p_solve", mode="before") + @classmethod + def validate_probability(cls, value: object) -> object: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError("p_solve must be a JSON number") + if isinstance(value, float) and not math.isfinite(value): + raise ValueError("p_solve must be finite") + return value + + +class CapabilityClassifierVerdict(BaseModel): + """Strict structured output returned by the classifier.""" + + candidates: tuple[CapabilityCandidateScore, ...] = Field(min_length=2) + + model_config = ConfigDict(extra="forbid") + + @model_validator(mode="after") + def validate_unique_models(self) -> Self: + models = tuple(candidate.model for candidate in self.candidates) + if len(models) != len(set(models)): + raise ValueError("classifier candidate model names must be unique") + return self + + +CapabilitySelectionReason = Literal[ + "cheapest_qualified", + "no_qualified_candidate", + "missing_candidate_price", + "classifier_error", + "invalid_classifier_verdict", +] diff --git a/litellm/router_strategy/capability_router/policy.py b/litellm/router_strategy/capability_router/policy.py new file mode 100644 index 00000000000..d76a831e9d3 --- /dev/null +++ b/litellm/router_strategy/capability_router/policy.py @@ -0,0 +1,80 @@ +"""Deterministic policy for selecting the cheapest qualified model.""" + +import math +from collections.abc import Mapping + +from pydantic import BaseModel, ConfigDict + +from .config import CapabilityClassifierVerdict, CapabilityRouterConfig, CapabilitySelectionReason + + +class CapabilityCandidateAssessment(BaseModel): + model: str + p_solve: float + reason: str + estimated_cost: float | None + qualified: bool + + model_config = ConfigDict(extra="forbid") + + +class CapabilityRoutingDecision(BaseModel): + selected_model: str + reason: CapabilitySelectionReason + candidates: tuple[CapabilityCandidateAssessment, ...] = () + + model_config = ConfigDict(extra="forbid") + + +def fallback_decision( + config: CapabilityRouterConfig, + reason: CapabilitySelectionReason, + candidates: tuple[CapabilityCandidateAssessment, ...] = (), +) -> CapabilityRoutingDecision: + return CapabilityRoutingDecision( + selected_model=config.fallback_model, + reason=reason, + candidates=candidates, + ) + + +def select_capability_model( + config: CapabilityRouterConfig, + verdict: CapabilityClassifierVerdict, + estimated_costs: Mapping[str, float | None], +) -> CapabilityRoutingDecision: + """Choose the cheapest candidate at or above the configured probability.""" + configured_models = tuple(candidate.model for candidate in config.candidates) + scores = {candidate.model: candidate for candidate in verdict.candidates} + if set(scores) != set(configured_models): + return fallback_decision(config, "invalid_classifier_verdict") + + assessments = tuple( + CapabilityCandidateAssessment( + model=model, + p_solve=scores[model].p_solve, + reason=scores[model].reason, + estimated_cost=estimated_costs.get(model), + qualified=scores[model].p_solve >= config.probability_threshold, + ) + for model in configured_models + ) + qualified = tuple(candidate for candidate in assessments if candidate.qualified) + if not qualified: + return fallback_decision(config, "no_qualified_candidate", assessments) + if any(candidate.estimated_cost is None for candidate in qualified): + return fallback_decision(config, "missing_candidate_price", assessments) + + order = {model: index for index, model in enumerate(configured_models)} + selected = min( + qualified, + key=lambda candidate: ( + candidate.estimated_cost if candidate.estimated_cost is not None else math.inf, + order[candidate.model], + ), + ) + return CapabilityRoutingDecision( + selected_model=selected.model, + reason="cheapest_qualified", + candidates=assessments, + ) diff --git a/litellm/router_strategy/capability_router/pricing.py b/litellm/router_strategy/capability_router/pricing.py new file mode 100644 index 00000000000..3048b437ea7 --- /dev/null +++ b/litellm/router_strategy/capability_router/pricing.py @@ -0,0 +1,79 @@ +"""Request-cost estimates used only by capability selection.""" + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from litellm._logging import verbose_router_logger +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.router_strategy.savings_baseline import canonical_model +from litellm.types.utils import Usage + +if TYPE_CHECKING: + from litellm.router import Router + + +@dataclass(frozen=True) +class _PricedModel: + model: str + deployment_id: str | None = None + + +def _deployment_model(deployment: Mapping[str, object]) -> _PricedModel | None: + params = deployment.get("litellm_params") + if not isinstance(params, Mapping): + return None + info_value = deployment.get("model_info") + info = info_value if isinstance(info_value, Mapping) else {} + model = info.get("base_model") or params.get("base_model") or params.get("model") + provider_value = params.get("custom_llm_provider") + provider = provider_value if isinstance(provider_value, str) else None + qualified = canonical_model(model, provider) if isinstance(model, str) else None + if qualified is None: + return None + deployment_id = info.get("id") + return _PricedModel(qualified, str(deployment_id) if deployment_id else None) + + +def _models_served_by(router: "Router", model_group: str) -> tuple[_PricedModel, ...]: + deployments = tuple(router.get_model_list(model_name=model_group) or ()) + if not deployments: + direct = canonical_model(model_group) + return (_PricedModel(direct),) if direct is not None else () + candidates = tuple( + candidate + for deployment in deployments + if (candidate := _deployment_model(deployment)) is not None + ) + return candidates if len(candidates) == len(deployments) else () + + +def _request_cost(router: "Router", candidate: _PricedModel, usage: Usage) -> float | None: + provider, _, model_name = candidate.model.partition("/") + try: + model_info = router.get_deployment_model_info(candidate.deployment_id or "", candidate.model) + if model_info is None: + return None + prompt_cost, completion_cost = generic_cost_per_token( + model=model_name or candidate.model, + usage=usage, + custom_llm_provider=provider, + model_info=model_info, + ) + except Exception as exc: # noqa: BLE001 - an unpriceable candidate uses the configured fallback + verbose_router_logger.debug("CapabilityRouter: no pricing for %s (%s)", candidate.model, exc) + return None + cost = prompt_cost + completion_cost + return cost if math.isfinite(cost) and cost >= 0 else None + + +def estimate_model_group_cost(router: "Router", model_group: str, usage: Usage) -> float | None: + """Return a conservative cost estimate for every deployment behind a group.""" + candidates = _models_served_by(router, model_group) + if not candidates: + return None + costs = tuple(_request_cost(router, candidate, usage) for candidate in candidates) + if any(cost is None for cost in costs): + return None + return max(cost for cost in costs if cost is not None) diff --git a/litellm/router_strategy/capability_router/prompts.py b/litellm/router_strategy/capability_router/prompts.py new file mode 100644 index 00000000000..b65080ff783 --- /dev/null +++ b/litellm/router_strategy/capability_router/prompts.py @@ -0,0 +1,45 @@ +"""Classifier prompt and response schema for capability routing.""" + +from typing import Any + +from .config import CapabilityRouterConfig + + +def build_classifier_prompt(config: CapabilityRouterConfig) -> str: + candidates = "\n".join( + f"- {candidate.model}: {candidate.description}" for candidate in config.candidates + ) + return f"""You route tasks between language models. + +For each candidate below, estimate the probability that it can complete the user's task correctly and completely. Base the estimate only on the task and the operator-provided description. Do not consider price; price is handled separately. + +Candidates: +{candidates} + +Return one score for every candidate using the exact model names. p_solve must be a number from 0 to 1. reason must briefly identify the capability that most affects the estimate.""" + + +def build_classifier_response_schema(config: CapabilityRouterConfig) -> dict[str, Any]: + model_names = [candidate.model for candidate in config.candidates] + return { + "type": "object", + "properties": { + "candidates": { + "type": "array", + "minItems": len(model_names), + "maxItems": len(model_names), + "items": { + "type": "object", + "properties": { + "model": {"type": "string", "enum": model_names}, + "p_solve": {"type": "number", "minimum": 0, "maximum": 1}, + "reason": {"type": "string", "minLength": 1}, + }, + "required": ["model", "p_solve", "reason"], + "additionalProperties": False, + }, + } + }, + "required": ["candidates"], + "additionalProperties": False, + } diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index a8aa543d735..5dac78d7bc0 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -22,9 +22,9 @@ from litellm.router_strategy.complexity_router.config import ( AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/" -StrategyRouterKind = Literal["semantic", "complexity", "adaptive", "quality"] +StrategyRouterKind = Literal["semantic", "complexity", "capability", "adaptive", "quality"] -StrategyRouterDependencyRole: TypeAlias = Literal["tier", "default", "classifier", "embedding"] +StrategyRouterDependencyRole: TypeAlias = Literal["tier", "candidate", "default", "classifier", "embedding"] @dataclass(frozen=True, slots=True) @@ -44,6 +44,7 @@ STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset( "auto_router_max_input_chars", "complexity_router_config", "complexity_router_default_model", + "capability_router_config", "adaptive_router_config", "quality_router_config", "quality_router_default_model", @@ -57,6 +58,7 @@ _REQUIRED_FIELD_GROUPS: Final[Mapping[StrategyRouterKind, tuple[tuple[str, ...], ("auto_router_embedding_model",), ), "complexity": (("complexity_router_config", "complexity_router_default_model"),), + "capability": (("capability_router_config",),), "adaptive": (("adaptive_router_config",),), "quality": (("quality_router_config", "quality_router_default_model"),), } @@ -74,6 +76,8 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None: remainder: Final = model[len(AUTO_ROUTER_MODEL_PREFIX) :] if remainder.startswith("complexity_router"): return "complexity" + if remainder.startswith("capability_router"): + return "capability" if remainder.startswith("adaptive_router"): return "adaptive" if remainder.startswith("quality_router"): @@ -143,6 +147,26 @@ def strategy_router_dependencies( ) ) ) + if kind == "capability": + capability = _mapping(litellm_params.get("capability_router_config")) + classifier = _mapping(capability.get("classifier")) + candidates = capability.get("candidates") + candidate_dependencies = ( + tuple( + dependency + for candidate in candidates + for dependency in _named(_mapping(candidate).get("model"), "candidate") + ) + if isinstance(candidates, Sequence) and not isinstance(candidates, (str, bytes)) + else () + ) + return tuple( + dict.fromkeys( + candidate_dependencies + + _named(capability.get("fallback_model"), "default") + + _named(classifier.get("model"), "classifier") + ) + ) complexity: Final = _mapping(litellm_params.get("complexity_router_config")) classifier: Final = _mapping(complexity.get("classifier_llm_config")) return tuple( @@ -191,6 +215,23 @@ def validate_complexity_router_config_write(complexity_router_config: Mapping[st return None +def validate_capability_router_config_write(capability_router_config: Mapping[str, object] | None) -> str | None: + """Reject a capability config the router itself would reject.""" + from pydantic import ValidationError + + from litellm.router_strategy.capability_router.config import CapabilityRouterConfig + + if capability_router_config is None: + return None + try: + _ = CapabilityRouterConfig.model_validate(capability_router_config) + except ValidationError as exc: + first = exc.errors()[0] + location = ".".join(str(part) for part in first.get("loc", ())) or "capability_router_config" + return f"capability_router_config is invalid at {location}: {first.get('msg', 'invalid value')}" + return None + + _COMPLEXITY_ROUTER_FIELDS: Final[frozenset[str]] = frozenset( field for group in _REQUIRED_FIELD_GROUPS["complexity"] for field in group ) diff --git a/litellm/types/router.py b/litellm/types/router.py index e0957383aac..61af5deb88b 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -348,6 +348,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): complexity_router_config: dict | None = None complexity_router_default_model: str | None = None + # capability-router params + capability_router_config: dict | None = None + # adaptive-router params adaptive_router_default_model: str | None = None adaptive_router_config: dict | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5783a39b30c..037282c5820 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2888,6 +2888,9 @@ RoutingDecisionCause = Literal[ "keyword", "quality_tier", "bandit", + "capability_classifier", + "capability_cache", + "capability_fallback", ] @@ -2910,7 +2913,7 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): """Per-request provenance for a pre-routing strategy (auto-router) decision.""" router_model_name: str - router_type: Literal["complexity", "adaptive", "quality"] + router_type: Literal["complexity", "adaptive", "quality", "capability"] routed_model: str cause: RoutingDecisionCause tier: str @@ -2931,6 +2934,12 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): savings_baseline_model: str savings_baseline_deployment_id: str tier_litellm_params: Mapping[str, object] # writable-ok: Pydantic warns on ReadOnly TypedDict fields + probability_threshold: float + candidate_probabilities: Mapping[str, float] + candidate_costs: Mapping[str, float] + qualified_models: Sequence[str] + fallback_reason: str | None + cached: bool # Fields whose values quote the caller's prompt. Dropped when an operator turns message @@ -2959,6 +2968,12 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "savings_baseline_model", "savings_baseline_deployment_id", "tier_litellm_params", + "probability_threshold", + "candidate_probabilities", + "candidate_costs", + "qualified_models", + "fallback_reason", + "cached", } ) @@ -3705,6 +3720,7 @@ all_litellm_params = ( "auto_router_max_input_chars", "complexity_router_config", "complexity_router_default_model", + "capability_router_config", "adaptive_router_config", "adaptive_router_default_model", "quality_router_config", diff --git a/tests/test_litellm/router_strategy/capability_router/test_capability_router.py b/tests/test_litellm/router_strategy/capability_router/test_capability_router.py new file mode 100644 index 00000000000..fd3fee10c5e --- /dev/null +++ b/tests/test_litellm/router_strategy/capability_router/test_capability_router.py @@ -0,0 +1,140 @@ +from unittest.mock import AsyncMock + +import pytest +from pydantic import ValidationError + +from litellm.router import Router +from litellm.router_strategy.capability_router.capability_router import CapabilityRouter +from litellm.router_strategy.capability_router.config import ( + CapabilityClassifierVerdict, + CapabilityRouterConfig, +) +from litellm.router_strategy.capability_router.policy import select_capability_model + + +def config() -> dict: + return { + "candidates": [ + {"model": "small", "description": "Reliable for short extraction tasks"}, + {"model": "frontier", "description": "Reliable for ambiguous multi-step tasks"}, + ], + "classifier": {"model": "classifier"}, + "probability_threshold": 0.7, + "fallback_model": "frontier", + "cache_ttl_seconds": 60, + } + + +def test_config_requires_unique_candidates_and_candidate_fallback() -> None: + duplicate = config() + duplicate["candidates"] = [ + {"model": "small", "description": "one"}, + {"model": "small", "description": "two"}, + ] + with pytest.raises(ValidationError, match="unique"): + CapabilityRouterConfig.model_validate(duplicate) + + missing_fallback = config() + missing_fallback["fallback_model"] = "other" + with pytest.raises(ValidationError, match="one of the candidate"): + CapabilityRouterConfig.model_validate(missing_fallback) + + +def test_policy_selects_cheapest_model_above_global_threshold() -> None: + parsed = CapabilityRouterConfig.model_validate(config()) + verdict = CapabilityClassifierVerdict.model_validate( + { + "candidates": [ + {"model": "small", "p_solve": 0.78, "reason": "clear bounded task"}, + {"model": "frontier", "p_solve": 0.95, "reason": "more capable"}, + ] + } + ) + + decision = select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05}) + + assert decision.selected_model == "small" + assert decision.reason == "cheapest_qualified" + assert [candidate.qualified for candidate in decision.candidates] == [True, True] + + +def test_policy_falls_back_if_no_model_qualifies_or_price_is_unknown() -> None: + parsed = CapabilityRouterConfig.model_validate(config()) + verdict = CapabilityClassifierVerdict.model_validate( + { + "candidates": [ + {"model": "small", "p_solve": 0.4, "reason": "too hard"}, + {"model": "frontier", "p_solve": 0.6, "reason": "uncertain"}, + ] + } + ) + assert select_capability_model(parsed, verdict, {"small": 0.01, "frontier": 0.05}).reason == ( + "no_qualified_candidate" + ) + + qualified = verdict.model_copy( + update={ + "candidates": tuple( + candidate.model_copy(update={"p_solve": 0.9}) for candidate in verdict.candidates + ) + } + ) + assert select_capability_model(parsed, qualified, {"small": None, "frontier": 0.05}).reason == ( + "missing_candidate_price" + ) + + +def test_router_registers_capability_strategy() -> None: + router = Router( + model_list=[ + {"model_name": "small", "litellm_params": {"model": "openai/test-small"}}, + {"model_name": "frontier", "litellm_params": {"model": "openai/test-frontier"}}, + {"model_name": "classifier", "litellm_params": {"model": "openai/test-classifier"}}, + { + "model_name": "cost-router", + "litellm_params": { + "model": "auto_router/capability_router", + "capability_router_config": config(), + }, + }, + ] + ) + + assert len(router.capability_routers["cost-router"]) == 1 + + +@pytest.mark.asyncio +async def test_same_user_turn_reuses_cached_decision() -> None: + router = Router(model_list=[]) + strategy = CapabilityRouter("cost-router", router, config()) + strategy._new_decision = AsyncMock( # type: ignore[method-assign] + return_value=( + select_capability_model( + strategy.config, + CapabilityClassifierVerdict.model_validate( + { + "candidates": [ + {"model": "small", "p_solve": 0.9, "reason": "fits"}, + {"model": "frontier", "p_solve": 0.95, "reason": "fits"}, + ] + } + ), + {"small": 0.01, "frontier": 0.05}, + ), + 0.001, + ) + ) + messages = [{"role": "user", "content": "Extract the invoice number"}] + kwargs = {"messages": messages, "metadata": {"user_api_key_hash": "key", "session_id": "session"}} + + first = await strategy.async_pre_routing_hook("cost-router", kwargs, messages) + second = await strategy.async_pre_routing_hook( + "cost-router", + {**kwargs, "messages": [*messages, {"role": "assistant", "content": "calling tool"}]}, + [*messages, {"role": "assistant", "content": "calling tool"}], + ) + + assert first.model == second.model == "small" + assert first.routing_decision is not None and first.routing_decision["cached"] is False + assert second.routing_decision is not None and second.routing_decision["cached"] is True + strategy._new_decision.assert_awaited_once()