diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 3732ffd734c..d78acaf48b0 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -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. diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 5d0395ea7f1..ada44db1df1 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -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, ) diff --git a/litellm/router.py b/litellm/router.py index dd0155ddff9..dd4a4c22326 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, []), diff --git a/litellm/router_strategy/capability_router/__init__.py b/litellm/router_strategy/capability_router/__init__.py index a22b9a2268e..bee8eddd3da 100644 --- a/litellm/router_strategy/capability_router/__init__.py +++ b/litellm/router_strategy/capability_router/__init__.py @@ -1,3 +1,3 @@ from .capability_router import CapabilityRouter -__all__ = ["CapabilityRouter"] +__all__ = ("CapabilityRouter",) diff --git a/litellm/router_strategy/capability_router/capability_router.py b/litellm/router_strategy/capability_router/capability_router.py index bf60a29775f..8b8164daa10 100644 --- a/litellm/router_strategy/capability_router/capability_router.py +++ b/litellm/router_strategy/capability_router/capability_router.py @@ -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, diff --git a/litellm/router_strategy/capability_router/config.py b/litellm/router_strategy/capability_router/config.py index 1489d80a86d..1d1e6098507 100644 --- a/litellm/router_strategy/capability_router/config.py +++ b/litellm/router_strategy/capability_router/config.py @@ -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", diff --git a/litellm/router_strategy/capability_router/policy.py b/litellm/router_strategy/capability_router/policy.py index 8eafeb14bc4..81b5acc8f14 100644 --- a/litellm/router_strategy/capability_router/policy.py +++ b/litellm/router_strategy/capability_router/policy.py @@ -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, diff --git a/litellm/router_strategy/capability_router/pricing.py b/litellm/router_strategy/capability_router/pricing.py index 3048b437ea7..2146d6ef67c 100644 --- a/litellm/router_strategy/capability_router/pricing.py +++ b/litellm/router_strategy/capability_router/pricing.py @@ -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) diff --git a/litellm/router_strategy/capability_router/prompts.py b/litellm/router_strategy/capability_router/prompts.py index 464d53bafb6..d7e3b8f2372 100644 --- a/litellm/router_strategy/capability_router/prompts.py +++ b/litellm/router_strategy/capability_router/prompts.py @@ -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 diff --git a/litellm/router_utils/auto_router_model_naming.py b/litellm/router_utils/auto_router_model_naming.py index 5dac78d7bc0..57d7832bd8b 100644 --- a/litellm/router_utils/auto_router_model_naming.py +++ b/litellm/router_utils/auto_router_model_naming.py @@ -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 diff --git a/litellm/types/router.py b/litellm/types/router.py index 61af5deb88b..ffa166618a3 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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 diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 037282c5820..f148d4eb4e2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 diff --git a/litellm/utils.py b/litellm/utils.py index 252b6756937..2aec3e44ec4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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: diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 1ec8be88c9b..178d98b4290 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -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