scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

chore(router): satisfy internal lint gates for capability router

Browse source
This commit is contained in:
Tin Chi Lo 2026-09-01 20:10:36 -07:00
parent d994c090c7
commit 41653a273a
14 changed files with 279 additions and 189 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/litellm_core_utils/token_counter.py
View file

@ -361,7 +361,7 @@ def token_counter(
messages: Sequence[AllMessageValues | Message] | None = None,
count_response_tokens: bool | None = False,
tools: list[ChatCompletionToolParam] | None = None,
tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
tool_choice: ChatCompletionNamedToolChoiceParam | Literal["none", "auto", "required"] | None = None,
use_default_image_token_count: bool | None = False,
default_token_count: int | None = None,
) -> int:
@ -507,7 +507,7 @@ def _count_messages(
def _count_extra(
count_function: TokenCounterFunction,
tools: list[ChatCompletionToolParam] | None,
tool_choice: ChatCompletionNamedToolChoiceParam | None,
tool_choice: ChatCompletionNamedToolChoiceParam | Literal["none", "auto", "required"] | None,
includes_system_message: bool,
) -> int:
"""Count extra tokens for function definitions and tool choices.

76
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -351,8 +351,8 @@ async def validate_complexity_router_config(
@router.post(
"/auto_router/validate_capability_router_config",
tags=["model management"],
dependencies=[Depends(user_api_key_auth)],
tags=["model management"], # mutable-ok: fastapi's decorator signature types tags as a list
dependencies=[Depends(user_api_key_auth)], # mutable-ok: fastapi's decorator signature types dependencies as a list
response_model=CapabilityRouterConfigValidationResponse,
status_code=status.HTTP_200_OK,
)
@ -370,6 +370,38 @@ async def validate_capability_router_config(
return CapabilityRouterConfigValidationResponse(valid=error is None, error=error)
async def _authorized_routing_test_strategy(
data: AutoRouterRoutingTestRequest,
user_api_key_dict: UserAPIKeyAuth,
llm_router: "Router",
) -> CapabilityRouter | ComplexityRouter:
"""Authorize the config's internal dry-run calls, then build the strategy under test."""
if data.capability_router_config is not None:
await _authorize_model_names_this_test_can_call(
models=(data.capability_router_config.classifier.model,),
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
return CapabilityRouter(
model_name=data.router_name,
litellm_router_instance=llm_router,
capability_router_config=data.capability_router_config.model_dump(exclude_none=True),
)
assert data.complexity_router_config is not None
await _authorize_models_this_test_can_call(
config=data.complexity_router_config,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
return ComplexityRouter(
model_name=data.router_name,
litellm_router_instance=llm_router,
complexity_router_config=data.complexity_router_config.model_dump(exclude_none=True),
default_model=data.default_model,
derive_savings_baseline=False,
)
@router.post(
"/auto_router/test_routing",
tags=["model management"], # mutable-ok: fastapi's decorator signature types tags as a list
@ -433,31 +465,11 @@ async def preview_auto_router_routing(
},
)
if data.capability_router_config is not None:
await _authorize_model_names_this_test_can_call(
models=(data.capability_router_config.classifier.model,),
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
strategy = CapabilityRouter(
model_name=data.router_name,
litellm_router_instance=llm_router,
capability_router_config=data.capability_router_config.model_dump(exclude_none=True),
)
else:
assert data.complexity_router_config is not None
await _authorize_models_this_test_can_call(
config=data.complexity_router_config,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
strategy = ComplexityRouter(
model_name=data.router_name,
litellm_router_instance=llm_router,
complexity_router_config=data.complexity_router_config.model_dump(exclude_none=True),
default_model=data.default_model,
derive_savings_baseline=False,
)
strategy: Final = await _authorized_routing_test_strategy(
data=data,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
request_kwargs: Final = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
data={ # mutable-ok: the request-metadata helper takes and returns request kwargs as a dict
@ -683,7 +695,7 @@ def _idle_router_groups(
@router.get(
"/auto_router/benchmarks",
tags=("auto router",),
tags=["auto router"], # mutable-ok: fastapi's decorator signature types tags as a list
dependencies=(Depends(user_api_key_auth),),
response_model=AutoRouterBenchmarksResponse,
)
@ -1353,7 +1365,7 @@ async def _shadow_eval_results(
@router.post(
"/auto_router/shadow_eval/start",
tags=("auto router",),
tags=["auto router"], # mutable-ok: fastapi's decorator signature types tags as a list
dependencies=(Depends(user_api_key_auth),),
response_model=ShadowEvalJobResponse,
status_code=status.HTTP_201_CREATED,
@ -1576,7 +1588,7 @@ async def start_shadow_eval(
@router.get(
"/auto_router/shadow_eval",
tags=("auto router",),
tags=["auto router"], # mutable-ok: fastapi's decorator signature types tags as a list
dependencies=(Depends(user_api_key_auth),),
response_model=list[ShadowEvalJobResponse],
)
@ -1626,7 +1638,7 @@ async def list_shadow_eval_jobs(
@router.get(
"/auto_router/shadow_eval/{job_id}",
tags=("auto router",),
tags=["auto router"], # mutable-ok: fastapi's decorator signature types tags as a list
dependencies=(Depends(user_api_key_auth),),
response_model=ShadowEvalJobResponse,
)
@ -1681,7 +1693,7 @@ async def get_shadow_eval_job(
@router.post(
"/auto_router/shadow_eval/{job_id}/stop",
tags=("auto router",),
tags=["auto router"], # mutable-ok: fastapi's decorator signature types tags as a list
dependencies=(Depends(user_api_key_auth),),
response_model=ShadowEvalJobResponse,
)

12
litellm/router.py
View file

@ -795,7 +795,9 @@ 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.capability_routers: dict[ # mutable-ok: registry mutated by the shared pre-routing helpers
str, list[TaggedPreRoutingStrategy[CapabilityRouter]]
] = {} # mutable-ok: registry mutated by the shared pre-routing helpers
self.complexity_routers: dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] = {}
self.adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] = {}
self.quality_routers: dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] = {}
@ -8575,9 +8577,7 @@ class Router:
config: Final = deployment.litellm_params.capability_router_config
if config is None:
raise ValueError(
"capability_router_config is required for capability-router deployments"
)
raise ValueError("capability_router_config is required for capability-router deployments")
capability_router: Final = CapabilityRouter(
model_name=deployment.model_name,
litellm_router_instance=self,
@ -8947,7 +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.capability_routers = {} # mutable-ok: registry mutated by the shared pre-routing helpers
self.complexity_routers = {}
self.auto_routers = {}
self._provider_unresolved_deployments = ()
@ -12577,7 +12577,7 @@ class Router:
"""
candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [
*self.auto_routers.get(model, []),
*self.capability_routers.get(model, []),
*self.capability_routers.get(model, ()),
*self.complexity_routers.get(model, []),
*self.adaptive_routers.get(model, []),
*self.quality_routers.get(model, []),

2
litellm/router_strategy/capability_router/__init__.py
View file

@ -1,3 +1,3 @@
from .capability_router import CapabilityRouter
__all__ = ["CapabilityRouter"]
__all__ = ("CapabilityRouter",)

192
litellm/router_strategy/capability_router/capability_router.py
View file

@ -8,7 +8,8 @@ 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 types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypedDict
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly
@ -33,7 +34,7 @@ from litellm.types.utils import (
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
from .prompts import ClassifierResponseSchema, build_classifier_prompt, build_classifier_response_schema
if TYPE_CHECKING:
from litellm.router import Router
@ -43,7 +44,7 @@ if TYPE_CHECKING:
class _JsonSchemaSpec(TypedDict):
name: ReadOnly[str]
strict: ReadOnly[bool]
schema: ReadOnly[dict[str, Any]]
schema: ReadOnly[ClassifierResponseSchema]
class _JsonSchemaResponseFormat(TypedDict):
@ -62,6 +63,16 @@ class _ClassifierProxyRequest(TypedDict):
body: ReadOnly[_ClassifierRequestBody]
class _ClassifierPayload(TypedDict):
conversation: ReadOnly[tuple[object, ...]]
available_tools: ReadOnly[tuple[str, ...]]
class _CacheKeyContext(TypedDict):
messages: ReadOnly[tuple[Mapping[str, object], ...]]
tools: ReadOnly[tuple[str, ...]]
@dataclass(frozen=True)
class _DecisionOutcome:
decision: CapabilityRoutingDecision
@ -84,15 +95,17 @@ def _hash(value: str) -> str:
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
hidden_params: Final = getattr(response, "_hidden_params", None)
if not isinstance(hidden_params, Mapping):
return None
value: Final = hidden_params.get("response_cost")
if isinstance(value, bool) or not isinstance(value, (int, float)):
return None
return float(value)
def _normalize_json(content: str) -> str:
normalized = content.strip()
normalized: Final = content.strip()
for prefix in ("```json", "```"):
if normalized.startswith(prefix) and normalized.endswith("```"):
return normalized[len(prefix) : -3].strip()
@ -100,50 +113,45 @@ def _normalize_json(content: str) -> str:
def _message_role(message: Mapping[str, object]) -> str:
role = message.get("role")
role: Final = 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(
last_user_index: Final = 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)
through_user: Final = messages[: last_user_index + 1]
recent: Final = through_user[-8:]
kept: Final = tuple(message for message in through_user if _message_role(message) == "system" or message in recent)
return tuple(message for index, message in enumerate(kept) if message not in kept[:index])
def _capped(value: object, cap: int) -> object:
if isinstance(value, str):
return value if len(value) <= cap else f"{value[:cap]}...[truncated {len(value) - cap} chars]"
if isinstance(value, Mapping):
return {key: _capped(item, cap) for key, item in value.items()}
return {key: _capped(item, cap) for key, item in value.items()} # mutable-ok: json.dumps needs a plain dict
if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
return tuple(_capped(item, cap) for item in value)
return value
def _tool_names(request_kwargs: Mapping[str, object]) -> tuple[str, ...]:
tools = request_kwargs.get("tools")
tools: Final = 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)
return tuple(
name
for tool in tools
if isinstance(tool, Mapping)
and isinstance((function := tool.get("function")), Mapping)
and isinstance((name := function.get("name")), str)
)
class CapabilityRouter(CustomLogger):
@ -172,33 +180,36 @@ class CapabilityRouter(CustomLogger):
@staticmethod
def _resolve_messages(
messages: list[dict[str, Any]] | None,
request_kwargs: dict[str, Any],
) -> list[dict[str, Any]]:
messages: list[dict[str, object]] | None, # mutable-ok: same shape the base hook receives
request_kwargs: dict[str, object], # mutable-ok: handed to resolve_structured_messages as-is
) -> tuple[Mapping[str, object], ...]:
from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
return resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or []
resolved: Final = resolve_structured_messages(messages=messages, request_kwargs=request_kwargs)
return tuple(resolved) if resolved else ()
@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
return MappingProxyType(
{
key: value
for carrier in ("metadata", "litellm_metadata")
if isinstance((entry := request_kwargs.get(carrier)), Mapping)
for key, value in entry.items()
}
)
def _classifier_payload(
self,
messages: Sequence[Mapping[str, object]],
request_kwargs: Mapping[str, object],
) -> str:
context = _classification_context(messages)
context: Final = _classification_context(messages)
if not context:
raise CapabilityClassifierFailure("No user task was available for capability classification")
cap: Final = self.config.classifier.max_message_chars
payload: Final = {
payload: Final[_ClassifierPayload] = {
"conversation": tuple(_capped(message, cap) for message in context),
"available_tools": _tool_names(request_kwargs),
}
@ -212,14 +223,14 @@ class CapabilityRouter(CustomLogger):
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 = {
metadata: Final = self._metadata(request_kwargs)
caller: Final = metadata.get("user_api_key_hash") or metadata.get("team_id") or "unscoped"
session: Final = metadata.get("session_id") or request_kwargs.get("litellm_session_id") or "no-session"
context: Final[_CacheKeyContext] = {
"messages": _classification_context(messages),
"tools": _tool_names(request_kwargs),
}
context_hash = _hash(json.dumps(context, default=str, sort_keys=True))
context_hash: Final = _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}"
@ -230,13 +241,15 @@ class CapabilityRouter(CustomLogger):
messages: Sequence[Mapping[str, object]],
request_kwargs: Mapping[str, object],
) -> tuple[CapabilityClassifierVerdict, float | None]:
classifier = self.config.classifier
classifier_messages: list[AllMessageValues] = [
classifier: Final = self.config.classifier
classifier_messages: Final[list[AllMessageValues]] = [ # mutable-ok: router.acompletion requires a list
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(
metadata: Final = forwarded_internal_call_metadata(
self._metadata(request_kwargs), AUTOROUTER_CLASSIFIER_CALL_ORIGIN
)
response: Final = await self.litellm_router_instance.acompletion(
model=classifier.model,
messages=classifier_messages,
response_format=self._response_format,
@ -252,8 +265,8 @@ class CapabilityRouter(CustomLogger):
)
),
)
classifier_cost = _response_cost(response)
content = response.choices[0].message.content
classifier_cost: Final = _response_cost(response)
content: Final = response.choices[0].message.content
if not isinstance(content, str) or not content.strip():
raise CapabilityClassifierFailure("Capability classifier returned empty content", classifier_cost)
try:
@ -269,15 +282,16 @@ class CapabilityRouter(CustomLogger):
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_value: Final = request_kwargs.get("tools")
tools: Final = _TOOLS_ADAPTER.validate_python(tools_value) if tools_value is not None else None
tool_choice_value: Final = request_kwargs.get("tool_choice")
tool_choice: Final = (
tool_choice_value
if tool_choice_value in ("none", "auto", "required", None)
else _NAMED_TOOL_CHOICE_ADAPTER.validate_python(tool_choice_value)
)
input_tokens: Final = litellm.token_counter(
messages=messages,
tools=tools,
tool_choice=tool_choice,
use_default_image_token_count=True,
@ -286,12 +300,12 @@ class CapabilityRouter(CustomLogger):
verbose_router_logger.warning("CapabilityRouter: token estimate failed (%s)", exc)
return None
requested_limits = tuple(
requested_limits: Final = 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))
output_tokens: Final = min((self.config.estimated_output_tokens, *requested_limits))
return Usage(
prompt_tokens=input_tokens,
completion_tokens=output_tokens,
@ -305,15 +319,17 @@ class CapabilityRouter(CustomLogger):
) -> 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
}
usage: Final = self._estimated_usage(messages, request_kwargs)
costs: Final = MappingProxyType(
{
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")
@ -324,10 +340,10 @@ class CapabilityRouter(CustomLogger):
def _cached_decision(self, value: object) -> CapabilityRoutingDecision | None:
try:
decision = CapabilityRoutingDecision.model_validate(value)
decision: Final = CapabilityRoutingDecision.model_validate(value)
except ValidationError:
return None
configured = {candidate.model for candidate in self.config.candidates}
configured: Final = frozenset(candidate.model for candidate in self.config.candidates)
return decision if decision.selected_model in configured else None
async def _decision(
@ -336,15 +352,17 @@ class CapabilityRouter(CustomLogger):
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))
cached: Final = 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())
lock: Final = 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)
rechecked: Final = self._cached_decision(
await self.litellm_router_instance.cache.async_get_cache(key=cache_key)
)
if rechecked is not None:
return _DecisionOutcome(rechecked, None, True)
decision, classifier_cost = await self._new_decision(messages, request_kwargs)
await self.litellm_router_instance.cache.async_set_cache(
key=cache_key,
@ -354,8 +372,8 @@ class CapabilityRouter(CustomLogger):
return _DecisionOutcome(decision, classifier_cost, False)
def _routing_record(self, outcome: _DecisionOutcome) -> StandardLoggingRoutingDecision:
decision = outcome.decision
record = StandardLoggingRoutingDecision(
decision: Final = outcome.decision
record: Final = StandardLoggingRoutingDecision(
router_model_name=self.model_name,
router_type="capability",
routed_model=decision.selected_model,
@ -368,13 +386,15 @@ class CapabilityRouter(CustomLogger):
),
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_probabilities={ # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
candidate.model: candidate.p_solve for candidate in decision.candidates
},
candidate_costs={ # mutable-ok: safe_dumps stringifies non-dict mappings in the spend log
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],
qualified_models=tuple(candidate.model for candidate in decision.candidates if candidate.qualified),
fallback_reason=(decision.reason if decision.reason != "cheapest_qualified" else None),
cached=outcome.cached,
)
@ -385,15 +405,15 @@ class CapabilityRouter(CustomLogger):
async def async_pre_routing_hook(
self,
model: str,
request_kwargs: dict,
messages: list[dict[str, Any]] | None = None,
input: str | list | None = None,
request_kwargs: dict, # mutable-ok: same shape the base hook receives
messages: list[dict[str, object]] | None = None, # mutable-ok: same shape the base hook receives
input: str | list | None = None, # mutable-ok: same shape the base hook receives
specific_deployment: bool | None = False,
) -> PreRoutingHookResponse:
from litellm.types.router import PreRoutingHookResponse
resolved_messages = self._resolve_messages(messages, request_kwargs)
outcome = await self._decision(
resolved_messages: Final = self._resolve_messages(messages, request_kwargs)
outcome: Final = await self._decision(
self._cache_key(resolved_messages, request_kwargs),
resolved_messages,
request_kwargs,

20
litellm/router_strategy/capability_router/config.py
View file

@ -1,21 +1,21 @@
"""Configuration and classifier output for capability routing."""
import math
from typing import Literal
from typing import Final, Literal, TypeAlias
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()
normalized: Final = 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)
normalized: Final = _nonblank(value, field_name)
if normalized.startswith("auto_router/"):
raise ValueError(f"{field_name} must name a provider model group")
return normalized
@ -83,15 +83,15 @@ class CapabilityRouterConfig(BaseModel):
@model_validator(mode="after")
def validate_candidates(self) -> Self:
models = tuple(candidate.model for candidate in self.candidates)
if len(models) != len(set(models)):
models: Final = tuple(candidate.model for candidate in self.candidates)
if len(models) != len(frozenset(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
CapabilityBoundary = Literal["supported", "uncertain", "unsupported", "unmatched"]
CapabilityBoundary: TypeAlias = Literal["supported", "uncertain", "unsupported", "unmatched"]
class CapabilityCandidateScore(BaseModel):
@ -113,7 +113,7 @@ class CapabilityCandidateScore(BaseModel):
@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")
raise ValueError("p_solve must be a JSON number") # noqa: TRY004 # pydantic needs ValueError
if isinstance(value, float) and not math.isfinite(value):
raise ValueError("p_solve must be finite")
return value
@ -128,13 +128,13 @@ class CapabilityClassifierVerdict(BaseModel):
@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)):
models: Final = tuple(candidate.model for candidate in self.candidates)
if len(models) != len(frozenset(models)):
raise ValueError("classifier candidate model names must be unique")
return self
CapabilitySelectionReason = Literal[
CapabilitySelectionReason: TypeAlias = Literal[
"cheapest_qualified",
"no_qualified_candidate",
"missing_candidate_price",

14
litellm/router_strategy/capability_router/policy.py
View file

@ -56,12 +56,12 @@ def select_capability_model(
estimated_costs: Mapping[str, float | None],
) -> CapabilityRoutingDecision:
"""Choose the cheapest candidate whose p_solve clears its boundary-stepped threshold."""
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):
configured_models: Final = tuple(candidate.model for candidate in config.candidates)
scores: Final = MappingProxyType({candidate.model: candidate for candidate in verdict.candidates})
if frozenset(scores) != frozenset(configured_models):
return fallback_decision(config, "invalid_classifier_verdict")
assessments = tuple(
assessments: Final = tuple(
CapabilityCandidateAssessment(
model=model,
p_solve=scores[model].p_solve,
@ -77,14 +77,14 @@ def select_capability_model(
)
for model in configured_models
)
qualified = tuple(candidate for candidate in assessments if candidate.qualified)
qualified: Final = 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(
order: Final = MappingProxyType({model: index for index, model in enumerate(configured_models)})
selected: Final = min(
qualified,
key=lambda candidate: (
candidate.estimated_cost if candidate.estimated_cost is not None else math.inf,

37
litellm/router_strategy/capability_router/pricing.py
View file

@ -3,7 +3,8 @@
import math
from collections.abc import Mapping
from dataclasses import dataclass
from typing import TYPE_CHECKING
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_router_logger
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
@ -21,30 +22,28 @@ class _PricedModel:
def _deployment_model(deployment: Mapping[str, object]) -> _PricedModel | None:
params = deployment.get("litellm_params")
params: Final = 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
info_value: Final = deployment.get("model_info")
info: Final = info_value if isinstance(info_value, Mapping) else MappingProxyType({})
model: Final = info.get("base_model") or params.get("base_model") or params.get("model")
provider_value: Final = params.get("custom_llm_provider")
provider: Final = provider_value if isinstance(provider_value, str) else None
qualified: Final = canonical_model(model, provider) if isinstance(model, str) else None
if qualified is None:
return None
deployment_id = info.get("id")
deployment_id: Final = 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 ())
deployments: Final = tuple(router.get_model_list(model_name=model_group) or ())
if not deployments:
direct = canonical_model(model_group)
direct: Final = 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
candidates: Final = tuple(
candidate for deployment in deployments if (candidate := _deployment_model(deployment)) is not None
)
return candidates if len(candidates) == len(deployments) else ()
@ -52,7 +51,7 @@ def _models_served_by(router: "Router", model_group: str) -> tuple[_PricedModel,
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)
model_info: Final = 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(
@ -64,16 +63,16 @@ def _request_cost(router: "Router", candidate: _PricedModel, usage: Usage) -> fl
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
cost: Final = 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)
candidates: Final = _models_served_by(router, model_group)
if not candidates:
return None
costs = tuple(_request_cost(router, candidate, usage) for candidate in candidates)
costs: Final = 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)

77
litellm/router_strategy/capability_router/prompts.py
View file

@ -1,14 +1,64 @@
"""Classifier prompt and response schema for capability routing."""
from typing import Any
from typing import Final, Literal, TypedDict
from typing_extensions import ReadOnly
from .config import CapabilityRouterConfig
CAPABILITY_BOUNDARIES = ("supported", "uncertain", "unsupported", "unmatched")
CAPABILITY_BOUNDARIES: Final = ("supported", "uncertain", "unsupported", "unmatched")
class _StringEnumSchema(TypedDict):
type: ReadOnly[Literal["string"]]
enum: ReadOnly[list[str]] # mutable-ok: provider json_schema transforms only rewrite list-typed arrays
class _NonBlankStringSchema(TypedDict):
type: ReadOnly[Literal["string"]]
minLength: ReadOnly[int]
class _UnitIntervalSchema(TypedDict):
type: ReadOnly[Literal["number"]]
minimum: ReadOnly[int]
maximum: ReadOnly[int]
class _CandidateProperties(TypedDict):
model: ReadOnly[_StringEnumSchema]
reason: ReadOnly[_NonBlankStringSchema]
capability_boundary: ReadOnly[_StringEnumSchema]
p_solve: ReadOnly[_UnitIntervalSchema]
class _CandidateItemSchema(TypedDict):
type: ReadOnly[Literal["object"]]
properties: ReadOnly[_CandidateProperties]
required: ReadOnly[list[str]] # mutable-ok: provider json_schema transforms only rewrite list-typed arrays
additionalProperties: ReadOnly[bool]
class _CandidateArraySchema(TypedDict):
type: ReadOnly[Literal["array"]]
minItems: ReadOnly[int]
maxItems: ReadOnly[int]
items: ReadOnly[_CandidateItemSchema]
class _VerdictProperties(TypedDict):
candidates: ReadOnly[_CandidateArraySchema]
class ClassifierResponseSchema(TypedDict):
type: ReadOnly[Literal["object"]]
properties: ReadOnly[_VerdictProperties]
required: ReadOnly[list[str]] # mutable-ok: provider json_schema transforms only rewrite list-typed arrays
additionalProperties: ReadOnly[bool]
def build_classifier_prompt(config: CapabilityRouterConfig) -> str:
candidates = "\n".join(f"- {candidate.model}: {candidate.description}" for candidate in config.candidates)
candidates: Final = "\n".join(f"- {candidate.model}: {candidate.description}" for candidate in config.candidates)
return f"""You forecast task outcomes for a model router.
For each candidate model below, forecast one binary event. SUCCESS means the candidate completes the newest user task correctly and completely in one fresh attempt, using only the tools available in the request. FAILURE is any other outcome. The two outcomes are exhaustive.
@ -28,9 +78,9 @@ Candidates:
Return one entry for every candidate using the exact model names."""
def build_classifier_response_schema(config: CapabilityRouterConfig) -> dict[str, Any]:
model_names = [candidate.model for candidate in config.candidates]
return {
def build_classifier_response_schema(config: CapabilityRouterConfig) -> ClassifierResponseSchema:
model_names: Final = [candidate.model for candidate in config.candidates] # mutable-ok: wire arrays are lists
schema: Final[ClassifierResponseSchema] = {
"type": "object",
"properties": {
"candidates": {
@ -42,14 +92,23 @@ def build_classifier_response_schema(config: CapabilityRouterConfig) -> dict[str
"properties": {
"model": {"type": "string", "enum": model_names},
"reason": {"type": "string", "minLength": 1},
"capability_boundary": {"type": "string", "enum": list(CAPABILITY_BOUNDARIES)},
"capability_boundary": {
"type": "string",
"enum": list(CAPABILITY_BOUNDARIES), # mutable-ok: wire arrays are lists
},
"p_solve": {"type": "number", "minimum": 0, "maximum": 1},
},
"required": ["model", "reason", "capability_boundary", "p_solve"],
"required": [ # mutable-ok: wire arrays are lists
"model",
"reason",
"capability_boundary",
"p_solve",
],
"additionalProperties": False,
},
}
},
"required": ["candidates"],
"required": ["candidates"], # mutable-ok: wire arrays are lists
"additionalProperties": False,
}
return schema

16
litellm/router_utils/auto_router_model_naming.py
View file

@ -22,7 +22,7 @@ from litellm.router_strategy.complexity_router.config import (
AUTO_ROUTER_MODEL_PREFIX: Final = "auto_router/"
StrategyRouterKind = Literal["semantic", "complexity", "capability", "adaptive", "quality"]
StrategyRouterKind: TypeAlias = Literal["semantic", "complexity", "capability", "adaptive", "quality"]
StrategyRouterDependencyRole: TypeAlias = Literal["tier", "candidate", "default", "classifier", "embedding"]
@ -148,10 +148,10 @@ 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 = (
capability: Final = _mapping(litellm_params.get("capability_router_config"))
capability_classifier: Final = _mapping(capability.get("classifier"))
candidates: Final = capability.get("candidates")
candidate_dependencies: Final = (
tuple(
dependency
for candidate in candidates
@ -164,7 +164,7 @@ def strategy_router_dependencies(
dict.fromkeys(
candidate_dependencies
+ _named(capability.get("fallback_model"), "default")
+ _named(classifier.get("model"), "classifier")
+ _named(capability_classifier.get("model"), "classifier")
)
)
complexity: Final = _mapping(litellm_params.get("complexity_router_config"))
@ -226,8 +226,8 @@ def validate_capability_router_config_write(capability_router_config: Mapping[st
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"
first: Final = exc.errors()[0]
location: Final = ".".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

2
litellm/types/router.py
View file

@ -349,7 +349,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
complexity_router_default_model: str | None = None
# capability-router params
capability_router_config: dict | None = None
capability_router_config: dict | None = None # mutable-ok: sibling router-config fields arrive via dict splats
# adaptive-router params
adaptive_router_default_model: str | None = None

12
litellm/types/utils.py
View file

@ -2934,12 +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
probability_threshold: float # writable-ok: Pydantic warns on ReadOnly TypedDict fields
candidate_probabilities: Mapping[str, float] # writable-ok: Pydantic warns on ReadOnly TypedDict fields
candidate_costs: Mapping[str, float] # writable-ok: Pydantic warns on ReadOnly TypedDict fields
qualified_models: Sequence[str] # writable-ok: Pydantic warns on ReadOnly TypedDict fields
fallback_reason: str | None # writable-ok: Pydantic warns on ReadOnly TypedDict fields
cached: bool # writable-ok: Pydantic warns on ReadOnly TypedDict fields
# Fields whose values quote the caller's prompt. Dropped when an operator turns message

2
litellm/utils.py
View file

@ -2307,7 +2307,7 @@ def token_counter(
messages: Sequence | None = None,
count_response_tokens: bool | None = False,
tools: list[ChatCompletionToolParam] | None = None,
tool_choice: ChatCompletionNamedToolChoiceParam | None = None,
tool_choice: ChatCompletionNamedToolChoiceParam | Literal["none", "auto", "required"] | None = None,
use_default_image_token_count: bool | None = False,
default_token_count: int | None = None,
) -> int:

2
tests/test_litellm/router_strategy/test_complexity_router.py
View file

@ -8395,7 +8395,7 @@ class TestSavingsBaselineOnDecision:
from litellm.proxy.management_endpoints import auto_router_endpoints
source = inspect.getsource(auto_router_endpoints.preview_auto_router_routing)
source = inspect.getsource(auto_router_endpoints._authorized_routing_test_strategy)
assert "derive_savings_baseline=False" in source