mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
feat(router): add Python Switchyard autorouting strategy
This commit is contained in:
Joshua Valluru
parent
c43a1addf2
commit
10f3dd75e5
34 changed files with 2775 additions and 13 deletions
|
|
@ -0,0 +1,27 @@
|
|||
# ADR 002: Virtual routing calls settle usage independently
|
||||
|
||||
Status: Proposed
|
||||
|
||||
## Context
|
||||
|
||||
Switchyard can call an efficient model, judge its response, and either return that response or call a capable model. The existing pre-routing hook only selects a deployment. Repeating an accepted response through another completion would generate a second answer and charge twice
|
||||
|
||||
## Decision
|
||||
|
||||
Dispatch Switchyard chat requests through an execution strategy before ordinary deployment selection. The strategy calls configured models through the owning Python router, with separate call IDs and inherited caller attribution. Each provider call records its own usage and cost. The virtual request does not record an additional provider charge
|
||||
|
||||
Child calls do not own the virtual request's admission reservation. The router host releases that reservation after execution finishes, fails, or is cancelled. When the final provider call streams, release waits for stream completion or closure
|
||||
|
||||
Return an accepted provider response unchanged, including usage extras described in ADR 001. Judged responses are buffered. A caller requesting streaming receives the accepted response through the existing response-to-stream conversion, without an additional billing wrapper. Ordinary final calls can stream directly
|
||||
|
||||
## Alternatives Considered
|
||||
|
||||
Returning only a model name through the pre-routing hook cannot reuse an answer already generated during escalation
|
||||
|
||||
Combining every internal call into a fabricated parent usage object would lose model-specific pricing and make provider usage fields ambiguous
|
||||
|
||||
## Consequences
|
||||
|
||||
Provider execution and billing remain in LiteLLM, while Switchyard owns routing decisions. Logs contain separate provider calls linked to the caller's session
|
||||
|
||||
Judged streaming has additional time to first token. Cancellation, reservation cleanup, and the sum of child charges require end-to-end verification. The initial execution strategy supports async Chat Completions only and rejects other endpoint paths until their adapters are implemented
|
||||
|
|
@ -7,7 +7,7 @@ target up at call time, so patches of the proxy functions apply.
|
|||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.router_host import CompressionPolicy, RouterHost
|
||||
|
||||
|
|
@ -45,6 +45,15 @@ class ProxyRouterHost(RouterHost):
|
|||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
async def finalize_strategy_request(self, request_kwargs: Mapping[str, object]) -> None:
|
||||
from litellm.litellm_core_utils.core_helpers import budget_reservation_from_metadata
|
||||
from litellm.proxy.spend_tracking.budget_reservation import release_or_invalidate_budget_reservation
|
||||
|
||||
for key in ("metadata", "litellm_metadata"):
|
||||
metadata: Final = request_kwargs.get(key)
|
||||
if isinstance(metadata, Mapping):
|
||||
await release_or_invalidate_budget_reservation(budget_reservation_from_metadata(metadata))
|
||||
|
||||
def compression_policy(
|
||||
self,
|
||||
llm_router: "Router | None",
|
||||
|
|
|
|||
|
|
@ -145,6 +145,7 @@ from litellm.router_strategy.auto_router.auto_router_model_naming import (
|
|||
validate_complexity_router_config_placement,
|
||||
validate_complexity_router_config_write,
|
||||
validate_strategy_router_model_write,
|
||||
validate_switchyard_router_config_write,
|
||||
)
|
||||
from litellm.router_strategy.auto_router.auto_router_tuning_baseline import (
|
||||
is_mutable_tuned_candidate,
|
||||
|
|
@ -374,6 +375,9 @@ def _strategy_router_write_violation(
|
|||
"""
|
||||
if incoming_params is None:
|
||||
return None
|
||||
switchyard_violation: Final = validate_switchyard_router_config_write(incoming_params.switchyard_router_config)
|
||||
if switchyard_violation is not None:
|
||||
return switchyard_violation
|
||||
config_violation: Final = validate_complexity_router_config_write(
|
||||
complexity_router_config=(
|
||||
_effective_complexity_router_config(incoming_params, existing_params)
|
||||
|
|
|
|||
|
|
@ -63,6 +63,9 @@ class RouterHost:
|
|||
response=httpx.Response(status_code=403, request=httpx.Request("POST", "https://litellm.invalid")),
|
||||
)
|
||||
|
||||
async def finalize_strategy_request(self, request_kwargs: Mapping[str, object]) -> None:
|
||||
"""Release admission held by a virtual request whose children settle their own usage."""
|
||||
|
||||
def compression_policy(
|
||||
self,
|
||||
llm_router: "Router | None",
|
||||
|
|
|
|||
|
|
@ -536,6 +536,7 @@ if TYPE_CHECKING:
|
|||
from litellm.router_strategy.quality_router.quality_router import (
|
||||
QualityRouter,
|
||||
)
|
||||
from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardAutoRouter
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import (
|
||||
AnthropicMessagesResponse,
|
||||
)
|
||||
|
|
@ -1769,6 +1770,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self._core_provider_default_deployment_ids: list[str] = []
|
||||
self._core_pattern_router = PatternMatchRouter()
|
||||
self._core_team_pattern_routers: dict[str, PatternMatchRouter] = {} # {"TEAM_ID": PatternMatchRouter}
|
||||
self._core_switchyard_routers: dict[ # mutable-ok: registry lifecycle mutates model entries
|
||||
str, list[TaggedPreRoutingStrategy[SwitchyardAutoRouter]]
|
||||
] = {}
|
||||
self._core_auto_routers: dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] = {}
|
||||
self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] = {}
|
||||
self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] = {}
|
||||
|
|
@ -3726,6 +3730,13 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
kwargs["stream"] = stream
|
||||
kwargs["original_function"] = self._acompletion
|
||||
|
||||
if self._core_switchyard_routers:
|
||||
from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardAutoRouter
|
||||
|
||||
selected: Final = self._select_pre_routing_strategy(model, kwargs)
|
||||
if selected is not None and isinstance(selected.strategy, SwitchyardAutoRouter):
|
||||
return await self._execute_switchyard(selected, model, kwargs)
|
||||
|
||||
self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs)
|
||||
start_time: Final = time.time()
|
||||
is_prompt_management_model: Final = self._is_prompt_management_model(model)
|
||||
|
|
@ -3773,6 +3784,100 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
)
|
||||
raise e
|
||||
|
||||
async def _execute_switchyard(
|
||||
self,
|
||||
selected: TaggedPreRoutingStrategy[PreRoutingStrategy],
|
||||
model: str,
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> ModelResponse | CustomStreamWrapper | AsyncIterable[ModelResponseStream]:
|
||||
from litellm.router_strategy.switchyard_router.execution import (
|
||||
ACTIVE_SWITCHYARD,
|
||||
FinalizingStream,
|
||||
SwitchyardReplay,
|
||||
child_call_kwargs,
|
||||
scoped_switchyard_headers,
|
||||
switchyard_execution,
|
||||
)
|
||||
from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardAutoRouter
|
||||
|
||||
strategy: Final = selected.strategy
|
||||
if not isinstance(strategy, SwitchyardAutoRouter):
|
||||
raise TypeError("Expected a Switchyard execution strategy")
|
||||
if ACTIVE_SWITCHYARD.get():
|
||||
raise ValueError("Switchyard targets cannot be auto-router deployments")
|
||||
marker: Final = self._selected_strategy_marker_deployment(model, selected.tags, request_kwargs)
|
||||
await router_host().authorize_strategy_inference(
|
||||
deployment=marker, request_kwargs=request_kwargs, llm_router=self._core_public()
|
||||
)
|
||||
defaults: Final = dict(self._forwardable_alias_marker_params(model, selected.tags, request_kwargs))
|
||||
forwarded: Final = tuple(key for key in defaults if key not in request_kwargs)
|
||||
values: Final = {
|
||||
**defaults,
|
||||
**request_kwargs,
|
||||
**({_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: forwarded} if forwarded else {}),
|
||||
}
|
||||
config: Final = strategy.config
|
||||
affinity: Final = (
|
||||
config.custom
|
||||
if config.algorithm == "custom"
|
||||
else config.stage.classifier
|
||||
if config.algorithm == "stage"
|
||||
else config.capability
|
||||
)
|
||||
message_hash_fallback: Final = affinity.message_hash_fallback if affinity is not None else False
|
||||
headers: Final = scoped_switchyard_headers(values, message_hash_fallback, strategy.session_namespace)
|
||||
|
||||
async def execute_model(target: str, payload: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper:
|
||||
child: Final = dict(child_call_kwargs(values, payload))
|
||||
decision: Final[StandardLoggingRoutingDecision] = {
|
||||
"router_model_name": selected.deployment.model_name if selected.deployment else model,
|
||||
"router_type": "switchyard",
|
||||
"routed_model": target,
|
||||
"signals": (f"algorithm:{config.algorithm}",),
|
||||
**(
|
||||
{"router_config_fingerprint": selected.definition_fingerprint}
|
||||
if selected.definition_fingerprint is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
self._record_routing_decision(child, decision)
|
||||
request_tags: Final = get_routing_tags_from_request_kwargs(values)
|
||||
if (
|
||||
selected.tags
|
||||
and request_tags
|
||||
and is_valid_deployment_tag(selected.tags, request_tags, self._core_tag_filtering_match_any)
|
||||
):
|
||||
self._stamp_or_clear_metadata_key(
|
||||
child,
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY,
|
||||
ConsumedRequestTagsStamp(model_group=target, tags=selected.tags),
|
||||
)
|
||||
if self._select_pre_routing_strategy(target, child) is not None:
|
||||
raise ValueError("Switchyard targets cannot be auto-router deployments")
|
||||
child_messages: Final = TypeAdapter(list[AllMessageValues]).validate_python(child.get("messages"))
|
||||
child_stream: Final = child.get("stream") is True
|
||||
child_options: Final = {key: value for key, value in child.items() if key not in {"messages", "stream"}}
|
||||
return await self.acompletion(model=target, messages=child_messages, stream=child_stream, **child_options)
|
||||
|
||||
async def finalize() -> None:
|
||||
await router_host().finalize_strategy_request(request_kwargs=request_kwargs)
|
||||
|
||||
try:
|
||||
with switchyard_execution():
|
||||
response: Final = await strategy.execute(values, execute_model, headers=headers)
|
||||
except BaseException:
|
||||
await finalize()
|
||||
raise
|
||||
if isinstance(response, CustomStreamWrapper):
|
||||
return FinalizingStream(
|
||||
response.__aiter__(), finalize, response.aclose, get_hidden_params_dict(response), response.logging_obj
|
||||
)
|
||||
await finalize()
|
||||
if request_kwargs.get("stream") is True and isinstance(response, ModelResponse):
|
||||
options: Final = request_kwargs.get("stream_options")
|
||||
return SwitchyardReplay(response, isinstance(options, Mapping) and options.get("include_usage") is True)
|
||||
return response
|
||||
|
||||
@staticmethod
|
||||
def _combine_fallback_usage(
|
||||
fallback_item: ModelResponseStream,
|
||||
|
|
@ -11126,6 +11231,17 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
"""Deployment tags used to disambiguate strategy registries keyed by model_name."""
|
||||
return tuple(deployment.litellm_params.tags or ())
|
||||
|
||||
def init_switchyard_router_deployment(self, deployment: Deployment) -> None:
|
||||
from litellm.router_strategy.switchyard_router import SwitchyardAutoRouter, SwitchyardRouterConfig
|
||||
|
||||
config: Final = SwitchyardRouterConfig.model_validate(deployment.litellm_params.switchyard_router_config)
|
||||
self._register_pre_routing_strategy(
|
||||
registry=self._core_switchyard_routers,
|
||||
deployment=deployment,
|
||||
strategy=SwitchyardAutoRouter(config),
|
||||
strategy_label="Switchyard",
|
||||
)
|
||||
|
||||
def init_auto_router_deployment(self, deployment: Deployment):
|
||||
"""
|
||||
Initialize the auto-router deployment.
|
||||
|
|
@ -11365,7 +11481,12 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
return
|
||||
model_name: Final = deployment.model_name
|
||||
tags: Final = self._deployment_tags(deployment)
|
||||
for registry in (self._core_auto_routers, self._core_complexity_routers, self._core_quality_routers):
|
||||
for registry in (
|
||||
self._core_auto_routers,
|
||||
self._core_complexity_routers,
|
||||
self._core_quality_routers,
|
||||
self._core_switchyard_routers,
|
||||
):
|
||||
self._unregister_pre_routing_strategy(registry, model_name, tags)
|
||||
if self._unregister_pre_routing_strategy(self._core_adaptive_routers, model_name, tags):
|
||||
self._sync_adaptive_router_hooks()
|
||||
|
|
@ -11575,6 +11696,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self._core_team_public_model_names = frozenset()
|
||||
# Reset per-strategy router registries so hot-reload doesn't leave
|
||||
# stale routers pointing at the old model_list.
|
||||
self._core_switchyard_routers = {}
|
||||
self._core_quality_routers = {}
|
||||
self._core_complexity_routers = {}
|
||||
self._core_auto_routers = {}
|
||||
|
|
@ -11770,6 +11892,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
#########################################################
|
||||
# Check if this is an auto-router deployment
|
||||
#########################################################
|
||||
if classify_strategy_router_model(deployment.litellm_params.model) == "switchyard":
|
||||
self.init_switchyard_router_deployment(deployment)
|
||||
|
||||
if self._is_auto_router_deployment(litellm_params=deployment.litellm_params):
|
||||
self.init_auto_router_deployment(deployment=deployment)
|
||||
|
||||
|
|
@ -16417,6 +16542,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self._core_complexity_routers,
|
||||
self._core_adaptive_routers,
|
||||
self._core_quality_routers,
|
||||
self._core_switchyard_routers,
|
||||
)
|
||||
if not any(registries):
|
||||
return None
|
||||
|
|
@ -16510,7 +16636,9 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
registered_model_name: str,
|
||||
request_kwargs: Mapping[str, object],
|
||||
) -> str:
|
||||
if is_native_compaction_call():
|
||||
from litellm.router_strategy.switchyard_router.execution import ACTIVE_SWITCHYARD
|
||||
|
||||
if is_native_compaction_call() or ACTIVE_SWITCHYARD.get():
|
||||
return registered_model_name
|
||||
if not any(
|
||||
(
|
||||
|
|
@ -16518,6 +16646,7 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
self._core_complexity_routers,
|
||||
self._core_adaptive_routers,
|
||||
self._core_quality_routers,
|
||||
self._core_switchyard_routers,
|
||||
)
|
||||
):
|
||||
return registered_model_name
|
||||
|
|
@ -16573,8 +16702,11 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
spend metadata is stamped before routing and the response carries the tier group the
|
||||
strategy picked.
|
||||
"""
|
||||
from litellm.router_strategy.switchyard_router.execution import ACTIVE_SWITCHYARD
|
||||
|
||||
self._restore_client_output_ceiling(request_kwargs)
|
||||
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
|
||||
if not ACTIVE_SWITCHYARD.get():
|
||||
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
|
||||
requested_registered_model_name: Final = self._core_get_model_from_alias(model=model) or model
|
||||
registered_model_name: Final = await self._resolve_claude_code_session_router(
|
||||
model=model,
|
||||
|
|
@ -16597,6 +16729,8 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
model=registered_model_name, request_kwargs=request_kwargs
|
||||
)
|
||||
if selected_strategy is None:
|
||||
if ACTIVE_SWITCHYARD.get():
|
||||
return None
|
||||
await arm_compaction(request_kwargs, None)
|
||||
self._stamp_or_clear_metadata_key(
|
||||
request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ the router silently dropping the deployment at load time under
|
|||
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
|
|
@ -26,7 +27,7 @@ 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", "adaptive", "quality", "switchyard"]
|
||||
|
||||
StrategyRouterDependencyRole: TypeAlias = Literal[
|
||||
"tier", "default", "classifier", "embedding", "evaluation", "compactor"
|
||||
|
|
@ -52,6 +53,7 @@ STRATEGY_ROUTER_PARAM_FIELDS: Final[frozenset[str]] = frozenset(
|
|||
"complexity_router_default_model",
|
||||
"adaptive_router_config",
|
||||
"quality_router_config",
|
||||
"switchyard_router_config",
|
||||
"quality_router_default_model",
|
||||
}
|
||||
)
|
||||
|
|
@ -64,6 +66,7 @@ _REQUIRED_FIELD_GROUPS: Final[Mapping[StrategyRouterKind, tuple[tuple[str, ...],
|
|||
),
|
||||
"complexity": (("complexity_router_config", "complexity_router_default_model"),),
|
||||
"adaptive": (("adaptive_router_config",),),
|
||||
"switchyard": (("switchyard_router_config",),),
|
||||
"quality": (("quality_router_config", "quality_router_default_model"),),
|
||||
}
|
||||
|
||||
|
|
@ -82,6 +85,8 @@ def classify_strategy_router_model(model: str) -> StrategyRouterKind | None:
|
|||
return "complexity"
|
||||
if remainder.startswith("adaptive_router"):
|
||||
return "adaptive"
|
||||
if remainder.startswith("switchyard"):
|
||||
return "switchyard"
|
||||
if remainder.startswith("quality_router"):
|
||||
return "quality"
|
||||
return "semantic"
|
||||
|
|
@ -141,6 +146,12 @@ def strategy_router_dependencies(
|
|||
return _named(litellm_params.get("auto_router_default_model"), "default") + _named(
|
||||
litellm_params.get("auto_router_embedding_model"), "embedding"
|
||||
)
|
||||
if kind == "switchyard":
|
||||
models: Final = _mapping(_mapping(litellm_params.get("switchyard_router_config")).get("models"))
|
||||
dependencies: Final = tuple(
|
||||
_pool(names, "classifier" if group == "judge" else "tier") for group, names in models.items()
|
||||
)
|
||||
return tuple(dict.fromkeys(chain.from_iterable(dependencies)))
|
||||
if kind == "adaptive":
|
||||
return _pool(_mapping(litellm_params.get("adaptive_router_config")).get("available_models"), "tier")
|
||||
if kind == "quality":
|
||||
|
|
@ -363,6 +374,22 @@ def capability_limit_violation(*, capability: GatedAutoRouterCapability, held: i
|
|||
)
|
||||
|
||||
|
||||
def validate_switchyard_router_config_write(config: Mapping[str, object] | None) -> str | None:
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig
|
||||
|
||||
if config is None:
|
||||
return None
|
||||
try:
|
||||
_ = SwitchyardRouterConfig.model_validate(config)
|
||||
except ValidationError as exc:
|
||||
first: Final = exc.errors()[0]
|
||||
location: Final = ".".join(str(part) for part in first.get("loc", ())) or "switchyard_router_config"
|
||||
return f"switchyard_router_config is invalid at {location}: {first.get('msg', 'invalid value')}"
|
||||
return None
|
||||
|
||||
|
||||
def validate_complexity_router_config_write(complexity_router_config: Mapping[str, object] | None) -> str | None:
|
||||
"""Reject a complexity config the router would refuse to build a deployment from.
|
||||
|
||||
|
|
|
|||
4
litellm/router_strategy/switchyard_router/__init__.py
Normal file
4
litellm/router_strategy/switchyard_router/__init__.py
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
from .config import SwitchyardRouterConfig
|
||||
from .switchyard_router import SwitchyardAutoRouter
|
||||
|
||||
__all__ = ["SwitchyardAutoRouter", "SwitchyardRouterConfig"]
|
||||
141
litellm/router_strategy/switchyard_router/config.py
Normal file
141
litellm/router_strategy/switchyard_router/config.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
|
||||
|
||||
class CapabilityConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
base_threshold: float = Field(default=0.5, ge=0, le=1)
|
||||
threshold_step: float = Field(default=0, ge=0, le=1)
|
||||
session_affinity: bool = False
|
||||
message_hash_fallback: bool = False
|
||||
recent_turn_window: int | None = Field(default=None, gt=0)
|
||||
max_output_tokens: int = Field(default=4096, gt=0)
|
||||
prompt: str | None = None
|
||||
response_format_type: Literal["json_schema", "json_object"] = "json_schema"
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_affinity(self) -> CapabilityConfig:
|
||||
if self.message_hash_fallback and not self.session_affinity:
|
||||
raise ValueError("message_hash_fallback requires session_affinity")
|
||||
return self
|
||||
|
||||
|
||||
class DeescalationConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
strong_min_calls: int = Field(gt=0)
|
||||
confirmations: int = Field(gt=0)
|
||||
strong_max_calls: int | None = Field(default=None, gt=0)
|
||||
weak_cooldown_calls: int = Field(default=0, ge=0)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_call_bounds(self) -> DeescalationConfig:
|
||||
if self.strong_max_calls is not None and self.strong_max_calls < self.strong_min_calls:
|
||||
raise ValueError("strong_max_calls must be at least strong_min_calls")
|
||||
return self
|
||||
|
||||
|
||||
class EscalationConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
confirmations: int = Field(default=2, gt=0)
|
||||
recent_turn_window: int = Field(default=28, gt=0)
|
||||
window_message_chars: int = Field(default=500, ge=50)
|
||||
max_output_tokens: int = Field(default=4096, gt=0)
|
||||
prompt: str | None = None
|
||||
response_format_type: Literal["json_schema", "json_object"] = "json_schema"
|
||||
deescalation: DeescalationConfig | None = None
|
||||
|
||||
|
||||
class CustomConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
prompt: str = Field(min_length=1)
|
||||
response_schema: Mapping[str, object]
|
||||
selector: str = Field(min_length=1)
|
||||
default_target: str = Field(min_length=1)
|
||||
session_affinity: bool = False
|
||||
message_hash_fallback: bool = False
|
||||
recent_turn_window: int | None = Field(default=None, gt=0)
|
||||
max_output_tokens: int = Field(default=4096, gt=0)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_affinity(self) -> CustomConfig:
|
||||
if self.message_hash_fallback and not self.session_affinity:
|
||||
raise ValueError("message_hash_fallback requires session_affinity")
|
||||
return self
|
||||
|
||||
|
||||
class StageConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
picker: Literal["efficient_first", "capable_first"] = "efficient_first"
|
||||
confidence_threshold: float = Field(default=0.5, ge=0, le=1)
|
||||
recent_window: int | None = Field(default=None, gt=0)
|
||||
escalation_note: str | None = None
|
||||
deescalation_note: str | None = None
|
||||
only_on_wrong_signal_escalation: bool = True
|
||||
capable_system_prompt: str | None = None
|
||||
efficient_system_prompt: str | None = None
|
||||
tool_semantics: Mapping[str, tuple[str, ...]] | None = None
|
||||
classifier: CapabilityConfig | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_notes_and_tools(self) -> StageConfig:
|
||||
if self.deescalation_note is not None and self.escalation_note is None:
|
||||
raise ValueError("deescalation_note requires escalation_note")
|
||||
if self.tool_semantics is not None and set(self.tool_semantics) - {"observe", "mutate", "plan", "new"}:
|
||||
raise ValueError("tool_semantics supports observe, mutate, plan, and new categories")
|
||||
return self
|
||||
|
||||
|
||||
class SwitchyardRouterConfig(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, extra="forbid")
|
||||
|
||||
algorithm: Literal["random", "capability", "custom", "escalation", "stage"] = "capability"
|
||||
models: Mapping[str, tuple[str, ...]]
|
||||
weights: tuple[float, ...] | None = None
|
||||
seed: int | None = Field(default=None, ge=0, le=18446744073709551615)
|
||||
fallback_enabled: bool = True
|
||||
capability: CapabilityConfig = Field(default_factory=CapabilityConfig)
|
||||
escalation: EscalationConfig = Field(default_factory=EscalationConfig)
|
||||
custom: CustomConfig | None = None
|
||||
stage: StageConfig = Field(default_factory=StageConfig)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_models(self) -> SwitchyardRouterConfig:
|
||||
if any(
|
||||
not group or not names or any(not name.strip() for name in names) for group, names in self.models.items()
|
||||
):
|
||||
raise ValueError("model groups and model names must be nonempty")
|
||||
required: Final = self.required_groups()
|
||||
if any(group not in self.models for group in required):
|
||||
raise ValueError(f"{self.algorithm} requires model groups {', '.join(required)}")
|
||||
if self.algorithm == "custom" and self.custom is None:
|
||||
raise ValueError("custom routing requires custom configuration")
|
||||
if self.custom is not None and self.custom.default_target not in self.models:
|
||||
raise ValueError("custom.default_target must name a configured model group")
|
||||
if self.weights is not None:
|
||||
import math
|
||||
|
||||
if self.algorithm != "random":
|
||||
raise ValueError("weights are only supported for random routing")
|
||||
if len(self.weights) != len(self.models.get("any", ())):
|
||||
raise ValueError("weights must match the any model group")
|
||||
if any(not math.isfinite(weight) or weight < 0 for weight in self.weights) or not any(self.weights):
|
||||
raise ValueError("weights must be finite, nonnegative, and have a positive total")
|
||||
return self
|
||||
|
||||
def required_groups(self) -> tuple[str, ...]:
|
||||
if self.algorithm == "random":
|
||||
return ("any",)
|
||||
if self.algorithm == "custom":
|
||||
return ("judge",)
|
||||
if self.algorithm == "stage" and self.stage.classifier is None:
|
||||
return ("efficient", "capable")
|
||||
return ("efficient", "capable", "judge")
|
||||
178
litellm/router_strategy/switchyard_router/execution.py
Normal file
178
litellm/router_strategy/switchyard_router/execution.py
Normal file
|
|
@ -0,0 +1,178 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
from litellm._uuid import uuid
|
||||
from litellm.litellm_core_utils.internal_call_metadata import forwarded_internal_call_metadata
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_streaming_response_async
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream
|
||||
|
||||
ACTIVE_SWITCHYARD: Final[ContextVar[bool]] = ContextVar("active_switchyard", default=False)
|
||||
_EXECUTION_CONTEXT_KEYS: Final = frozenset(
|
||||
{
|
||||
"litellm_trace_id",
|
||||
"litellm_session_id",
|
||||
"timeout",
|
||||
"request_timeout",
|
||||
"num_retries",
|
||||
"disable_fallbacks",
|
||||
"fallbacks",
|
||||
"context_window_fallbacks",
|
||||
"content_policy_fallbacks",
|
||||
"tags",
|
||||
"enable_tag_filtering",
|
||||
"proxy_server_request",
|
||||
"turn_off_message_logging",
|
||||
"user",
|
||||
"guardrails",
|
||||
}
|
||||
)
|
||||
_CALL_LOCAL_KEYS: Final = frozenset(
|
||||
{
|
||||
"model",
|
||||
"original_function",
|
||||
"litellm_call_id",
|
||||
"litellm_logging_obj",
|
||||
"metadata",
|
||||
"litellm_metadata",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def switchyard_execution() -> Generator[None]:
|
||||
token: Final = ACTIVE_SWITCHYARD.set(True)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
ACTIVE_SWITCHYARD.reset(token)
|
||||
|
||||
|
||||
def child_call_kwargs(parent: Mapping[str, object], payload: Mapping[str, object]) -> Mapping[str, object]:
|
||||
metadata: Final = parent.get("metadata")
|
||||
litellm_metadata: Final = parent.get("litellm_metadata")
|
||||
return {
|
||||
**{key: value for key, value in parent.items() if key in _EXECUTION_CONTEXT_KEYS},
|
||||
**{key: value for key, value in payload.items() if key not in _CALL_LOCAL_KEYS},
|
||||
"metadata": forwarded_internal_call_metadata(
|
||||
metadata if isinstance(metadata, Mapping) else None, "switchyard_router"
|
||||
),
|
||||
**(
|
||||
{"litellm_metadata": forwarded_internal_call_metadata(litellm_metadata, "switchyard_router")}
|
||||
if isinstance(litellm_metadata, Mapping)
|
||||
else {}
|
||||
),
|
||||
"litellm_call_id": str(uuid.uuid4()),
|
||||
}
|
||||
|
||||
|
||||
def scoped_switchyard_headers(
|
||||
request: Mapping[str, object], message_hash_fallback: bool = False, sdk_namespace: str | None = None
|
||||
) -> Mapping[str, str]:
|
||||
proxy_request: Final = request.get("proxy_server_request")
|
||||
raw_headers: Final = proxy_request.get("headers") if isinstance(proxy_request, Mapping) else None
|
||||
headers: Final = raw_headers if isinstance(raw_headers, Mapping) else {}
|
||||
metadata: Final = request.get("litellm_metadata") or request.get("metadata")
|
||||
caller: Final = metadata if isinstance(metadata, Mapping) else {}
|
||||
caller_key: Final = caller.get("user_api_key_hash")
|
||||
sdk_session: Final = request.get("litellm_session_id")
|
||||
sdk_context: Final = (
|
||||
sdk_namespace is not None
|
||||
and proxy_request is None
|
||||
and "user_api_key_auth" not in caller
|
||||
and isinstance(sdk_session, str)
|
||||
and bool(sdk_session)
|
||||
)
|
||||
caller_scope: Final = (
|
||||
(caller_key, caller.get("user_api_key_team_id"))
|
||||
if isinstance(caller_key, str) and caller_key
|
||||
else (sdk_namespace,)
|
||||
if sdk_context
|
||||
else (str(uuid.uuid4()),)
|
||||
)
|
||||
scope: Final = json.dumps(caller_scope, sort_keys=True)
|
||||
normalized: Final = {key.lower(): value for key, value in headers.items() if isinstance(key, str)}
|
||||
aliases: Final = (
|
||||
"x-switchyard-session-id",
|
||||
"x-claude-code-session-id",
|
||||
"x-nemo-relay-session-id",
|
||||
"x-session-id",
|
||||
"session-id",
|
||||
)
|
||||
session: Final = (
|
||||
sdk_session
|
||||
if sdk_context
|
||||
else next((normalized[key] for key in aliases if isinstance(normalized.get(key), str)), None)
|
||||
)
|
||||
fallback: Final = (
|
||||
json.dumps(request.get("messages"), sort_keys=True, default=str) if message_hash_fallback else str(uuid.uuid4())
|
||||
)
|
||||
return {"x-switchyard-session-id": sha256(f"{scope}:{session or fallback}".encode()).hexdigest()}
|
||||
|
||||
|
||||
class SwitchyardReplay:
|
||||
def __init__(self, response: ModelResponse, include_usage: bool) -> None:
|
||||
self._hidden_params = response._hidden_params # pyright: ignore[reportPrivateUsage] # proxy response metadata contract
|
||||
self.response: Final = response
|
||||
self.include_usage: Final = include_usage
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ModelResponseStream]:
|
||||
return self._frames()
|
||||
|
||||
async def _frames(self) -> AsyncIterator[ModelResponseStream]:
|
||||
payload: Final = self.response.model_dump(exclude={"usage"})
|
||||
async for frame in convert_to_streaming_response_async(payload):
|
||||
yield frame
|
||||
if self.include_usage:
|
||||
usage_frame: Final = ModelResponseStream(
|
||||
id=self.response.id,
|
||||
created=self.response.created,
|
||||
model=self.response.model,
|
||||
choices=[],
|
||||
usage=self.response.usage,
|
||||
)
|
||||
yield usage_frame
|
||||
|
||||
|
||||
class FinalizingStream:
|
||||
def __init__(
|
||||
self,
|
||||
stream: AsyncIterator[ModelResponseStream],
|
||||
finalize: Callable[[], Awaitable[None]],
|
||||
close: Callable[[], Awaitable[None]],
|
||||
hidden_params: object,
|
||||
logging_obj: object,
|
||||
) -> None:
|
||||
self.stream: Final = stream
|
||||
self.finalize: Final = finalize
|
||||
self.close: Final = close
|
||||
self._hidden_params = hidden_params
|
||||
self.logging_obj = logging_obj
|
||||
self.closed = False
|
||||
|
||||
def __aiter__(self) -> FinalizingStream:
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> ModelResponseStream:
|
||||
if self.closed:
|
||||
raise StopAsyncIteration
|
||||
try:
|
||||
with switchyard_execution():
|
||||
return await self.stream.__anext__()
|
||||
except BaseException:
|
||||
await self.aclose()
|
||||
raise
|
||||
|
||||
async def aclose(self) -> None:
|
||||
if self.closed:
|
||||
return
|
||||
self.closed = True
|
||||
try:
|
||||
await self.close()
|
||||
finally:
|
||||
await self.finalize()
|
||||
249
litellm/router_strategy/switchyard_router/native.py
Normal file
249
litellm/router_strategy/switchyard_router/native.py
Normal file
|
|
@ -0,0 +1,249 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
|
||||
from importlib import import_module
|
||||
from typing import Final, Protocol, runtime_checkable
|
||||
|
||||
from .config import CapabilityConfig, SwitchyardRouterConfig
|
||||
|
||||
|
||||
class Aggregate(Protocol):
|
||||
@property
|
||||
def response(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class ModelCall(Protocol):
|
||||
@property
|
||||
def models(self) -> Sequence[str]: ...
|
||||
|
||||
@property
|
||||
def request(self) -> Mapping[str, object]: ...
|
||||
|
||||
async def respond(self, work: Awaitable[Aggregate]) -> None: ...
|
||||
|
||||
|
||||
class ModelStep(Protocol):
|
||||
@property
|
||||
def call(self) -> ModelCall: ...
|
||||
|
||||
|
||||
class Outcome(Protocol):
|
||||
@property
|
||||
def selected_model_ids(self) -> Sequence[str]: ...
|
||||
|
||||
@property
|
||||
def request(self) -> Mapping[str, object]: ...
|
||||
|
||||
@property
|
||||
def fallback_request(self) -> Mapping[str, object] | None: ...
|
||||
|
||||
@property
|
||||
def response(self) -> Aggregate | None: ...
|
||||
|
||||
|
||||
class DoneStep(Protocol):
|
||||
@property
|
||||
def outcome(self) -> Outcome: ...
|
||||
|
||||
|
||||
class StepTypes(Protocol):
|
||||
CallModel: type[ModelStep]
|
||||
Done: type[DoneStep]
|
||||
|
||||
|
||||
class ResponseTypes(Protocol):
|
||||
Agg: Callable[[Mapping[str, object]], Aggregate]
|
||||
|
||||
|
||||
class Algorithm(Protocol):
|
||||
def run_stream(
|
||||
self,
|
||||
request: Mapping[str, object],
|
||||
models: Mapping[str, Sequence[str]],
|
||||
*,
|
||||
headers: Mapping[str, str] | None = None,
|
||||
) -> AsyncIterator[ModelStep | DoneStep]: ...
|
||||
|
||||
|
||||
class ClassifierFactory(Protocol):
|
||||
def capability(self, *, config: object) -> object: ...
|
||||
|
||||
def escalation(self, *, config: object) -> object: ...
|
||||
|
||||
def custom(self, *, default_target: str, config: object) -> object: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class NativeApi(Protocol):
|
||||
Step: StepTypes
|
||||
LlmResponse: ResponseTypes
|
||||
LlmClassifierConfig: ClassifierFactory
|
||||
ContextWindowExceededError: type[Exception]
|
||||
LibsyError: type[Exception]
|
||||
|
||||
def TaskClassifierConfig(
|
||||
self,
|
||||
base_threshold: float,
|
||||
*,
|
||||
threshold_step: float,
|
||||
session_affinity: bool,
|
||||
message_hash_fallback: bool,
|
||||
recent_turn_window: int | None,
|
||||
max_output_tokens: int,
|
||||
prompt: str | None,
|
||||
response_format_type: str,
|
||||
) -> object: ...
|
||||
|
||||
def CustomClassifierConfig(
|
||||
self,
|
||||
prompt: str,
|
||||
response_schema: Mapping[str, object],
|
||||
selector: str,
|
||||
*,
|
||||
session_affinity: bool,
|
||||
message_hash_fallback: bool,
|
||||
recent_turn_window: int | None,
|
||||
max_output_tokens: int,
|
||||
) -> object: ...
|
||||
|
||||
def EscalationClassifierConfig(
|
||||
self,
|
||||
*,
|
||||
confirmations: int,
|
||||
recent_turn_window: int,
|
||||
window_message_chars: int,
|
||||
deescalation: object | None,
|
||||
max_output_tokens: int,
|
||||
prompt: str | None,
|
||||
response_format_type: str,
|
||||
) -> object: ...
|
||||
|
||||
def DeescalationConfig(
|
||||
self,
|
||||
*,
|
||||
strong_min_calls: int,
|
||||
confirmations: int,
|
||||
strong_max_calls: int | None,
|
||||
weak_cooldown_calls: int,
|
||||
) -> object: ...
|
||||
|
||||
def LlmFallback(self, *, config: object) -> object: ...
|
||||
|
||||
def random(self, *, weights: Sequence[float] | None, seed: int | None) -> Algorithm: ...
|
||||
|
||||
def llm_classifier(self, config: object) -> Algorithm: ...
|
||||
|
||||
def stage_router(
|
||||
self,
|
||||
*,
|
||||
picker: str,
|
||||
confidence_threshold: float,
|
||||
recent_window: int | None,
|
||||
escalation_note: str | None,
|
||||
deescalation_note: str | None,
|
||||
only_on_wrong_signal_escalation: bool,
|
||||
capable_system_prompt: str | None,
|
||||
efficient_system_prompt: str | None,
|
||||
tool_semantics: Mapping[str, Sequence[str]] | None,
|
||||
classifier: object | None,
|
||||
) -> Algorithm: ...
|
||||
|
||||
|
||||
class NativeExtension(Protocol):
|
||||
@property
|
||||
def libsy(self) -> NativeApi: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class NativeLoader(Protocol):
|
||||
def load_native(self) -> NativeExtension: ...
|
||||
|
||||
|
||||
def load_native() -> NativeApi:
|
||||
try:
|
||||
module: Final = import_module("switchyard_rust._native")
|
||||
except ImportError as error:
|
||||
raise ImportError("Switchyard routing requires the optional dependency nemo-switchyard==0.4.0a1") from error
|
||||
if not isinstance(module, NativeLoader):
|
||||
raise TypeError("Installed Switchyard does not expose the required libsy API")
|
||||
return module.load_native().libsy
|
||||
|
||||
|
||||
def capability_config(api: NativeApi, config: CapabilityConfig) -> object:
|
||||
return api.TaskClassifierConfig(
|
||||
base_threshold=config.base_threshold,
|
||||
threshold_step=config.threshold_step,
|
||||
session_affinity=config.session_affinity,
|
||||
message_hash_fallback=config.message_hash_fallback,
|
||||
recent_turn_window=config.recent_turn_window,
|
||||
max_output_tokens=config.max_output_tokens,
|
||||
prompt=config.prompt,
|
||||
response_format_type=config.response_format_type,
|
||||
)
|
||||
|
||||
|
||||
def build_algorithm(api: NativeApi, config: SwitchyardRouterConfig) -> Algorithm:
|
||||
if config.algorithm == "random":
|
||||
return api.random(weights=config.weights, seed=config.seed)
|
||||
if config.algorithm == "stage":
|
||||
stage: Final = config.stage
|
||||
fallback: Final = (
|
||||
api.LlmFallback(config=capability_config(api, stage.classifier)) if stage.classifier is not None else None
|
||||
)
|
||||
return api.stage_router(
|
||||
picker=stage.picker,
|
||||
confidence_threshold=stage.confidence_threshold,
|
||||
recent_window=stage.recent_window,
|
||||
escalation_note=stage.escalation_note,
|
||||
deescalation_note=stage.deescalation_note,
|
||||
only_on_wrong_signal_escalation=stage.only_on_wrong_signal_escalation,
|
||||
capable_system_prompt=stage.capable_system_prompt,
|
||||
efficient_system_prompt=stage.efficient_system_prompt,
|
||||
tool_semantics=stage.tool_semantics,
|
||||
classifier=fallback,
|
||||
)
|
||||
if config.algorithm == "capability":
|
||||
return api.llm_classifier(api.LlmClassifierConfig.capability(config=capability_config(api, config.capability)))
|
||||
if config.algorithm == "custom":
|
||||
custom: Final = config.custom
|
||||
if custom is None:
|
||||
raise ValueError("custom routing requires custom configuration")
|
||||
return api.llm_classifier(
|
||||
api.LlmClassifierConfig.custom(
|
||||
default_target=custom.default_target,
|
||||
config=api.CustomClassifierConfig(
|
||||
prompt=custom.prompt,
|
||||
response_schema=custom.response_schema,
|
||||
selector=custom.selector,
|
||||
session_affinity=custom.session_affinity,
|
||||
message_hash_fallback=custom.message_hash_fallback,
|
||||
recent_turn_window=custom.recent_turn_window,
|
||||
max_output_tokens=custom.max_output_tokens,
|
||||
),
|
||||
)
|
||||
)
|
||||
escalation: Final = config.escalation
|
||||
release: Final = escalation.deescalation
|
||||
deescalation: Final = (
|
||||
api.DeescalationConfig(
|
||||
strong_min_calls=release.strong_min_calls,
|
||||
confirmations=release.confirmations,
|
||||
strong_max_calls=release.strong_max_calls,
|
||||
weak_cooldown_calls=release.weak_cooldown_calls,
|
||||
)
|
||||
if release is not None
|
||||
else None
|
||||
)
|
||||
return api.llm_classifier(
|
||||
api.LlmClassifierConfig.escalation(
|
||||
config=api.EscalationClassifierConfig(
|
||||
confirmations=escalation.confirmations,
|
||||
recent_turn_window=escalation.recent_turn_window,
|
||||
window_message_chars=escalation.window_message_chars,
|
||||
deescalation=deescalation,
|
||||
max_output_tokens=escalation.max_output_tokens,
|
||||
prompt=escalation.prompt,
|
||||
response_format_type=escalation.response_format_type,
|
||||
)
|
||||
)
|
||||
)
|
||||
154
litellm/router_strategy/switchyard_router/switchyard_router.py
Normal file
154
litellm/router_strategy/switchyard_router/switchyard_router.py
Normal file
|
|
@ -0,0 +1,154 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import Final, TypeAlias
|
||||
from uuid import uuid4
|
||||
|
||||
from litellm.exceptions import (
|
||||
APIConnectionError,
|
||||
APIError,
|
||||
ContextWindowExceededError,
|
||||
InternalServerError,
|
||||
RateLimitError,
|
||||
ServiceUnavailableError,
|
||||
Timeout,
|
||||
)
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.types.router import PreRoutingHookResponse
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
from .config import SwitchyardRouterConfig
|
||||
from .native import Aggregate, NativeApi, build_algorithm, load_native
|
||||
from .translation import chat_request, normalize_request, normalize_response
|
||||
|
||||
Executor: TypeAlias = Callable[[str, Mapping[str, object]], Awaitable[ModelResponse | CustomStreamWrapper]]
|
||||
|
||||
|
||||
def eligible_fallback(error: Exception) -> bool:
|
||||
return isinstance(
|
||||
error,
|
||||
(
|
||||
APIConnectionError,
|
||||
Timeout,
|
||||
RateLimitError,
|
||||
ContextWindowExceededError,
|
||||
InternalServerError,
|
||||
ServiceUnavailableError,
|
||||
),
|
||||
) or (isinstance(error, APIError) and error.status_code >= 500)
|
||||
|
||||
|
||||
class _Execution:
|
||||
def __init__(
|
||||
self, config: SwitchyardRouterConfig, api: NativeApi, request: Mapping[str, object], executor: Executor
|
||||
) -> None:
|
||||
self.config = config
|
||||
self.api = api
|
||||
self.request = request
|
||||
self.executor = executor
|
||||
self.responses: Mapping[str, ModelResponse] = {}
|
||||
self.provider_error: Exception | None = None
|
||||
|
||||
async def call_candidates(
|
||||
self,
|
||||
candidates: Sequence[str],
|
||||
request: Mapping[str, object],
|
||||
*,
|
||||
stream: bool,
|
||||
fallback_request: Mapping[str, object] | None = None,
|
||||
intermediate: bool = False,
|
||||
) -> ModelResponse | CustomStreamWrapper:
|
||||
if not candidates:
|
||||
raise ValueError("Switchyard returned no model candidates")
|
||||
target: Final = candidates[0]
|
||||
if not any(target in models for models in self.config.models.values()):
|
||||
raise ValueError("Switchyard selected an unconfigured model")
|
||||
payload: Final = chat_request(
|
||||
request,
|
||||
self.request,
|
||||
model=target,
|
||||
stream=stream,
|
||||
judge=intermediate and target in self.config.models.get("judge", ()) and bool(request.get("instructions")),
|
||||
)
|
||||
try:
|
||||
return await self.executor(target, payload)
|
||||
except Exception as error:
|
||||
if not self.config.fallback_enabled or len(candidates) == 1 or not eligible_fallback(error):
|
||||
raise
|
||||
return await self.call_candidates(
|
||||
candidates[1:], fallback_request or request, stream=stream, intermediate=intermediate
|
||||
)
|
||||
|
||||
async def respond(self, candidates: Sequence[str], request: Mapping[str, object]) -> Aggregate:
|
||||
try:
|
||||
response: Final = await self.call_candidates(candidates, request, stream=False, intermediate=True)
|
||||
except Exception as error:
|
||||
self.provider_error = error
|
||||
if isinstance(error, ContextWindowExceededError):
|
||||
raise self.api.ContextWindowExceededError(str(error)) from error
|
||||
raise
|
||||
self.provider_error = None
|
||||
if not isinstance(response, ModelResponse):
|
||||
raise TypeError("Switchyard intermediate calls must return aggregate chat responses")
|
||||
response_id: Final = str(uuid4())
|
||||
self.responses = {**self.responses, response_id: response}
|
||||
return self.api.LlmResponse.Agg(normalize_response(response, response_id))
|
||||
|
||||
def completed_response(self, aggregate: Aggregate) -> ModelResponse:
|
||||
response_id: Final = aggregate.response.get("id")
|
||||
if not isinstance(response_id, str) or response_id not in self.responses:
|
||||
raise ValueError("Switchyard returned an unknown completed response")
|
||||
return self.responses[response_id]
|
||||
|
||||
|
||||
class SwitchyardAutoRouter:
|
||||
def __init__(self, config: SwitchyardRouterConfig, *, native_api: NativeApi | None = None) -> None:
|
||||
self.config = config
|
||||
self.session_namespace: Final[str] = str(uuid4())
|
||||
self._api = native_api if native_api is not None else load_native()
|
||||
self._algorithm = build_algorithm(self._api, config)
|
||||
|
||||
async def async_pre_routing_hook(
|
||||
self,
|
||||
model: str,
|
||||
request_kwargs: Mapping[str, object],
|
||||
messages: Sequence[Mapping[str, object]] | None = None,
|
||||
input: str | Sequence[object] | None = None,
|
||||
specific_deployment: bool | None = False,
|
||||
) -> PreRoutingHookResponse | None:
|
||||
raise ValueError("Switchyard supports async chat completions only")
|
||||
|
||||
async def execute(
|
||||
self, request: Mapping[str, object], executor: Executor, headers: Mapping[str, str] | None = None
|
||||
) -> ModelResponse | CustomStreamWrapper:
|
||||
normalized: Final = normalize_request(request)
|
||||
execution: Final = _Execution(self.config, self._api, request, executor)
|
||||
try:
|
||||
return await self._run(execution, normalized, request.get("stream") is True, headers)
|
||||
except self._api.LibsyError:
|
||||
if execution.provider_error is not None:
|
||||
raise execution.provider_error
|
||||
raise
|
||||
|
||||
async def _run(
|
||||
self,
|
||||
execution: _Execution,
|
||||
normalized: Mapping[str, object],
|
||||
stream: bool,
|
||||
headers: Mapping[str, str] | None,
|
||||
) -> ModelResponse | CustomStreamWrapper:
|
||||
async for step in self._algorithm.run_stream(normalized, self.config.models, headers=headers):
|
||||
if isinstance(step, self._api.Step.CallModel):
|
||||
await step.call.respond(execution.respond(step.call.models, step.call.request))
|
||||
elif isinstance(step, self._api.Step.Done):
|
||||
if step.outcome.response is not None:
|
||||
return execution.completed_response(step.outcome.response)
|
||||
return await execution.call_candidates(
|
||||
step.outcome.selected_model_ids,
|
||||
step.outcome.request,
|
||||
stream=stream,
|
||||
fallback_request=step.outcome.fallback_request,
|
||||
)
|
||||
else:
|
||||
raise TypeError("Switchyard returned a step unsupported by the configured chat routing algorithms")
|
||||
raise ValueError("Switchyard ended without a routing outcome")
|
||||
226
litellm/router_strategy/switchyard_router/translation.py
Normal file
226
litellm/router_strategy/switchyard_router/translation.py
Normal file
|
|
@ -0,0 +1,226 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
_MAPPING: Final = TypeAdapter(Mapping[str, object])
|
||||
_SEQUENCE: Final = TypeAdapter(tuple[object, ...])
|
||||
_CONTROLLED: Final = frozenset(
|
||||
{
|
||||
"model",
|
||||
"messages",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"top_k",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"response_format",
|
||||
"reasoning_effort",
|
||||
"stream",
|
||||
"stream_options",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def mapping(value: object) -> Mapping[str, object]:
|
||||
return _MAPPING.validate_python(value)
|
||||
|
||||
|
||||
def sequence(value: object) -> tuple[object, ...]:
|
||||
return _SEQUENCE.validate_python(value)
|
||||
|
||||
|
||||
def text_blocks(content: object) -> tuple[Mapping[str, object], ...]:
|
||||
if content is None:
|
||||
return ()
|
||||
if isinstance(content, str):
|
||||
return ({"type": "text", "text": content},)
|
||||
blocks: Final = tuple(mapping(block) for block in sequence(content))
|
||||
if any(block.get("type") != "text" or not isinstance(block.get("text"), str) for block in blocks):
|
||||
raise ValueError("Switchyard currently supports text and function-tool chat content only")
|
||||
return blocks
|
||||
|
||||
|
||||
def normalized_tool_call(raw: object) -> Mapping[str, object]:
|
||||
call: Final = mapping(raw)
|
||||
function: Final = mapping(call.get("function"))
|
||||
arguments: Final = function.get("arguments")
|
||||
if call.get("type") != "function" or not isinstance(arguments, str):
|
||||
raise ValueError("Switchyard supports function tool calls with JSON arguments only")
|
||||
return {"type": "tool_call", "id": call["id"], "name": function["name"], "arguments": json.loads(arguments)}
|
||||
|
||||
|
||||
def normalized_message(raw: object) -> Mapping[str, object]:
|
||||
message: Final = mapping(raw)
|
||||
role: Final = message.get("role")
|
||||
blocks: Final = text_blocks(message.get("content"))
|
||||
if role == "tool":
|
||||
return {
|
||||
"role": "tool",
|
||||
"content": [
|
||||
{"type": "tool_result", "tool_call_id": message["tool_call_id"], "content": blocks, "is_error": None}
|
||||
],
|
||||
}
|
||||
if role not in ("system", "developer", "user", "assistant"):
|
||||
raise ValueError("Switchyard received an unsupported chat role")
|
||||
calls: Final = tuple(normalized_tool_call(call) for call in sequence(message.get("tool_calls") or ()))
|
||||
refusal: Final = message.get("refusal")
|
||||
refusals: Final = ({"type": "refusal", "text": refusal},) if isinstance(refusal, str) else ()
|
||||
return {"role": role, "content": (*blocks, *calls, *refusals)}
|
||||
|
||||
|
||||
def normalized_tool(raw: object) -> Mapping[str, object]:
|
||||
tool: Final = mapping(raw)
|
||||
if tool.get("type") != "function":
|
||||
raise ValueError("Switchyard currently supports function tools only")
|
||||
function: Final = mapping(tool.get("function"))
|
||||
return {
|
||||
"name": function["name"],
|
||||
"description": function.get("description"),
|
||||
"parameters": function.get("parameters", {}),
|
||||
"strict": function.get("strict"),
|
||||
}
|
||||
|
||||
|
||||
def normalized_tool_choice(raw: object) -> Mapping[str, object] | None:
|
||||
if raw is None:
|
||||
return None
|
||||
if raw in ("auto", "none", "required"):
|
||||
return {"type": raw}
|
||||
choice: Final = mapping(raw)
|
||||
if choice.get("type") != "function":
|
||||
raise ValueError("Switchyard currently supports function tool_choice only")
|
||||
return {"type": "tool", "data": {"name": mapping(choice.get("function"))["name"]}}
|
||||
|
||||
|
||||
def normalize_request(request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
messages: Final = tuple(normalized_message(message) for message in sequence(request.get("messages")))
|
||||
if not messages:
|
||||
raise ValueError("Switchyard requires at least one message")
|
||||
if request.get("n", 1) != 1:
|
||||
raise ValueError("Switchyard requires n=1")
|
||||
return {
|
||||
"model": "auto",
|
||||
"instructions": [],
|
||||
"messages": messages,
|
||||
"tools": tuple(normalized_tool(tool) for tool in sequence(request.get("tools") or ())),
|
||||
"tool_choice": normalized_tool_choice(request.get("tool_choice")),
|
||||
"sampling": {
|
||||
"temperature": request.get("temperature"),
|
||||
"top_p": request.get("top_p"),
|
||||
"top_k": request.get("top_k"),
|
||||
},
|
||||
"output": {
|
||||
"max_output_tokens": request.get("max_completion_tokens", request.get("max_tokens")),
|
||||
"response_format": request.get("response_format"),
|
||||
},
|
||||
"reasoning": {"effort": request.get("reasoning_effort"), "raw": None},
|
||||
"stream": False,
|
||||
}
|
||||
|
||||
|
||||
def chat_message(raw: object, originals: Sequence[object]) -> Mapping[str, object]:
|
||||
message: Final = mapping(raw)
|
||||
for original in originals:
|
||||
if json.dumps(normalized_message(original), sort_keys=True) == json.dumps(message, sort_keys=True):
|
||||
return mapping(original)
|
||||
blocks: Final = tuple(mapping(block) for block in sequence(message.get("content", ())))
|
||||
text: Final = "".join(str(block["text"]) for block in blocks if block.get("type") == "text")
|
||||
calls: Final = tuple(
|
||||
{
|
||||
"id": block["id"],
|
||||
"type": "function",
|
||||
"function": {"name": block["name"], "arguments": json.dumps(block["arguments"])},
|
||||
}
|
||||
for block in blocks
|
||||
if block.get("type") == "tool_call"
|
||||
)
|
||||
if message.get("role") == "tool":
|
||||
if len(blocks) != 1 or blocks[0].get("type") != "tool_result":
|
||||
raise ValueError("Switchyard produced an unsupported tool result")
|
||||
result: Final = blocks[0]
|
||||
return {
|
||||
"role": "tool",
|
||||
"tool_call_id": result["tool_call_id"],
|
||||
"content": tuple(mapping(block) for block in sequence(result["content"])),
|
||||
}
|
||||
if any(block.get("type") not in ("text", "tool_call", "refusal") for block in blocks):
|
||||
raise ValueError("Switchyard produced unsupported normalized message content")
|
||||
refusals: Final = tuple(str(block["text"]) for block in blocks if block.get("type") == "refusal")
|
||||
return {
|
||||
"role": message["role"],
|
||||
"content": text or None,
|
||||
**({"tool_calls": calls} if calls else {}),
|
||||
**({"refusal": "".join(refusals)} if refusals else {}),
|
||||
}
|
||||
|
||||
|
||||
def chat_tool_choice(raw: object) -> object:
|
||||
if raw is None:
|
||||
return None
|
||||
choice: Final = mapping(raw)
|
||||
if choice.get("type") in ("auto", "none", "required"):
|
||||
return choice["type"]
|
||||
if choice.get("type") == "tool":
|
||||
return {"type": "function", "function": mapping(choice["data"])}
|
||||
raise ValueError("Switchyard produced unsupported tool_choice")
|
||||
|
||||
|
||||
def chat_request(
|
||||
normalized: Mapping[str, object], original: Mapping[str, object], *, model: str, stream: bool, judge: bool
|
||||
) -> Mapping[str, object]:
|
||||
originals: Final = sequence(original.get("messages", ()))
|
||||
instructions: Final = sequence(normalized.get("instructions", ()))
|
||||
messages: Final = sequence(normalized.get("messages", ()))
|
||||
sampling: Final = mapping(normalized.get("sampling", {}))
|
||||
output: Final = mapping(normalized.get("output", {}))
|
||||
reasoning: Final = mapping(normalized.get("reasoning", {}))
|
||||
token_key: Final = "max_completion_tokens" if not judge and "max_completion_tokens" in original else "max_tokens"
|
||||
options: Final[Mapping[str, object]] = {
|
||||
**sampling,
|
||||
token_key: output.get("max_output_tokens"),
|
||||
"response_format": output.get("response_format"),
|
||||
"reasoning_effort": reasoning.get("effort"),
|
||||
"tool_choice": chat_tool_choice(normalized.get("tool_choice")),
|
||||
}
|
||||
tools: Final = tuple(
|
||||
{"type": "function", "function": {k: v for k, v in mapping(tool).items() if v is not None}}
|
||||
for tool in sequence(normalized.get("tools", ()))
|
||||
)
|
||||
passthrough: Final = {} if judge else {key: value for key, value in original.items() if key not in _CONTROLLED}
|
||||
return {
|
||||
**passthrough,
|
||||
**{key: value for key, value in options.items() if value is not None},
|
||||
"model": model,
|
||||
"messages": [chat_message(message, originals) for message in (*instructions, *messages)],
|
||||
"stream": stream,
|
||||
**({"tools": tools} if tools else {}),
|
||||
**({"stream_options": original["stream_options"]} if stream and "stream_options" in original else {}),
|
||||
}
|
||||
|
||||
|
||||
def normalize_response(response: ModelResponse, response_id: str) -> Mapping[str, object]:
|
||||
body: Final = mapping(BaseModel.model_dump(response))
|
||||
choices: Final = tuple(mapping(choice) for choice in sequence(body.get("choices")))
|
||||
return {"id": response_id, "model": body.get("model"), "outputs": tuple(normalized_output(c) for c in choices)}
|
||||
|
||||
|
||||
def normalized_output(choice: Mapping[str, object]) -> Mapping[str, object]:
|
||||
reasons: Final = {
|
||||
"stop": "end_turn",
|
||||
"length": "max_tokens",
|
||||
"tool_calls": "tool_use",
|
||||
"content_filter": "content_filter",
|
||||
}
|
||||
reason: Final = choice.get("finish_reason")
|
||||
return {
|
||||
**normalized_message(choice.get("message")),
|
||||
"stop_reason": reasons.get(str(reason), "unknown") if reason is not None else None,
|
||||
}
|
||||
|
|
@ -217,6 +217,7 @@ class SpecializedRouterOptions:
|
|||
complexity_router_default_model: str | None = None
|
||||
adaptive_router_config: Mapping[str, object] | None = None
|
||||
adaptive_router_default_model: str | None = None
|
||||
switchyard_router_config: Mapping[str, object] | None = None
|
||||
quality_router_config: Mapping[str, object] | None = None
|
||||
quality_router_default_model: str | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -676,6 +676,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams):
|
|||
adaptive_router_default_model: str | None = None
|
||||
adaptive_router_config: dict | None = None
|
||||
# quality-router params
|
||||
switchyard_router_config: Mapping[str, object] | None = None
|
||||
quality_router_config: dict | None = None
|
||||
quality_router_default_model: str | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -3284,6 +3284,7 @@ ClassifierFailureReason = Literal[
|
|||
|
||||
|
||||
InternalCallOrigin = Literal[
|
||||
"switchyard_router",
|
||||
"autorouter_classifier",
|
||||
"autorouter_compaction",
|
||||
"shadow_eval_router",
|
||||
|
|
@ -3314,7 +3315,7 @@ class StandardLoggingRoutingDecision(TypedDict, total=False):
|
|||
"""Per-request provenance for a pre-routing strategy (auto-router) decision."""
|
||||
|
||||
router_model_name: str
|
||||
router_type: ReadOnly[Literal["complexity", "adaptive", "quality", "semantic"]]
|
||||
router_type: ReadOnly[Literal["complexity", "adaptive", "quality", "semantic", "switchyard"]]
|
||||
router_config_id: ReadOnly[str]
|
||||
router_config_updated_at: ReadOnly[str]
|
||||
router_config_fingerprint: ReadOnly[str]
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ Documentation = "https://docs.litellm.ai"
|
|||
# sub-packages and the opentelemetry trio move in lockstep, and grpcio is
|
||||
# supply-chain-pinned to a vetted, aged release.
|
||||
[project.optional-dependencies]
|
||||
switchyard = ["nemo-switchyard==0.4.0a1"]
|
||||
proxy = [
|
||||
"gunicorn>=23.0.0,<24.0",
|
||||
"uvicorn>=0.33.0,<1.0",
|
||||
|
|
@ -202,6 +203,7 @@ litellm-proxy = "litellm_cli.main:litellm_proxy_cli"
|
|||
|
||||
[dependency-groups]
|
||||
dev = [
|
||||
"nemo-switchyard==0.4.0a1",
|
||||
"numpy>=1.26.0,<3.0",
|
||||
"diff-cover==9.7.2",
|
||||
"hypothesis==6.165.10",
|
||||
|
|
|
|||
61
tests/unit/router_strategy/switchyard_router/test_config.py
Normal file
61
tests/unit/router_strategy/switchyard_router/test_config.py
Normal file
|
|
@ -0,0 +1,61 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra",
|
||||
[
|
||||
{"weights": [1]},
|
||||
{"weights": [-1, 2]},
|
||||
{"weights": [0, 0]},
|
||||
{"weights": [float("inf"), 1]},
|
||||
{"models": {"any": []}},
|
||||
{"models": {"any": [""]}},
|
||||
{"seed": -1},
|
||||
{"unknown": True},
|
||||
],
|
||||
)
|
||||
def test_invalid_random_configuration_is_rejected(extra: dict[str, object]) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
SwitchyardRouterConfig.model_validate({"algorithm": "random", "models": {"any": ["a", "b"]}, **extra})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra",
|
||||
[
|
||||
{"capability": {"message_hash_fallback": True}},
|
||||
{"capability": {"base_threshold": 1.1}},
|
||||
{"escalation": {"deescalation": {"strong_min_calls": 5, "strong_max_calls": 2, "confirmations": 1}}},
|
||||
{"stage": {"deescalation_note": "return"}},
|
||||
{"stage": {"tool_semantics": {"unknown": ["tool"]}}},
|
||||
],
|
||||
)
|
||||
def test_inconsistent_classifier_configuration_is_rejected(extra: dict[str, object]) -> None:
|
||||
with pytest.raises(ValidationError):
|
||||
SwitchyardRouterConfig.model_validate(
|
||||
{
|
||||
"models": {"efficient": ["a"], "capable": ["b"], "judge": ["j"]},
|
||||
**extra,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_custom_configuration_roundtrip_keeps_schema_and_fallback_order() -> None:
|
||||
source: Final[dict[str, object]] = {
|
||||
"algorithm": "custom",
|
||||
"models": {"judge": ["j"], "code": ["a", "b"]},
|
||||
"custom": {
|
||||
"prompt": "Choose a target",
|
||||
"selector": "/target",
|
||||
"default_target": "code",
|
||||
"response_schema": {"type": "object", "properties": {"target": {"type": "string"}}},
|
||||
},
|
||||
}
|
||||
parsed: Final = SwitchyardRouterConfig.model_validate(source)
|
||||
restored: Final = SwitchyardRouterConfig.model_validate_json(parsed.model_dump_json())
|
||||
assert restored == parsed
|
||||
assert restored.models["code"] == ("a", "b")
|
||||
143
tests/unit/router_strategy/switchyard_router/test_execution.py
Normal file
143
tests/unit/router_strategy/switchyard_router/test_execution.py
Normal file
|
|
@ -0,0 +1,143 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.switchyard_router.execution import (
|
||||
SwitchyardReplay,
|
||||
child_call_kwargs,
|
||||
scoped_switchyard_headers,
|
||||
)
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
def test_child_calls_preserve_caller_without_sharing_log_identity_or_budget_reservation():
|
||||
parent: Final = {
|
||||
"metadata": {"user_api_key_hash": "caller", "user_api_key_budget_reservation": "reserved"},
|
||||
"litellm_call_id": "parent",
|
||||
"litellm_logging_obj": object(),
|
||||
"litellm_trace_id": "trace",
|
||||
}
|
||||
payload: Final = {"model": "fast", "messages": [{"role": "user", "content": "hello"}], "stream": False}
|
||||
first: Final = child_call_kwargs(parent, payload)
|
||||
second: Final = child_call_kwargs(parent, payload)
|
||||
assert first["litellm_call_id"] != second["litellm_call_id"]
|
||||
assert first["litellm_trace_id"] == "trace"
|
||||
assert "litellm_logging_obj" not in first
|
||||
assert first["metadata"] == {"user_api_key_hash": "caller", "internal_call_origin": "switchyard_router"}
|
||||
assert parent["metadata"]["user_api_key_budget_reservation"] == "reserved"
|
||||
|
||||
|
||||
def test_session_cache_keys_are_scoped_and_never_forward_credentials():
|
||||
request: Final = {
|
||||
"metadata": {"user_api_key_hash": "first"},
|
||||
"proxy_server_request": {"headers": {"authorization": "secret", "x-session-id": "same-session"}},
|
||||
}
|
||||
first: Final = scoped_switchyard_headers(request)
|
||||
other: Final = scoped_switchyard_headers({**request, "metadata": {"user_api_key_hash": "second"}})
|
||||
assert tuple(first) == ("x-switchyard-session-id",)
|
||||
assert first == scoped_switchyard_headers(request)
|
||||
assert first != other
|
||||
assert "secret" not in str(first)
|
||||
assert "same-session" not in str(first)
|
||||
|
||||
|
||||
def test_message_hash_fallback_is_scoped_without_client_session_header():
|
||||
request: Final = {"metadata": {"user_api_key_hash": "first"}, "messages": [{"role": "user", "content": "same"}]}
|
||||
first: Final = scoped_switchyard_headers(request, True)
|
||||
assert first == scoped_switchyard_headers(request, True)
|
||||
assert first != scoped_switchyard_headers({**request, "metadata": {"user_api_key_hash": "second"}}, True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("include_usage", [True, False])
|
||||
async def test_buffered_stream_replays_content_and_preserves_original_usage(include_usage: bool):
|
||||
response: Final = ModelResponse(
|
||||
id="answer",
|
||||
model="fast",
|
||||
created=123,
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hello world"}, "finish_reason": "stop"}],
|
||||
usage={"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6},
|
||||
)
|
||||
replay: Final = SwitchyardReplay(response, include_usage)
|
||||
frames: Final = [frame async for frame in replay]
|
||||
assert "".join(choice.delta.content or "" for frame in frames for choice in frame.choices) == "hello world"
|
||||
assert any(choice.finish_reason == "stop" for frame in frames for choice in frame.choices)
|
||||
if include_usage:
|
||||
assert frames[-1].usage == response.usage
|
||||
assert frames[-1].choices == []
|
||||
else:
|
||||
assert all(getattr(frame, "usage", None) is None for frame in frames)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure", [False, True])
|
||||
async def test_stream_finalizes_once_on_completion_or_failure(failure: bool):
|
||||
from collections.abc import AsyncIterator
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.router_strategy.switchyard_router.execution import FinalizingStream
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
async def source() -> AsyncIterator[ModelResponseStream]:
|
||||
yield ModelResponseStream(id="first")
|
||||
if failure:
|
||||
raise ValueError("provider failed")
|
||||
|
||||
stream: Final = source()
|
||||
finalize: Final = AsyncMock()
|
||||
wrapped: Final = FinalizingStream(stream, finalize, stream.aclose, {}, None)
|
||||
assert (await anext(wrapped)).id == "first"
|
||||
with pytest.raises(ValueError if failure else StopAsyncIteration):
|
||||
await anext(wrapped)
|
||||
await wrapped.aclose()
|
||||
finalize.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_finalizes_on_disconnect_before_first_chunk():
|
||||
from collections.abc import AsyncIterator
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.router_strategy.switchyard_router.execution import FinalizingStream
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
async def source() -> AsyncIterator[ModelResponseStream]:
|
||||
yield ModelResponseStream(id="unread")
|
||||
|
||||
stream: Final = source()
|
||||
finalize: Final = AsyncMock()
|
||||
wrapped: Final = FinalizingStream(stream, finalize, stream.aclose, {}, None)
|
||||
await wrapped.aclose()
|
||||
with pytest.raises(StopAsyncIteration):
|
||||
await anext(wrapped)
|
||||
finalize.assert_awaited_once()
|
||||
|
||||
|
||||
def test_judge_payload_does_not_inherit_parent_generation_controls():
|
||||
parent: Final = {
|
||||
"tools": [{"type": "function", "function": {"name": "run"}}],
|
||||
"tool_choice": "required",
|
||||
"max_completion_tokens": 1000,
|
||||
"response_format": {"type": "json_object"},
|
||||
"temperature": 0.9,
|
||||
"timeout": 10,
|
||||
"litellm_trace_id": "trace",
|
||||
}
|
||||
child: Final = child_call_kwargs(parent, {"messages": [{"role": "user", "content": "classify"}]})
|
||||
assert not {"tools", "tool_choice", "max_completion_tokens", "response_format", "temperature"}.intersection(child)
|
||||
assert child["timeout"] == 10
|
||||
assert child["litellm_trace_id"] == "trace"
|
||||
|
||||
|
||||
def test_sdk_sessions_are_stable_within_router_and_isolated_between_routers():
|
||||
request: Final = {"litellm_session_id": "conversation"}
|
||||
first: Final = scoped_switchyard_headers(request, sdk_namespace="router-one")
|
||||
assert first == scoped_switchyard_headers(request, sdk_namespace="router-one")
|
||||
assert first != scoped_switchyard_headers(request, sdk_namespace="router-two")
|
||||
|
||||
|
||||
def test_proxy_missing_identity_cannot_reuse_sdk_scope():
|
||||
request: Final = {"litellm_session_id": "conversation", "proxy_server_request": {"headers": {}}}
|
||||
assert scoped_switchyard_headers(request, sdk_namespace="router") != scoped_switchyard_headers(
|
||||
request, sdk_namespace="router"
|
||||
)
|
||||
|
|
@ -0,0 +1,253 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.exceptions import APIError, AuthenticationError, ServiceUnavailableError
|
||||
from litellm.router_strategy.switchyard_router import SwitchyardAutoRouter, SwitchyardRouterConfig
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
def response(model: str, text: str) -> ModelResponse:
|
||||
return ModelResponse(
|
||||
model=model,
|
||||
choices=[{"message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
|
||||
usage={"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7, "provider_extra_usage": {"count": 2}},
|
||||
)
|
||||
|
||||
|
||||
class ScriptedProvider:
|
||||
def __init__(self, replies: Mapping[str, ModelResponse | Exception]):
|
||||
self.replies = replies
|
||||
self.calls: tuple[tuple[str, Mapping[str, object]], ...] = ()
|
||||
|
||||
async def __call__(self, model: str, request: Mapping[str, object]) -> ModelResponse:
|
||||
self.calls = (*self.calls, (model, request))
|
||||
reply: Final = self.replies[model]
|
||||
if isinstance(reply, Exception):
|
||||
raise reply
|
||||
return reply
|
||||
|
||||
|
||||
def classifier_config(algorithm: str, **extra: object) -> SwitchyardRouterConfig:
|
||||
return SwitchyardRouterConfig.model_validate(
|
||||
{
|
||||
"algorithm": algorithm,
|
||||
"models": {"efficient": ["fast"], "capable": ["strong"], "judge": ["judge"]},
|
||||
**extra,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_weighted_routing_executes_only_selected_target_and_preserves_response() -> None:
|
||||
answer: Final = response("fast", "done")
|
||||
provider: Final = ScriptedProvider({"fast": answer})
|
||||
router: Final = SwitchyardAutoRouter(
|
||||
SwitchyardRouterConfig.model_validate(
|
||||
{
|
||||
"algorithm": "random",
|
||||
"models": {"any": ["fast", "slow"]},
|
||||
"weights": [1, 0],
|
||||
"seed": 42,
|
||||
}
|
||||
)
|
||||
)
|
||||
actual: Final = await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider)
|
||||
assert actual is answer
|
||||
assert tuple(model for model, _ in provider.calls) == ("fast",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("probability", "selected"), [(0.9, "fast"), (0.1, "strong")])
|
||||
async def test_capability_judge_selects_model_and_does_not_receive_user_stop(probability: float, selected: str) -> None:
|
||||
answer: Final = response(selected, "answer")
|
||||
provider: Final = ScriptedProvider(
|
||||
{
|
||||
"judge": response(
|
||||
"judge",
|
||||
'{"p_solve":'
|
||||
+ str(probability)
|
||||
+ ',"crux":"bounded task","primary_rule":"SUP-1","capability_boundary":"supported"}',
|
||||
),
|
||||
selected: answer,
|
||||
}
|
||||
)
|
||||
router: Final = SwitchyardAutoRouter(classifier_config("capability"))
|
||||
actual: Final = await router.execute({"messages": [{"role": "user", "content": "hello"}], "stop": ["}"]}, provider)
|
||||
assert actual is answer
|
||||
assert tuple(model for model, _ in provider.calls) == ("judge", selected)
|
||||
assert "stop" not in provider.calls[0][1]
|
||||
assert provider.calls[1][1]["stop"] == ["}"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("verdict", "selected", "calls"),
|
||||
[
|
||||
('{"escalate":false,"category":"none","new_evidence":false,"reason":"progress"}', "fast", ("fast", "judge")),
|
||||
(
|
||||
'{"escalate":true,"category":"repetition","new_evidence":true,"reason":"stuck"}',
|
||||
"strong",
|
||||
("fast", "judge", "strong"),
|
||||
),
|
||||
],
|
||||
)
|
||||
async def test_escalation_reuses_accepted_answer_or_executes_capable_model(
|
||||
verdict: str, selected: str, calls: tuple[str, ...]
|
||||
) -> None:
|
||||
efficient: Final = response("fast", "draft")
|
||||
capable: Final = response("strong", "answer")
|
||||
provider: Final = ScriptedProvider({"judge": response("judge", verdict), "fast": efficient, "strong": capable})
|
||||
router: Final = SwitchyardAutoRouter(classifier_config("escalation", escalation={"confirmations": 1}))
|
||||
actual: Final = await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider)
|
||||
assert actual is (efficient if selected == "fast" else capable)
|
||||
assert tuple(model for model, _ in provider.calls) == calls
|
||||
assert actual.usage.model_dump()["provider_extra_usage"] == {"count": 2}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_custom_classifier_selects_runtime_group() -> None:
|
||||
answer: Final = response("code", "result")
|
||||
provider: Final = ScriptedProvider({"judge": response("judge", '{"target":"coding"}'), "code": answer})
|
||||
router: Final = SwitchyardAutoRouter(
|
||||
SwitchyardRouterConfig.model_validate(
|
||||
{
|
||||
"algorithm": "custom",
|
||||
"models": {"judge": ["judge"], "coding": ["code"], "general": ["fast"]},
|
||||
"custom": {
|
||||
"prompt": "Choose target",
|
||||
"selector": "/target",
|
||||
"default_target": "general",
|
||||
"response_schema": {
|
||||
"type": "object",
|
||||
"required": ["target"],
|
||||
"properties": {"target": {"type": "string", "enum": ["coding", "general"]}},
|
||||
},
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
actual: Final = await router.execute({"messages": [{"role": "user", "content": "write code"}]}, provider)
|
||||
assert actual is answer
|
||||
assert tuple(model for model, _ in provider.calls) == ("judge", "code")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"error", [APIError(503, "unavailable", "fake", "fast"), ServiceUnavailableError("unavailable", "fake", "fast")]
|
||||
)
|
||||
async def test_service_failure_uses_configured_fallback(error: Exception) -> None:
|
||||
answer: Final = response("backup", "answer")
|
||||
provider: Final = ScriptedProvider({"fast": error, "backup": answer})
|
||||
router: Final = SwitchyardAutoRouter(
|
||||
classifier_config(
|
||||
"stage",
|
||||
models={"efficient": ["fast", "backup"], "capable": ["strong"]},
|
||||
)
|
||||
)
|
||||
actual: Final = await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider)
|
||||
assert actual is answer
|
||||
assert tuple(model for model, _ in provider.calls) == ("fast", "backup")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_authentication_failure_is_not_retried_on_another_model() -> None:
|
||||
error: Final = AuthenticationError("denied", "fake", "fast")
|
||||
provider: Final = ScriptedProvider({"fast": error})
|
||||
router: Final = SwitchyardAutoRouter(
|
||||
classifier_config(
|
||||
"stage",
|
||||
models={"efficient": ["fast", "backup"], "capable": ["strong"]},
|
||||
)
|
||||
)
|
||||
with pytest.raises(AuthenticationError) as raised:
|
||||
await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider)
|
||||
assert raised.value is error
|
||||
assert tuple(model for model, _ in provider.calls) == ("fast",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_intermediate_authentication_error_preserves_public_error_type() -> None:
|
||||
error: Final = AuthenticationError("denied", "fake", "fast")
|
||||
provider: Final = ScriptedProvider({"fast": error})
|
||||
router: Final = SwitchyardAutoRouter(classifier_config("escalation", escalation={"confirmations": 1}))
|
||||
with pytest.raises(AuthenticationError) as raised:
|
||||
await router.execute({"messages": [{"role": "user", "content": "hello"}]}, provider)
|
||||
assert raised.value is error
|
||||
assert tuple(model for model, _ in provider.calls) == ("fast",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_escalation_confirmations_are_isolated_between_sessions() -> None:
|
||||
efficient: Final = response("fast", "draft")
|
||||
capable: Final = response("strong", "answer")
|
||||
provider: Final = ScriptedProvider(
|
||||
{
|
||||
"fast": efficient,
|
||||
"strong": capable,
|
||||
"judge": response(
|
||||
"judge", '{"escalate":true,"category":"repetition","new_evidence":true,"reason":"stuck"}'
|
||||
),
|
||||
}
|
||||
)
|
||||
router: Final = SwitchyardAutoRouter(classifier_config("escalation", escalation={"confirmations": 2}))
|
||||
first: Final = await router.execute(
|
||||
{"messages": [{"role": "user", "content": "first"}]}, provider, headers={"x-switchyard-session-id": "session-a"}
|
||||
)
|
||||
separate: Final = await router.execute(
|
||||
{"messages": [{"role": "user", "content": "second"}]},
|
||||
provider,
|
||||
headers={"x-switchyard-session-id": "session-b"},
|
||||
)
|
||||
confirmed: Final = await router.execute(
|
||||
{"messages": [{"role": "user", "content": "third"}]}, provider, headers={"x-switchyard-session-id": "session-a"}
|
||||
)
|
||||
assert (first is efficient, separate is efficient, confirmed is capable) == (True, True, True)
|
||||
assert tuple(model for model, _ in provider.calls) == ("fast", "judge", "fast", "judge", "fast", "judge", "strong")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_escalation_buffers_intermediate_calls_for_stream_requests() -> None:
|
||||
answer: Final = response("fast", "draft")
|
||||
provider: Final = ScriptedProvider(
|
||||
{
|
||||
"fast": answer,
|
||||
"judge": response("judge", '{"escalate":false,"category":"none","new_evidence":false,"reason":"progress"}'),
|
||||
}
|
||||
)
|
||||
router: Final = SwitchyardAutoRouter(classifier_config("escalation", escalation={"confirmations": 1}))
|
||||
actual: Final = await router.execute({"messages": [{"role": "user", "content": "hello"}], "stream": True}, provider)
|
||||
assert actual is answer
|
||||
assert tuple((model, payload["stream"]) for model, payload in provider.calls) == (("fast", False), ("judge", False))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelling_judge_cancels_provider_work() -> None:
|
||||
started: Final = asyncio.Event()
|
||||
cancelled: Final = asyncio.Event()
|
||||
never: Final = asyncio.Event()
|
||||
|
||||
async def pending_judge(model: str, request: Mapping[str, object]) -> ModelResponse:
|
||||
started.set()
|
||||
try:
|
||||
await never.wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
return response(model, "unreachable")
|
||||
|
||||
router: Final = SwitchyardAutoRouter(classifier_config("capability"))
|
||||
task: Final = asyncio.create_task(
|
||||
router.execute({"messages": [{"role": "user", "content": "hello"}]}, pending_judge)
|
||||
)
|
||||
entered: Final = asyncio.create_task(started.wait())
|
||||
done, _ = await asyncio.wait((task, entered), return_when=asyncio.FIRST_COMPLETED)
|
||||
if task in done:
|
||||
entered.cancel()
|
||||
await task
|
||||
pytest.fail("Routing finished before the judge executed")
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
assert cancelled.is_set()
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.router_strategy.switchyard_router.translation import chat_request, normalize_request, normalize_response
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
|
||||
def test_chat_roundtrip_preserves_tools_history_and_provider_options() -> None:
|
||||
request: Final[dict[str, object]] = {
|
||||
"messages": [
|
||||
{"role": "system", "content": "Be precise"},
|
||||
{"role": "user", "content": "Read the file", "name": "reader"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{"id": "call-1", "type": "function", "function": {"name": "read", "arguments": '{"path":"a"}'}}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call-1", "content": "contents"},
|
||||
],
|
||||
"tools": [{"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": {"type": "function", "function": {"name": "read"}},
|
||||
"temperature": 0.2,
|
||||
"max_completion_tokens": 50,
|
||||
"parallel_tool_calls": False,
|
||||
"stop": ["END"],
|
||||
"metadata": {"trace": "abc"},
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
actual: Final = chat_request(normalize_request(request), request, model="fast", stream=True, judge=False)
|
||||
assert actual == {
|
||||
**request,
|
||||
"model": "fast",
|
||||
"tools": ({"type": "function", "function": {"name": "read", "parameters": {"type": "object"}}},),
|
||||
}
|
||||
|
||||
|
||||
def test_text_response_with_null_tool_calls_normalizes() -> None:
|
||||
response: Final = ModelResponse(
|
||||
model="fast",
|
||||
choices=[
|
||||
{
|
||||
"message": {"role": "assistant", "content": "answer", "tool_calls": None},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
)
|
||||
assert normalize_response(response, "internal-id") == {
|
||||
"id": "internal-id",
|
||||
"model": "fast",
|
||||
"outputs": (
|
||||
{"role": "assistant", "content": ({"type": "text", "text": "answer"},), "stop_reason": "end_turn"},
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def test_null_tools_are_normalized_to_empty_catalog() -> None:
|
||||
request: Final = normalize_request({"messages": [{"role": "user", "content": "hello"}], "tools": None})
|
||||
assert request["tools"] == ()
|
||||
|
||||
|
||||
def test_judge_rewrite_removes_user_tools_limits_and_stop_sequences() -> None:
|
||||
request: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": "task"}],
|
||||
"stop": ["}"],
|
||||
"max_completion_tokens": 1,
|
||||
"tools": [{"type": "function", "function": {"name": "read", "parameters": {}}}],
|
||||
}
|
||||
rewritten: Final[dict[str, object]] = {
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "Judge task"}]}],
|
||||
"instructions": [{"role": "system", "content": [{"type": "text", "text": "Return JSON"}]}],
|
||||
"output": {"max_output_tokens": 200, "response_format": {"type": "json_object"}},
|
||||
}
|
||||
assert chat_request(rewritten, request, model="judge", stream=False, judge=True) == {
|
||||
"model": "judge",
|
||||
"stream": False,
|
||||
"messages": [{"role": "system", "content": "Return JSON"}, {"role": "user", "content": "Judge task"}],
|
||||
"max_tokens": 200,
|
||||
"response_format": {"type": "json_object"},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content",
|
||||
[
|
||||
[{"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}],
|
||||
[{"type": "input_audio", "input_audio": {"data": "abc"}}],
|
||||
],
|
||||
)
|
||||
def test_unsupported_content_is_rejected_before_execution(content: object) -> None:
|
||||
with pytest.raises(ValueError, match="text and function-tool"):
|
||||
normalize_request({"messages": [{"role": "user", "content": content}]})
|
||||
|
|
@ -790,3 +790,20 @@ def test_inferred_custom_tiers_claim_existing_gateway_capability(
|
|||
) -> None:
|
||||
capabilities: Final = claimed_capabilities({"classifier_type": classifier, "tiers": tiers})
|
||||
assert any(capability.key == "tier_or_classifier_prompt" for capability in capabilities) is expected
|
||||
|
||||
|
||||
def test_switchyard_dependencies_include_targets_and_judges():
|
||||
params: Final = {
|
||||
"model": "auto_router/switchyard",
|
||||
"switchyard_router_config": {"models": {"efficient": ["fast"], "capable": ["strong"], "judge": ["judge"]}},
|
||||
}
|
||||
dependencies: Final = strategy_router_dependencies(params)
|
||||
assert {(entry.model_name, entry.role) for entry in dependencies} == {
|
||||
("fast", "tier"),
|
||||
("strong", "tier"),
|
||||
("judge", "classifier"),
|
||||
}
|
||||
assert (
|
||||
validate_strategy_router_model_write("auto_router/switchyard", frozenset({"switchyard_router_config"})) is None
|
||||
)
|
||||
assert validate_strategy_router_model_write("auto_router/switchyard", frozenset()) is not None
|
||||
|
|
|
|||
|
|
@ -116,6 +116,8 @@ export interface AutoRouterDeployment extends AutoRouterCandidateDeployment {
|
|||
auto_router_embedding_model?: string | null;
|
||||
adaptive_router_config?: unknown;
|
||||
adaptive_router_default_model?: string | null;
|
||||
switchyard_router_config?: unknown;
|
||||
switchyard_router_default_model?: string | null;
|
||||
quality_router_config?: unknown;
|
||||
quality_router_default_model?: string | null;
|
||||
} | null;
|
||||
|
|
|
|||
|
|
@ -43,6 +43,27 @@ const semanticDeployment = {
|
|||
};
|
||||
|
||||
describe("autoRouterRows", () => {
|
||||
it("shows Switchyard targets and permits editing database-backed routers", () => {
|
||||
const deployment = {
|
||||
model_name: "switchyard",
|
||||
model_info: { id: "switchyard-id", db_model: true },
|
||||
litellm_params: {
|
||||
model: "auto_router/switchyard_router",
|
||||
switchyard_router_config: {
|
||||
algorithm: "capability",
|
||||
models: { efficient: ["small"], capable: ["large", "small"], judge: ["judge"] },
|
||||
},
|
||||
},
|
||||
};
|
||||
const row = toAutoRouterRow(deployment, 0, ADMIN, null);
|
||||
const expected = {
|
||||
kind: "switchyard",
|
||||
typeLabel: "Switchyard",
|
||||
canEdit: true,
|
||||
targets: ["small", "large", "judge"],
|
||||
};
|
||||
expect(row).toMatchObject(expected);
|
||||
});
|
||||
it("classifies a complexity router and unions its tier models as targets", () => {
|
||||
const row = toAutoRouterRow(complexityDeployment, 0, ADMIN, null);
|
||||
|
||||
|
|
|
|||
|
|
@ -83,6 +83,10 @@ const configManaged = (label: string, config: Record<string, unknown>): Presenta
|
|||
|
||||
/** How each strategy renders itself, given its own config object. */
|
||||
const PRESENTERS: Record<AutoRouterKind, (config: Record<string, unknown>) => Presentation> = {
|
||||
switchyard: (config) => ({
|
||||
typeLabel: "Switchyard",
|
||||
targets: dedupe(Object.values(asRecord(config.models)).flatMap(asStringArray)),
|
||||
}),
|
||||
complexity: (config) => ({
|
||||
typeLabel: complexityTypeLabel(config),
|
||||
targets: dedupe(Object.values(asRecord(config.tiers)).flatMap(normalizeTierModels)),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,316 @@
|
|||
import React, { useId, useState } from "react";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { Textarea } from "@/components/ui/textarea";
|
||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import {
|
||||
SWITCHYARD_ALGORITHMS,
|
||||
SWITCHYARD_DEFAULTS,
|
||||
restoreSwitchyardDefaults,
|
||||
switchyardRecord,
|
||||
switchyardRequiredGroups,
|
||||
type SwitchyardAlgorithm,
|
||||
type SwitchyardDraft,
|
||||
} from "./switchyard_config";
|
||||
|
||||
type Setting = {
|
||||
key: string;
|
||||
label: string;
|
||||
kind: "number" | "boolean" | "text" | "json" | "format" | "picker";
|
||||
hint?: string;
|
||||
};
|
||||
const CLASSIFIER_SETTINGS: Setting[] = [
|
||||
{ key: "session_affinity", label: "Session affinity", kind: "boolean" },
|
||||
{ key: "message_hash_fallback", label: "Message hash fallback", kind: "boolean" },
|
||||
{ key: "recent_turn_window", label: "Recent turn window", kind: "number", hint: "Leave empty to use all turns" },
|
||||
{ key: "max_output_tokens", label: "Classifier output token limit", kind: "number" },
|
||||
];
|
||||
const PROMPT_SETTINGS: Setting[] = [
|
||||
{ key: "prompt", label: "Classifier prompt", kind: "text", hint: "Leave empty to use the built-in prompt" },
|
||||
{ key: "response_format_type", label: "Classifier response format", kind: "format" },
|
||||
];
|
||||
const SETTINGS: Record<SwitchyardAlgorithm, Setting[]> = {
|
||||
random: [
|
||||
{
|
||||
key: "weights",
|
||||
label: "Model weights (JSON)",
|
||||
kind: "json",
|
||||
hint: "One weight per model in the any group in the order above. Leave empty for equal weights",
|
||||
},
|
||||
{ key: "seed", label: "Random seed", kind: "number", hint: "Optional nonnegative integer" },
|
||||
],
|
||||
capability: [
|
||||
{
|
||||
key: "base_threshold",
|
||||
label: "Capability threshold",
|
||||
kind: "number",
|
||||
hint: "LiteLLM preset 0.5, between 0 and 1",
|
||||
},
|
||||
{ key: "threshold_step", label: "Threshold step", kind: "number" },
|
||||
...CLASSIFIER_SETTINGS,
|
||||
...PROMPT_SETTINGS,
|
||||
],
|
||||
custom: [
|
||||
{ key: "prompt", label: "Custom classifier prompt", kind: "text" },
|
||||
{ key: "response_schema", label: "Response schema (JSON)", kind: "json" },
|
||||
{ key: "selector", label: "Target selector", kind: "text" },
|
||||
{ key: "default_target", label: "Default target group", kind: "text" },
|
||||
...CLASSIFIER_SETTINGS,
|
||||
],
|
||||
escalation: [
|
||||
{ key: "confirmations", label: "Escalation confirmations", kind: "number" },
|
||||
{ key: "recent_turn_window", label: "Recent turn window", kind: "number" },
|
||||
{ key: "window_message_chars", label: "Characters per history message", kind: "number" },
|
||||
{ key: "max_output_tokens", label: "Classifier output token limit", kind: "number" },
|
||||
...PROMPT_SETTINGS,
|
||||
{
|
||||
key: "deescalation",
|
||||
label: "De-escalation (JSON)",
|
||||
kind: "json",
|
||||
hint: "Optional object with strong_min_calls, confirmations, strong_max_calls and weak_cooldown_calls. Leave empty to disable",
|
||||
},
|
||||
],
|
||||
stage: [
|
||||
{ key: "picker", label: "Stage picker", kind: "picker" },
|
||||
{
|
||||
key: "confidence_threshold",
|
||||
label: "Confidence threshold",
|
||||
kind: "number",
|
||||
hint: "LiteLLM preset 0.5, between 0 and 1",
|
||||
},
|
||||
{ key: "recent_window", label: "Recent stage window", kind: "number" },
|
||||
{ key: "escalation_note", label: "Escalation note", kind: "text" },
|
||||
{ key: "deescalation_note", label: "De-escalation note", kind: "text" },
|
||||
{ key: "only_on_wrong_signal_escalation", label: "Escalate only on wrong signal", kind: "boolean" },
|
||||
{ key: "capable_system_prompt", label: "Capable model system prompt", kind: "text" },
|
||||
{ key: "efficient_system_prompt", label: "Efficient model system prompt", kind: "text" },
|
||||
{
|
||||
key: "tool_semantics",
|
||||
label: "Tool semantics (JSON)",
|
||||
kind: "json",
|
||||
hint: "Map tool categories to lists of tool names",
|
||||
},
|
||||
{
|
||||
key: "classifier",
|
||||
label: "Fallback capability classifier (JSON)",
|
||||
kind: "json",
|
||||
hint: "Optional capability configuration. Requires a judge model group",
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
function SettingInput({
|
||||
setting,
|
||||
id,
|
||||
value,
|
||||
onChange,
|
||||
}: {
|
||||
setting: Setting;
|
||||
id: string;
|
||||
value: unknown;
|
||||
onChange: (next: unknown) => void;
|
||||
}) {
|
||||
if (setting.kind === "boolean")
|
||||
return (
|
||||
<input id={id} type="checkbox" checked={Boolean(value)} onChange={(event) => onChange(event.target.checked)} />
|
||||
);
|
||||
if (setting.kind === "number")
|
||||
return (
|
||||
<Input
|
||||
id={id}
|
||||
type="number"
|
||||
step="any"
|
||||
value={typeof value === "number" ? value : ""}
|
||||
onChange={(event) => onChange(event.target.value === "" ? null : Number(event.target.value))}
|
||||
/>
|
||||
);
|
||||
if (setting.kind === "format" || setting.kind === "picker") {
|
||||
const options = setting.kind === "format" ? ["json_schema", "json_object"] : ["efficient_first", "capable_first"];
|
||||
return (
|
||||
<Select value={String(value ?? "")} onValueChange={(next) => next && onChange(next)}>
|
||||
<SelectTrigger id={id}>
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{options.map((option) => (
|
||||
<SelectItem key={option} value={option}>
|
||||
{option}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
);
|
||||
}
|
||||
const textValue = typeof value === "string" ? value : JSON.stringify(value, null, 2);
|
||||
return (
|
||||
<Textarea
|
||||
id={id}
|
||||
value={value == null ? "" : textValue}
|
||||
onChange={(event) => onChange(event.target.value || null)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
export default function SwitchyardConfig({
|
||||
value,
|
||||
onChange,
|
||||
modelNames = [],
|
||||
}: {
|
||||
value: SwitchyardDraft;
|
||||
onChange: (value: SwitchyardDraft) => void;
|
||||
modelNames?: string[];
|
||||
}) {
|
||||
const id = useId();
|
||||
const [newGroup, setNewGroup] = useState("");
|
||||
const algorithm = value.algorithm as SwitchyardAlgorithm;
|
||||
const knownAlgorithm = SWITCHYARD_ALGORITHMS.includes(algorithm);
|
||||
const config = algorithm === "random" ? value : switchyardRecord(value[algorithm]);
|
||||
const defaults = SWITCHYARD_DEFAULTS[algorithm] ?? {};
|
||||
const models = switchyardRecord(value.models);
|
||||
const requiredGroups = switchyardRequiredGroups(algorithm, Boolean(switchyardRecord(value.stage).classifier));
|
||||
const groups = Array.from(new Set([...requiredGroups, ...Object.keys(models)]));
|
||||
const update = (key: string, next: unknown) =>
|
||||
onChange(
|
||||
algorithm === "random" ? { ...value, [key]: next } : { ...value, [algorithm]: { ...config, [key]: next } },
|
||||
);
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
<div className="space-y-2">
|
||||
<Label htmlFor={`${id}-algorithm`}>Switchyard algorithm</Label>
|
||||
<Select
|
||||
value={String(value.algorithm ?? "")}
|
||||
onValueChange={(next) => next && onChange({ ...value, algorithm: next })}
|
||||
>
|
||||
<SelectTrigger id={`${id}-algorithm`} className="w-full">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{SWITCHYARD_ALGORITHMS.map((item) => (
|
||||
<SelectItem key={item} value={item}>
|
||||
{item.charAt(0).toUpperCase() + item.slice(1)}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Switching algorithms keeps your other settings. Model names must refer to configured LiteLLM models.
|
||||
</p>
|
||||
</div>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Supports chat completions. Escalation may buffer responses before streaming the accepted answer.
|
||||
</p>
|
||||
<fieldset className="space-y-3">
|
||||
<legend className="font-medium">Model groups</legend>
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Enter ordered, comma-separated model names for each group. Later models are fallback candidates. Stage routing
|
||||
only needs a judge when its fallback classifier is enabled.
|
||||
</p>
|
||||
<datalist id={`${id}-models`}>
|
||||
{modelNames.map((name) => (
|
||||
<option key={name} value={name} />
|
||||
))}
|
||||
</datalist>
|
||||
{groups.map((group) => (
|
||||
<div key={group} className="space-y-1">
|
||||
<Label htmlFor={`${id}-group-${group}`}>{group} models</Label>
|
||||
<div className="flex gap-2">
|
||||
<Input
|
||||
id={`${id}-group-${group}`}
|
||||
list={`${id}-models`}
|
||||
value={Array.isArray(models[group]) ? models[group].join(", ") : ""}
|
||||
placeholder="model-name, fallback-model"
|
||||
onChange={(event) =>
|
||||
onChange({
|
||||
...value,
|
||||
models: { ...models, [group]: event.target.value.split(",").map((name) => name.trim()) },
|
||||
})
|
||||
}
|
||||
/>
|
||||
{!requiredGroups.includes(group) && (
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
aria-label={`Remove ${group} group`}
|
||||
onClick={() =>
|
||||
onChange({
|
||||
...value,
|
||||
models: Object.fromEntries(Object.entries(models).filter(([name]) => name !== group)),
|
||||
})
|
||||
}
|
||||
>
|
||||
Remove
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
<div className="flex gap-2">
|
||||
<Input
|
||||
aria-label="New model group name"
|
||||
value={newGroup}
|
||||
onChange={(event) => setNewGroup(event.target.value)}
|
||||
placeholder="New target group name"
|
||||
/>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
disabled={!newGroup.trim() || groups.includes(newGroup.trim())}
|
||||
onClick={() => {
|
||||
onChange({ ...value, models: { ...models, [newGroup.trim()]: [] } });
|
||||
setNewGroup("");
|
||||
}}
|
||||
>
|
||||
Add group
|
||||
</Button>
|
||||
</div>
|
||||
</fieldset>
|
||||
<details className="rounded-md border p-4">
|
||||
<summary className="cursor-pointer font-medium">Advanced settings</summary>
|
||||
<div className="mt-4 space-y-4">
|
||||
<p className="text-xs text-muted-foreground">
|
||||
Switchyard defaults are shown unless noted as LiteLLM presets. Empty optional fields use the default
|
||||
behavior.
|
||||
</p>
|
||||
<Label>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={value.fallback_enabled !== false}
|
||||
onChange={(event) => onChange({ ...value, fallback_enabled: event.target.checked })}
|
||||
/>{" "}
|
||||
Enable model failure fallback
|
||||
</Label>
|
||||
{(SETTINGS[algorithm] ?? []).map((setting) => {
|
||||
const current = config[setting.key] === undefined ? defaults[setting.key] : config[setting.key];
|
||||
const fieldId = `${id}-${setting.key}`;
|
||||
return (
|
||||
<div key={setting.key} className="space-y-1">
|
||||
<Label htmlFor={fieldId}>{setting.label}</Label>
|
||||
<SettingInput
|
||||
setting={setting}
|
||||
id={fieldId}
|
||||
value={current}
|
||||
onChange={(next) => update(setting.key, next)}
|
||||
/>
|
||||
{setting.hint && <p className="text-xs text-muted-foreground">{setting.hint}</p>}
|
||||
{setting.key === "prompt" && algorithm !== "custom" && (
|
||||
<Button type="button" variant="link" onClick={() => update("prompt", null)}>
|
||||
Restore built-in prompt
|
||||
</Button>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
disabled={!knownAlgorithm}
|
||||
onClick={() => onChange(restoreSwitchyardDefaults(value))}
|
||||
>
|
||||
Restore defaults
|
||||
</Button>
|
||||
</div>
|
||||
</details>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
import React from "react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
import { fireEvent, renderWithProviders, screen, waitFor, chooseSelectOption } from "../../../tests/test-utils";
|
||||
import SwitchyardStrategyPicker from "./SwitchyardStrategyPicker";
|
||||
|
||||
const { create } = vi.hoisted(() => ({ create: vi.fn().mockResolvedValue({}) }));
|
||||
vi.mock("../networking", () => ({ modelCreateCall: create }));
|
||||
vi.mock("../llm_calls/fetch_models", () => ({
|
||||
fetchAvailableModels: vi.fn().mockResolvedValue([]),
|
||||
fetchAutoRouterModels: vi.fn().mockResolvedValue([]),
|
||||
}));
|
||||
|
||||
const props = { handleOk: vi.fn(), accessToken: "test", userRole: "Admin" };
|
||||
|
||||
describe("Switchyard creation", () => {
|
||||
it("creates the Switchyard strategy with basic models and advanced tuning", async () => {
|
||||
renderWithProviders(
|
||||
<SwitchyardStrategyPicker {...props}>
|
||||
<p>Complexity setup</p>
|
||||
</SwitchyardStrategyPicker>,
|
||||
);
|
||||
await chooseSelectOption(userEvent.setup(), screen.getByLabelText("Routing strategy"), "Switchyard");
|
||||
fireEvent.change(screen.getByLabelText("efficient models"), { target: { value: "small" } });
|
||||
fireEvent.change(screen.getByLabelText("capable models"), { target: { value: "large, backup" } });
|
||||
fireEvent.change(screen.getByLabelText("judge models"), { target: { value: "judge" } });
|
||||
fireEvent.click(screen.getByText("Advanced settings"));
|
||||
fireEvent.change(screen.getByLabelText("Capability threshold"), { target: { value: "0.6" } });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Add Auto Router" }));
|
||||
await waitFor(() =>
|
||||
expect(create).toHaveBeenCalledWith("test", {
|
||||
model_name: "switchyard-router",
|
||||
litellm_params: {
|
||||
model: "auto_router/switchyard_router",
|
||||
switchyard_router_config: {
|
||||
algorithm: "capability",
|
||||
models: { efficient: ["small"], capable: ["large", "backup"], judge: ["judge"] },
|
||||
fallback_enabled: true,
|
||||
capability: { base_threshold: 0.6 },
|
||||
},
|
||||
},
|
||||
model_info: {},
|
||||
}),
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
import React, { useState } from "react";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { Card, CardContent } from "@/components/ui/card";
|
||||
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { canCreateAutoRouterForTeam, canModifyModel } from "@/utils/modelPermissions";
|
||||
import TeamDropdown from "../common_components/team_dropdown";
|
||||
import { modelCreateCall } from "../networking";
|
||||
import { fetchAutoRouterModels, fetchAvailableModels } from "../llm_calls/fetch_models";
|
||||
import SwitchyardConfig from "./SwitchyardConfig";
|
||||
import { newSwitchyardConfig, switchyardConfigError, switchyardConfigSchema } from "./switchyard_config";
|
||||
import type { AddAutoRouterTabProps } from "./add_auto_router_tab";
|
||||
|
||||
function AddSwitchyardRouter({
|
||||
accessToken,
|
||||
handleOk,
|
||||
createScope,
|
||||
userId,
|
||||
userRole,
|
||||
teams = null,
|
||||
}: AddAutoRouterTabProps) {
|
||||
const [name, setName] = useState("switchyard-router");
|
||||
const [teamId, setTeamId] = useState<string | null>(null);
|
||||
const [config, setConfig] = useState(newSwitchyardConfig);
|
||||
const [saving, setSaving] = useState(false);
|
||||
const requiresTeam = createScope === "team-required";
|
||||
const actor = { userRole, userID: userId ?? null, isViewOnly: false };
|
||||
const { data: models = [] } = useQuery({
|
||||
queryKey: ["availableModels", "switchyard", accessToken, teamId],
|
||||
queryFn: () => (requiresTeam ? fetchAutoRouterModels(accessToken, teamId) : fetchAvailableModels(accessToken)),
|
||||
enabled: Boolean(accessToken) && (!requiresTeam || Boolean(teamId)),
|
||||
});
|
||||
const teamError = requiresTeam && !teamId ? "Select a team to continue" : null;
|
||||
const nameError = name.trim() ? null : "Auto router name is required";
|
||||
const memberManaged = requiresTeam && Boolean(teamId) && !canModifyModel(actor, teams, { teamId, isDbModel: true });
|
||||
const permissionError = memberManaged ? "Switchyard routers must be created by a team or proxy administrator" : null;
|
||||
const error = teamError ?? permissionError ?? nameError ?? switchyardConfigError(config);
|
||||
const save = async () => {
|
||||
if (error || saving) return;
|
||||
setSaving(true);
|
||||
try {
|
||||
const payload = {
|
||||
model_name: name,
|
||||
litellm_params: {
|
||||
model: "auto_router/switchyard_router",
|
||||
switchyard_router_config: switchyardConfigSchema.parse(config),
|
||||
},
|
||||
model_info: requiresTeam ? { team_id: teamId } : {},
|
||||
};
|
||||
await modelCreateCall(accessToken, payload);
|
||||
toast.success(`Successfully created Auto Router: ${name}`);
|
||||
handleOk();
|
||||
} catch (cause) {
|
||||
toast.fromError(cause);
|
||||
} finally {
|
||||
setSaving(false);
|
||||
}
|
||||
};
|
||||
return (
|
||||
<Card>
|
||||
<CardContent>
|
||||
<form
|
||||
className="space-y-4"
|
||||
onSubmit={(event) => {
|
||||
event.preventDefault();
|
||||
void save();
|
||||
}}
|
||||
>
|
||||
<div className="space-y-2">
|
||||
<Label htmlFor="switchyard-create-name">Auto Router Name</Label>
|
||||
<Input id="switchyard-create-name" value={name} onChange={(event) => setName(event.target.value)} />
|
||||
</div>
|
||||
{requiresTeam && (
|
||||
<div className="space-y-2">
|
||||
<Label htmlFor="switchyard-team">Select Team</Label>
|
||||
<TeamDropdown
|
||||
id="switchyard-team"
|
||||
value={teamId}
|
||||
onChange={setTeamId}
|
||||
filterTeam={(team) => canCreateAutoRouterForTeam(actor, team)}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
<SwitchyardConfig
|
||||
value={config}
|
||||
onChange={setConfig}
|
||||
modelNames={models.filter((model) => model.mode !== "embedding").map((model) => model.model_group)}
|
||||
/>
|
||||
{error && (
|
||||
<p role="alert" className="text-sm text-destructive">
|
||||
{error}
|
||||
</p>
|
||||
)}
|
||||
<Button type="submit" disabled={Boolean(error) || saving}>
|
||||
Add Auto Router
|
||||
</Button>
|
||||
</form>
|
||||
</CardContent>
|
||||
</Card>
|
||||
);
|
||||
}
|
||||
|
||||
export default function SwitchyardStrategyPicker({
|
||||
children,
|
||||
...props
|
||||
}: AddAutoRouterTabProps & { children: React.ReactNode }) {
|
||||
const [strategy, setStrategy] = useState("complexity");
|
||||
return (
|
||||
<div className="space-y-4">
|
||||
<div className="space-y-2">
|
||||
<Label htmlFor="auto-router-strategy">Routing strategy</Label>
|
||||
<Select value={strategy} onValueChange={(next) => next && setStrategy(next)}>
|
||||
<SelectTrigger id="auto-router-strategy" className="w-full">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
<SelectItem value="complexity">Complexity</SelectItem>
|
||||
<SelectItem value="switchyard">Switchyard</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
{strategy === "switchyard" ? <AddSwitchyardRouter {...props} /> : children}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import SwitchyardStrategyPicker from "./SwitchyardStrategyPicker";
|
||||
import { AutoRouterAvailabilityContext, useAutoRouterAvailability } from "./AutoRouterAvailability";
|
||||
import { decisionDeploymentName } from "./jev_classifier_config";
|
||||
import DecisionModelSelect from "./DecisionModelSelect";
|
||||
|
|
@ -91,7 +92,7 @@ import {
|
|||
type FamilyPreference,
|
||||
} from "./auto_setup";
|
||||
|
||||
interface AddAutoRouterTabProps {
|
||||
export interface AddAutoRouterTabProps {
|
||||
handleOk: () => void;
|
||||
accessToken: string;
|
||||
userRole: string;
|
||||
|
|
@ -843,4 +844,10 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
|
|||
);
|
||||
};
|
||||
|
||||
export default AddAutoRouterTab;
|
||||
const AddAutoRouterWithStrategy: React.FC<AddAutoRouterTabProps> = (props) => (
|
||||
<SwitchyardStrategyPicker {...props}>
|
||||
<AddAutoRouterTab {...props} />
|
||||
</SwitchyardStrategyPicker>
|
||||
);
|
||||
|
||||
export default AddAutoRouterWithStrategy;
|
||||
|
|
|
|||
|
|
@ -13,10 +13,12 @@
|
|||
* prefixes are matched first, and only a bare `auto_router/<name>` is the semantic router.
|
||||
*/
|
||||
|
||||
export type AutoRouterKind = "complexity" | "adaptive" | "quality" | "semantic";
|
||||
export type AutoRouterKind = "complexity" | "adaptive" | "quality" | "semantic" | "switchyard";
|
||||
|
||||
export interface AutoRouterParams {
|
||||
model?: string | null;
|
||||
switchyard_router_config?: unknown;
|
||||
switchyard_router_default_model?: string | null;
|
||||
complexity_router_config?: unknown;
|
||||
complexity_router_default_model?: string | null;
|
||||
auto_router_config?: unknown;
|
||||
|
|
@ -42,6 +44,14 @@ const startsWith = (params: AutoRouterParams, prefix: string): boolean => params
|
|||
|
||||
/** Ordered; the semantic entry matches anything left and must stay last. */
|
||||
export const AUTO_ROUTER_STRATEGIES: readonly AutoRouterStrategy[] = [
|
||||
{
|
||||
kind: "switchyard",
|
||||
label: "Switchyard",
|
||||
configKey: "switchyard_router_config",
|
||||
defaultModelKey: "switchyard_router_default_model",
|
||||
hasEditor: true,
|
||||
matches: (p) => startsWith(p, "auto_router/switchyard_router") || p.switchyard_router_config != null,
|
||||
},
|
||||
{
|
||||
kind: "complexity",
|
||||
label: "Complexity",
|
||||
|
|
@ -87,8 +97,9 @@ export const isComplexityRouter = (params: AutoRouterParams | null | undefined):
|
|||
/** Any `auto_router/*` deployment, whatever its strategy. Use for listing and filtering. */
|
||||
export const isAutoRouterDeployment = (params: AutoRouterParams | null | undefined): boolean =>
|
||||
params?.model?.startsWith("auto_router/") === true ||
|
||||
params?.complexity_router_config != null ||
|
||||
params?.auto_router_config != null;
|
||||
[params?.complexity_router_config, params?.auto_router_config, params?.switchyard_router_config].some(
|
||||
(config) => config != null,
|
||||
);
|
||||
|
||||
/**
|
||||
* Whether EditAutoRouterModal understands this deployment. It only speaks complexity and
|
||||
|
|
|
|||
|
|
@ -0,0 +1,86 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
import {
|
||||
hydrateSwitchyardConfig,
|
||||
restoreSwitchyardDefaults,
|
||||
SWITCHYARD_DEFAULTS,
|
||||
switchyardConfigError,
|
||||
switchyardConfigSchema,
|
||||
} from "./switchyard_config";
|
||||
|
||||
const models = { efficient: ["small"], capable: ["large", "backup"], judge: ["judge"] };
|
||||
|
||||
describe("Switchyard configuration", () => {
|
||||
it("preserves unknown configuration and parses JSON editing fields on save", () => {
|
||||
const stored = {
|
||||
algorithm: "escalation",
|
||||
models,
|
||||
future_option: "keep",
|
||||
escalation: { future_judge_option: 9, deescalation: { strong_min_calls: 3, confirmations: 2 } },
|
||||
custom: null,
|
||||
};
|
||||
const draft = hydrateSwitchyardConfig(JSON.stringify(stored));
|
||||
const edited = {
|
||||
...draft,
|
||||
escalation: { ...stored.escalation, confirmations: 4, deescalation: '{"strong_min_calls":5,"confirmations":3}' },
|
||||
};
|
||||
expect(switchyardConfigSchema.parse(edited)).toEqual({
|
||||
...stored,
|
||||
escalation: { ...stored.escalation, confirmations: 4, deescalation: { strong_min_calls: 5, confirmations: 3 } },
|
||||
});
|
||||
});
|
||||
|
||||
it("restores tuning defaults while preserving models and future settings", () => {
|
||||
const customConfig = {
|
||||
algorithm: "capability",
|
||||
models,
|
||||
future: true,
|
||||
capability: { base_threshold: 0.9, prompt: "override", future_setting: 4 },
|
||||
};
|
||||
const restored = restoreSwitchyardDefaults(customConfig);
|
||||
expect(restored).toMatchObject({
|
||||
models,
|
||||
future: true,
|
||||
capability: { ...SWITCHYARD_DEFAULTS.capability, future_setting: 4 },
|
||||
});
|
||||
});
|
||||
|
||||
it("validates random weights against candidates rather than group count", () => {
|
||||
expect(
|
||||
switchyardConfigError({ algorithm: "random", models: { any: ["small", "large"] }, weights: "[1, 2]" }),
|
||||
).toBeNull();
|
||||
expect(
|
||||
switchyardConfigError({ algorithm: "random", models: { any: ["small", "large"] }, weights: "[1]" }),
|
||||
).toContain("one weight per model");
|
||||
expect(switchyardConfigError({ algorithm: "random", models: { any: ["small"] }, weights: "[0]" })).toContain(
|
||||
"positive weight",
|
||||
);
|
||||
});
|
||||
|
||||
it("requires judge for classifiers but permits stage without a classifier", () => {
|
||||
const solvers = { efficient: ["small"], capable: ["large"] };
|
||||
expect(switchyardConfigError({ algorithm: "capability", models: solvers })).toContain("judge");
|
||||
expect(switchyardConfigError({ algorithm: "stage", models: solvers })).toBeNull();
|
||||
expect(switchyardConfigError({ algorithm: "stage", models: solvers, stage: { classifier: "{}" } })).toContain(
|
||||
"judge",
|
||||
);
|
||||
});
|
||||
|
||||
it("rejects invalid JSON, missing custom fields, and incompatible affinity", () => {
|
||||
expect(switchyardConfigError({ algorithm: "random", models: { any: ["small"] }, weights: "bad" })).toContain(
|
||||
"weights",
|
||||
);
|
||||
expect(switchyardConfigError({ algorithm: "custom", models })).toContain("response_schema");
|
||||
expect(
|
||||
switchyardConfigError({ algorithm: "capability", models, capability: { message_hash_fallback: true } }),
|
||||
).toContain("requires session affinity");
|
||||
expect(switchyardConfigError({ algorithm: "capability", models, capability: { base_threshold: 2 } })).toContain(
|
||||
"base_threshold",
|
||||
);
|
||||
});
|
||||
|
||||
it("omits random weights when switching to a classifier and keeps the draft intact", () => {
|
||||
const draft = { algorithm: "capability", models, weights: [1, 2] };
|
||||
expect(switchyardConfigSchema.parse(draft)).not.toHaveProperty("weights");
|
||||
expect(draft.weights).toEqual([1, 2]);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,202 @@
|
|||
import { z } from "zod";
|
||||
|
||||
export const SWITCHYARD_ALGORITHMS = ["random", "capability", "custom", "escalation", "stage"] as const;
|
||||
export type SwitchyardAlgorithm = (typeof SWITCHYARD_ALGORITHMS)[number];
|
||||
export type SwitchyardDraft = Record<string, unknown>;
|
||||
|
||||
export const switchyardRecord = (value: unknown): Record<string, unknown> =>
|
||||
value !== null && typeof value === "object" && !Array.isArray(value) ? (value as Record<string, unknown>) : {};
|
||||
|
||||
const jsonValue = (value: unknown): unknown => {
|
||||
if (typeof value !== "string") return value;
|
||||
try {
|
||||
return JSON.parse(value);
|
||||
} catch {
|
||||
return value;
|
||||
}
|
||||
};
|
||||
const positiveInteger = z.number().int().positive();
|
||||
const optionalWindow = positiveInteger.nullable().optional();
|
||||
const responseFormat = z.enum(["json_schema", "json_object"]).optional();
|
||||
const optionalPrompt = z.string().nullable().optional();
|
||||
const classifierFields = {
|
||||
session_affinity: z.boolean().optional(),
|
||||
message_hash_fallback: z.boolean().optional(),
|
||||
recent_turn_window: optionalWindow,
|
||||
max_output_tokens: positiveInteger.optional(),
|
||||
};
|
||||
const capabilityFields = {
|
||||
...classifierFields,
|
||||
base_threshold: z.number().min(0).max(1).optional(),
|
||||
threshold_step: z.number().nonnegative().optional(),
|
||||
prompt: optionalPrompt,
|
||||
response_format_type: responseFormat,
|
||||
};
|
||||
const affinityError = (
|
||||
value: { message_hash_fallback?: boolean; session_affinity?: boolean },
|
||||
ctx: z.RefinementCtx,
|
||||
) => {
|
||||
if (value.message_hash_fallback && !value.session_affinity)
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
path: ["message_hash_fallback"],
|
||||
message: "Message hash fallback requires session affinity",
|
||||
});
|
||||
};
|
||||
const capabilitySchema = z.object(capabilityFields).passthrough().superRefine(affinityError);
|
||||
const deescalationFields = {
|
||||
strong_min_calls: positiveInteger,
|
||||
confirmations: positiveInteger,
|
||||
strong_max_calls: positiveInteger.nullable().optional(),
|
||||
weak_cooldown_calls: z.number().int().nonnegative().optional(),
|
||||
};
|
||||
const deescalationSchema = z
|
||||
.object(deescalationFields)
|
||||
.passthrough()
|
||||
.superRefine((value, ctx) => {
|
||||
if (value.strong_max_calls != null && value.strong_max_calls < value.strong_min_calls)
|
||||
ctx.addIssue({ code: "custom", message: "Maximum calls must be at least minimum calls" });
|
||||
});
|
||||
const modelsSchema = z.record(z.string().min(1), z.array(z.string().trim().min(1)).min(1));
|
||||
const escalationFields = {
|
||||
confirmations: positiveInteger.optional(),
|
||||
recent_turn_window: positiveInteger.optional(),
|
||||
window_message_chars: z.number().int().min(50).optional(),
|
||||
max_output_tokens: positiveInteger.optional(),
|
||||
prompt: optionalPrompt,
|
||||
response_format_type: responseFormat,
|
||||
deescalation: z.preprocess(jsonValue, deescalationSchema.nullable().optional()),
|
||||
};
|
||||
const customFields = {
|
||||
...classifierFields,
|
||||
prompt: z.string().min(1).optional(),
|
||||
response_schema: z.preprocess(jsonValue, z.record(z.string(), z.unknown()).optional()),
|
||||
selector: z.string().min(1).optional(),
|
||||
default_target: z.string().min(1).optional(),
|
||||
};
|
||||
const stageFields = {
|
||||
picker: z.enum(["efficient_first", "capable_first"]).optional(),
|
||||
confidence_threshold: z.number().min(0).max(1).optional(),
|
||||
recent_window: optionalWindow,
|
||||
escalation_note: optionalPrompt,
|
||||
deescalation_note: optionalPrompt,
|
||||
only_on_wrong_signal_escalation: z.boolean().optional(),
|
||||
capable_system_prompt: optionalPrompt,
|
||||
efficient_system_prompt: optionalPrompt,
|
||||
tool_semantics: z.preprocess(jsonValue, z.record(z.string(), z.array(z.string())).nullable().optional()),
|
||||
classifier: z.preprocess(jsonValue, capabilitySchema.nullable().optional()),
|
||||
};
|
||||
const configFields = {
|
||||
algorithm: z.enum(SWITCHYARD_ALGORITHMS),
|
||||
models: z.preprocess(jsonValue, modelsSchema),
|
||||
weights: z.preprocess(jsonValue, z.array(z.number().nonnegative()).min(1).nullable().optional()),
|
||||
seed: z.number().int().nonnegative().nullable().optional(),
|
||||
fallback_enabled: z.boolean().optional(),
|
||||
capability: capabilitySchema.optional(),
|
||||
escalation: z.object(escalationFields).passthrough().optional(),
|
||||
custom: z.object(customFields).passthrough().superRefine(affinityError).nullable().optional(),
|
||||
stage: z.object(stageFields).passthrough().optional(),
|
||||
};
|
||||
export function switchyardRequiredGroups(algorithm: string, hasStageClassifier: boolean): string[] {
|
||||
if (algorithm === "random") return ["any"];
|
||||
if (algorithm === "custom") return ["judge"];
|
||||
if (algorithm === "stage" && !hasStageClassifier) return ["efficient", "capable"];
|
||||
return ["efficient", "capable", "judge"];
|
||||
}
|
||||
export const switchyardConfigSchema = z
|
||||
.object(configFields)
|
||||
.passthrough()
|
||||
.superRefine((config, ctx) => {
|
||||
for (const group of switchyardRequiredGroups(config.algorithm, Boolean(config.stage?.classifier))) {
|
||||
if (!config.models[group]?.length)
|
||||
ctx.addIssue({ code: "custom", path: ["models", group], message: `Select a model for ${group}` });
|
||||
}
|
||||
if (!Object.keys(config.models).length)
|
||||
ctx.addIssue({ code: "custom", path: ["models"], message: "Add at least one target group" });
|
||||
if (config.algorithm !== "random" || !config.weights) return;
|
||||
const validWeights =
|
||||
config.weights.length === config.models.any?.length && config.weights.some((weight) => weight > 0);
|
||||
if (!validWeights)
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
path: ["weights"],
|
||||
message: "Provide one weight per model in the any group, with at least one positive weight",
|
||||
});
|
||||
})
|
||||
.superRefine((config, ctx) => {
|
||||
if (config.algorithm !== "custom") return;
|
||||
for (const field of ["prompt", "response_schema", "selector", "default_target"] as const) {
|
||||
if (!config.custom?.[field])
|
||||
ctx.addIssue({ code: "custom", path: ["custom", field], message: `${field} is required for custom routing` });
|
||||
}
|
||||
if (config.custom?.default_target && !config.models[config.custom.default_target])
|
||||
ctx.addIssue({
|
||||
code: "custom",
|
||||
path: ["custom", "default_target"],
|
||||
message: "Default target must name a model group",
|
||||
});
|
||||
})
|
||||
.transform((config) =>
|
||||
config.algorithm === "random"
|
||||
? config
|
||||
: Object.fromEntries(Object.entries(config).filter(([key]) => key !== "weights")),
|
||||
);
|
||||
|
||||
export const SWITCHYARD_DEFAULTS: Record<SwitchyardAlgorithm, Record<string, unknown>> = {
|
||||
random: {},
|
||||
capability: {
|
||||
base_threshold: 0.5,
|
||||
threshold_step: 0,
|
||||
session_affinity: false,
|
||||
message_hash_fallback: false,
|
||||
recent_turn_window: null,
|
||||
max_output_tokens: 4096,
|
||||
prompt: null,
|
||||
response_format_type: "json_schema",
|
||||
},
|
||||
escalation: {
|
||||
confirmations: 2,
|
||||
recent_turn_window: 28,
|
||||
window_message_chars: 500,
|
||||
max_output_tokens: 4096,
|
||||
prompt: null,
|
||||
response_format_type: "json_schema",
|
||||
deescalation: null,
|
||||
},
|
||||
custom: { session_affinity: false, message_hash_fallback: false, recent_turn_window: null, max_output_tokens: 4096 },
|
||||
stage: {
|
||||
picker: "efficient_first",
|
||||
confidence_threshold: 0.5,
|
||||
recent_window: null,
|
||||
escalation_note: null,
|
||||
deescalation_note: null,
|
||||
only_on_wrong_signal_escalation: true,
|
||||
capable_system_prompt: null,
|
||||
efficient_system_prompt: null,
|
||||
tool_semantics: null,
|
||||
classifier: null,
|
||||
},
|
||||
};
|
||||
|
||||
export const newSwitchyardConfig = (): SwitchyardDraft => ({
|
||||
algorithm: "capability",
|
||||
models: {},
|
||||
fallback_enabled: true,
|
||||
});
|
||||
export const restoreSwitchyardDefaults = (config: SwitchyardDraft): SwitchyardDraft => {
|
||||
const algorithm = config.algorithm as SwitchyardAlgorithm;
|
||||
if (algorithm === "random") return { ...config, weights: null, seed: null, fallback_enabled: true };
|
||||
return {
|
||||
...config,
|
||||
fallback_enabled: true,
|
||||
[algorithm]: { ...switchyardRecord(config[algorithm]), ...SWITCHYARD_DEFAULTS[algorithm] },
|
||||
};
|
||||
};
|
||||
|
||||
export const hydrateSwitchyardConfig = (value: unknown): SwitchyardDraft => switchyardRecord(jsonValue(value));
|
||||
export const switchyardConfigError = (config: SwitchyardDraft): string | null => {
|
||||
const result = switchyardConfigSchema.safeParse(config);
|
||||
return result.success
|
||||
? null
|
||||
: result.error.issues.map((issue) => `${issue.path.join(".")}: ${issue.message}`).join(". ");
|
||||
};
|
||||
|
|
@ -0,0 +1,92 @@
|
|||
import React from "react";
|
||||
import { describe, expect, it, vi, beforeEach } from "vitest";
|
||||
import { fireEvent, renderWithProviders, screen, waitFor } from "../../../tests/test-utils";
|
||||
import EditSwitchyardRouter from "./EditSwitchyardRouter";
|
||||
|
||||
const { patch } = vi.hoisted(() => ({ patch: vi.fn().mockResolvedValue({}) }));
|
||||
vi.mock("../networking", () => ({ modelPatchUpdateCall: patch }));
|
||||
vi.mock("../llm_calls/fetch_models", () => ({
|
||||
fetchAvailableModels: vi.fn().mockResolvedValue([]),
|
||||
fetchAutoRouterModels: vi.fn().mockResolvedValue([]),
|
||||
}));
|
||||
const modelData = {
|
||||
model_name: "switchyard",
|
||||
model_info: { id: "router-id" },
|
||||
litellm_params: {
|
||||
model: "auto_router/switchyard_router",
|
||||
switchyard_router_config: {
|
||||
algorithm: "capability",
|
||||
models: { efficient: ["small"], capable: ["large"], judge: ["judge"] },
|
||||
capability: { base_threshold: 0.7, future_option: "preserve" },
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
beforeEach(() => vi.clearAllMocks());
|
||||
describe("Switchyard editor", () => {
|
||||
it("saves advanced edits under the strategy key and preserves unrecognized settings", async () => {
|
||||
const onSuccess = vi.fn();
|
||||
renderWithProviders(
|
||||
<EditSwitchyardRouter
|
||||
isVisible
|
||||
accessToken="test"
|
||||
modelData={modelData}
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={onSuccess}
|
||||
/>,
|
||||
);
|
||||
fireEvent.click(screen.getByText("Advanced settings"));
|
||||
fireEvent.change(screen.getByLabelText("Capability threshold"), { target: { value: "0.8" } });
|
||||
fireEvent.click(screen.getByRole("button", { name: "Save Changes" }));
|
||||
await waitFor(() =>
|
||||
expect(patch).toHaveBeenCalledWith(
|
||||
"test",
|
||||
{
|
||||
model_name: "switchyard",
|
||||
litellm_params: {
|
||||
switchyard_router_config: {
|
||||
...modelData.litellm_params.switchyard_router_config,
|
||||
capability: { base_threshold: 0.8, future_option: "preserve" },
|
||||
},
|
||||
},
|
||||
},
|
||||
"router-id",
|
||||
),
|
||||
);
|
||||
expect(onSuccess).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("blocks invalid tuning and permits restoring defaults before save", () => {
|
||||
renderWithProviders(
|
||||
<EditSwitchyardRouter
|
||||
isVisible
|
||||
accessToken="test"
|
||||
modelData={modelData}
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
fireEvent.click(screen.getByText("Advanced settings"));
|
||||
fireEvent.change(screen.getByLabelText("Capability threshold"), { target: { value: "2" } });
|
||||
expect(screen.getByRole("button", { name: "Save Changes" })).toBeDisabled();
|
||||
fireEvent.click(screen.getByRole("button", { name: "Restore defaults" }));
|
||||
expect(screen.getByRole("button", { name: "Save Changes" })).toBeEnabled();
|
||||
});
|
||||
|
||||
it("blocks unsupported member-managed updates", () => {
|
||||
renderWithProviders(
|
||||
<EditSwitchyardRouter
|
||||
isVisible
|
||||
isMemberManaged
|
||||
accessToken="test"
|
||||
modelData={modelData}
|
||||
onCancel={vi.fn()}
|
||||
onSuccess={vi.fn()}
|
||||
/>,
|
||||
);
|
||||
expect(screen.getByLabelText("Auto Router Name")).toBeDisabled();
|
||||
expect(screen.getByRole("button", { name: "Save Changes" })).toBeDisabled();
|
||||
expect(screen.getByRole("alert")).toHaveTextContent("must be edited by a team or proxy administrator");
|
||||
expect(patch).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,109 @@
|
|||
import React, { useState } from "react";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import {
|
||||
Dialog,
|
||||
DialogContent,
|
||||
DialogDescription,
|
||||
DialogFooter,
|
||||
DialogHeader,
|
||||
DialogTitle,
|
||||
} from "@/components/ui/dialog";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Input } from "@/components/ui/input";
|
||||
import { Label } from "@/components/ui/label";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { modelPatchUpdateCall } from "../networking";
|
||||
import { fetchAutoRouterModels, fetchAvailableModels } from "../llm_calls/fetch_models";
|
||||
import SwitchyardConfig from "../add_model/SwitchyardConfig";
|
||||
import { hydrateSwitchyardConfig, switchyardConfigError, switchyardConfigSchema } from "../add_model/switchyard_config";
|
||||
|
||||
interface SwitchyardDeployment {
|
||||
model_name?: string;
|
||||
litellm_params?: Record<string, unknown>;
|
||||
model_info?: { id?: string; team_id?: string | null; [key: string]: unknown };
|
||||
}
|
||||
|
||||
export default function EditSwitchyardRouter({
|
||||
isVisible,
|
||||
onCancel,
|
||||
onSuccess,
|
||||
modelData,
|
||||
accessToken,
|
||||
isMemberManaged = false,
|
||||
}: {
|
||||
isVisible: boolean;
|
||||
onCancel: () => void;
|
||||
onSuccess: (value: SwitchyardDeployment) => void;
|
||||
modelData: SwitchyardDeployment;
|
||||
accessToken: string;
|
||||
isMemberManaged?: boolean;
|
||||
}) {
|
||||
const [name, setName] = useState(modelData.model_name ?? "");
|
||||
const [config, setConfig] = useState(() =>
|
||||
hydrateSwitchyardConfig(modelData.litellm_params?.switchyard_router_config),
|
||||
);
|
||||
const [saving, setSaving] = useState(false);
|
||||
const teamId = modelData.model_info?.team_id;
|
||||
const { data: models = [] } = useQuery({
|
||||
queryKey: ["availableModels", "switchyard", accessToken, teamId],
|
||||
queryFn: () => (teamId ? fetchAutoRouterModels(accessToken, teamId) : fetchAvailableModels(accessToken)),
|
||||
enabled: isVisible && Boolean(accessToken),
|
||||
});
|
||||
const nameError = !name.trim() ? "Auto router name is required" : switchyardConfigError(config);
|
||||
const error = isMemberManaged ? "Switchyard routers must be edited by a team or proxy administrator" : nameError;
|
||||
const save = async () => {
|
||||
if (error || saving || !modelData.model_info?.id) return;
|
||||
setSaving(true);
|
||||
try {
|
||||
const params = { switchyard_router_config: switchyardConfigSchema.parse(config) };
|
||||
const patch = isMemberManaged ? { litellm_params: params } : { model_name: name, litellm_params: params };
|
||||
await modelPatchUpdateCall(accessToken, patch, modelData.model_info.id);
|
||||
onSuccess({
|
||||
...modelData,
|
||||
model_name: isMemberManaged ? modelData.model_name : name,
|
||||
litellm_params: { ...modelData.litellm_params, ...params },
|
||||
});
|
||||
toast.success("Auto router configuration updated successfully");
|
||||
onCancel();
|
||||
} catch (cause) {
|
||||
toast.fromError(cause);
|
||||
} finally {
|
||||
setSaving(false);
|
||||
}
|
||||
};
|
||||
return (
|
||||
<Dialog open={isVisible} onOpenChange={(open) => !open && onCancel()}>
|
||||
<DialogContent className="max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[760px]">
|
||||
<DialogHeader>
|
||||
<DialogTitle>Edit Switchyard auto router</DialogTitle>
|
||||
<DialogDescription>Configure the routing algorithm and model groups</DialogDescription>
|
||||
</DialogHeader>
|
||||
<Label htmlFor="switchyard-router-name">Auto Router Name</Label>
|
||||
<Input
|
||||
id="switchyard-router-name"
|
||||
value={name}
|
||||
disabled={isMemberManaged}
|
||||
onChange={(event) => setName(event.target.value)}
|
||||
/>
|
||||
<SwitchyardConfig
|
||||
value={config}
|
||||
onChange={setConfig}
|
||||
modelNames={models.filter((model) => model.mode !== "embedding").map((model) => model.model_group)}
|
||||
/>
|
||||
{error && (
|
||||
<p role="alert" className="text-sm text-destructive">
|
||||
{error}
|
||||
</p>
|
||||
)}
|
||||
<DialogFooter>
|
||||
<Button variant="outline" onClick={onCancel}>
|
||||
Cancel
|
||||
</Button>
|
||||
<Button disabled={Boolean(error) || saving} onClick={() => void save()}>
|
||||
Save Changes
|
||||
</Button>
|
||||
</DialogFooter>
|
||||
</DialogContent>
|
||||
</Dialog>
|
||||
);
|
||||
}
|
||||
|
|
@ -1,3 +1,4 @@
|
|||
import EditSwitchyardRouter from "./EditSwitchyardRouter";
|
||||
import { useQuery } from "@tanstack/react-query";
|
||||
import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
|
||||
import { autoRouterListKey, fetchAllModelDeployments } from "@/app/(dashboard)/hooks/models/useModels";
|
||||
|
|
@ -34,7 +35,7 @@ import {
|
|||
getCustomTierRowsError,
|
||||
resolveComplexityDefaultModel,
|
||||
} from "../add_model/tier_rows";
|
||||
import { isComplexityRouter } from "../add_model/auto_router_strategies";
|
||||
import { autoRouterStrategy, isComplexityRouter } from "../add_model/auto_router_strategies";
|
||||
import {
|
||||
type BuildComplexityRouterConfigParams,
|
||||
type StoredComplexityRouterConfig,
|
||||
|
|
@ -806,4 +807,11 @@ const EditAutoRouterModal: React.FC<EditAutoRouterModalProps> = ({
|
|||
);
|
||||
};
|
||||
|
||||
export default EditAutoRouterModal;
|
||||
const AutoRouterEditor: React.FC<EditAutoRouterModalProps> = (props) =>
|
||||
autoRouterStrategy(props.modelData?.litellm_params).kind === "switchyard" ? (
|
||||
props.isVisible && <EditSwitchyardRouter key={props.modelData?.model_info?.id} {...props} />
|
||||
) : (
|
||||
<EditAutoRouterModal {...props} />
|
||||
);
|
||||
|
||||
export default AutoRouterEditor;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue