mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(router): add cached capability selection
This commit is contained in:
moe-berri
Tin Chi Lo
parent
92d453373a
commit
655838b1da
11 changed files with 978 additions and 4 deletions
|
|
@ -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, []),
|
||||
|
|
|
|||
3
litellm/router_strategy/capability_router/__init__.py
Normal file
3
litellm/router_strategy/capability_router/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .capability_router import CapabilityRouter
|
||||
|
||||
__all__ = ["CapabilityRouter"]
|
||||
391
litellm/router_strategy/capability_router/capability_router.py
Normal file
391
litellm/router_strategy/capability_router/capability_router.py
Normal file
|
|
@ -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),
|
||||
)
|
||||
137
litellm/router_strategy/capability_router/config.py
Normal file
137
litellm/router_strategy/capability_router/config.py
Normal file
|
|
@ -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",
|
||||
]
|
||||
80
litellm/router_strategy/capability_router/policy.py
Normal file
80
litellm/router_strategy/capability_router/policy.py
Normal file
|
|
@ -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,
|
||||
)
|
||||
79
litellm/router_strategy/capability_router/pricing.py
Normal file
79
litellm/router_strategy/capability_router/pricing.py
Normal file
|
|
@ -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)
|
||||
45
litellm/router_strategy/capability_router/prompts.py
Normal file
45
litellm/router_strategy/capability_router/prompts.py
Normal file
|
|
@ -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,
|
||||
}
|
||||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
Loading…
Add table
Reference in a new issue