mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(router): validate routing_groups at save time and keep invalid DB groups from blocking SSO load
Overlapping routing_groups persisted from the Admin UI raised inside Router._init_routing_groups during the DB config reconcile, which skipped loading SSO, guardrails and the other DB-backed settings while leaving the proxy healthy. /config/update now returns 400 for overlapping models, duplicate names, the reserved default name and unknown strategies before writing, the Router builds every group selector before replacing its state so a rejected update keeps the previous groups routing, and the proxy applies routing_groups separately from the other router settings so an already persisted invalid value is logged and skipped instead of aborting the reconcile. The Admin UI modal blocks picking a model another group owns. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
171888b716
commit
9f990c4f86
10 changed files with 501 additions and 81 deletions
|
|
@ -143,6 +143,7 @@ from litellm.router_utils.auto_router_tuning_baseline import (
|
|||
snapshot_tuning_baselines,
|
||||
tuning_limit_violation,
|
||||
)
|
||||
from litellm.router_utils.routing_groups import parse_routing_groups
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
|
|
@ -759,6 +760,7 @@ from litellm.types.router import (
|
|||
ClassifierPlugin,
|
||||
DeploymentTypedDict,
|
||||
RouterGeneralSettings,
|
||||
RoutingGroup,
|
||||
RoutingPlugin,
|
||||
SearchToolTypedDict,
|
||||
updateDeployment,
|
||||
|
|
@ -6896,7 +6898,27 @@ class ProxyConfig:
|
|||
combined_router_settings = db_router_settings.param_value
|
||||
|
||||
if combined_router_settings:
|
||||
llm_router.update_settings(**combined_router_settings)
|
||||
self._apply_router_settings(llm_router, combined_router_settings)
|
||||
|
||||
@staticmethod
|
||||
def _apply_router_settings(llm_router: Router, router_settings: Mapping[str, object]) -> None:
|
||||
"""
|
||||
`routing_groups` is applied on its own so a value persisted before
|
||||
save-time validation existed cannot abort the reconcile that also loads
|
||||
SSO, guardrails and the other DB-backed settings. The router keeps the
|
||||
groups it already holds when the new value is rejected.
|
||||
"""
|
||||
llm_router.update_settings(**{k: v for k, v in router_settings.items() if k != "routing_groups"})
|
||||
if "routing_groups" not in router_settings:
|
||||
return
|
||||
try:
|
||||
llm_router.update_settings(routing_groups=router_settings["routing_groups"])
|
||||
except (TypeError, ValueError) as invalid_groups:
|
||||
verbose_proxy_logger.error(
|
||||
"Ignoring invalid router_settings.routing_groups from config/DB, all other router settings still "
|
||||
"apply. Fix the routing groups in the Admin UI to load them: %s",
|
||||
invalid_groups,
|
||||
)
|
||||
|
||||
def _add_general_settings_from_db_config(
|
||||
self, config_data: dict, general_settings: dict, proxy_logging_obj: ProxyLogging
|
||||
|
|
@ -16903,6 +16925,12 @@ async def update_config(
|
|||
)
|
||||
},
|
||||
)
|
||||
try:
|
||||
parse_routing_groups(
|
||||
TypeAdapter(list[RoutingGroup] | None).validate_python(raw_router_settings.get("routing_groups"))
|
||||
)
|
||||
except (ValidationError, ValueError) as invalid_groups:
|
||||
raise HTTPException(status_code=400, detail={"error": str(invalid_groups)})
|
||||
|
||||
if prisma_client is None:
|
||||
raise Exception("No DB Connected")
|
||||
|
|
|
|||
|
|
@ -218,6 +218,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import (
|
|||
increment_deployment_failures_for_current_minute,
|
||||
increment_deployment_successes_for_current_minute,
|
||||
)
|
||||
from litellm.router_utils.routing_groups import parse_routing_groups, validate_routing_strategy
|
||||
from litellm.scheduler import FlowItem, Scheduler
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -1244,20 +1245,9 @@ class Router:
|
|||
return strategy.value
|
||||
return strategy
|
||||
|
||||
def _validate_routing_strategy(self, routing_strategy: RoutingStrategy | str | None) -> None:
|
||||
# See: https://github.com/BerriAI/litellm/issues/11330
|
||||
valid_strategy_strings: Final = ["simple-shuffle", "lar1"] + [s.value for s in RoutingStrategy]
|
||||
if routing_strategy is None:
|
||||
return
|
||||
is_valid_string: Final = isinstance(routing_strategy, str) and routing_strategy in valid_strategy_strings
|
||||
is_valid_enum: Final = isinstance(routing_strategy, RoutingStrategy)
|
||||
if not is_valid_string and not is_valid_enum:
|
||||
raise ValueError(
|
||||
f"Invalid routing_strategy: '{routing_strategy}'. "
|
||||
f"Valid options: {valid_strategy_strings}. "
|
||||
f"Check 'router_settings.routing_strategy' in your config.yaml "
|
||||
f"or the 'routing_strategy' parameter if using the Router SDK directly."
|
||||
)
|
||||
@staticmethod
|
||||
def _validate_routing_strategy(routing_strategy: RoutingStrategy | str | None) -> None:
|
||||
validate_routing_strategy(routing_strategy)
|
||||
|
||||
def _build_strategy_selector(
|
||||
self,
|
||||
|
|
@ -1274,11 +1264,6 @@ class Router:
|
|||
match self._normalize_strategy(strategy):
|
||||
case RoutingStrategy.LEAST_BUSY.value:
|
||||
selector = LeastBusyLoggingHandler(router_cache=self.cache)
|
||||
if register_callbacks:
|
||||
if isinstance(litellm.input_callback, list):
|
||||
litellm.logging_callback_manager.add_litellm_input_callback(selector)
|
||||
else:
|
||||
litellm.input_callback = [selector]
|
||||
case RoutingStrategy.USAGE_BASED_ROUTING.value:
|
||||
selector = LowestTPMLoggingHandler(
|
||||
router_cache=self.cache,
|
||||
|
|
@ -1302,11 +1287,21 @@ class Router:
|
|||
case _:
|
||||
pass
|
||||
|
||||
if selector is not None and register_callbacks and isinstance(litellm.callbacks, list):
|
||||
litellm.logging_callback_manager.add_litellm_callback(selector)
|
||||
if selector is not None and register_callbacks:
|
||||
self._register_router_selector(selector)
|
||||
|
||||
return selector
|
||||
|
||||
@staticmethod
|
||||
def _register_router_selector(selector: RouterStrategySelector) -> None:
|
||||
if isinstance(selector, LeastBusyLoggingHandler):
|
||||
if isinstance(litellm.input_callback, list):
|
||||
litellm.logging_callback_manager.add_litellm_input_callback(selector)
|
||||
else:
|
||||
litellm.input_callback = [selector]
|
||||
if isinstance(litellm.callbacks, list):
|
||||
litellm.logging_callback_manager.add_litellm_callback(selector)
|
||||
|
||||
def _unregister_router_selectors(self, selectors: Sequence[object]) -> None:
|
||||
"""
|
||||
Drop router-owned strategy selectors from litellm's global callback
|
||||
|
|
@ -1397,75 +1392,69 @@ class Router:
|
|||
at most one explicit group. Constructs per-group strategy selectors so
|
||||
groups with different `routing_strategy_args` track independent state.
|
||||
|
||||
Validation and selector construction run to completion before any
|
||||
router state changes, so a rejected input raises with the previously
|
||||
loaded groups still routing.
|
||||
|
||||
Models not claimed by any explicit group are served by the implicit
|
||||
`"default"` group, whose selectors are the `self.<strategy>_logger`
|
||||
attributes set up in `routing_strategy_init`.
|
||||
"""
|
||||
group_selectors: Final[Mapping[str, Mapping[str, RouterStrategySelector]]] = getattr(
|
||||
self, "_group_selectors", {}
|
||||
)
|
||||
self._unregister_router_selectors([sel for selectors in group_selectors.values() for sel in selectors.values()])
|
||||
|
||||
self._routing_groups: dict[str, RoutingGroup] = {}
|
||||
self._model_to_group: dict[str, str] = {}
|
||||
self._group_selectors: dict[str, dict[str, RouterStrategySelector]] = {}
|
||||
self._invalidate_model_group_info_cache()
|
||||
self._invalidate_access_groups_cache()
|
||||
|
||||
if not groups_input:
|
||||
self._replace_routing_groups(())
|
||||
return
|
||||
|
||||
known_model_names: Final = {m.get("model_name") for m in (self.model_list or []) if m.get("model_name")}
|
||||
known_model_names: Final = frozenset(m["model_name"] for m in (self.model_list or ()) if m.get("model_name"))
|
||||
groups: Final = parse_routing_groups(groups_input, known_model_names=known_model_names)
|
||||
|
||||
seen_group_names: Final[set] = set()
|
||||
for raw in groups_input:
|
||||
group = raw if isinstance(raw, RoutingGroup) else RoutingGroup(**raw)
|
||||
|
||||
if not group.group_name:
|
||||
raise ValueError("routing_groups: group_name must be non-empty.")
|
||||
if group.group_name == "default":
|
||||
raise ValueError("routing_groups: 'default' is reserved for the implicit fallback group.")
|
||||
if group.group_name in known_model_names or group.group_name in (self.model_group_alias or {}):
|
||||
alias_names: Final = frozenset(self.model_group_alias or ())
|
||||
for group in groups:
|
||||
if group.group_name in known_model_names or group.group_name in alias_names:
|
||||
verbose_router_logger.warning(
|
||||
"routing_groups: group_name '%s' is shadowed by an existing model_name or model_group_alias; "
|
||||
"the group's strategy still applies to its members, but the name is not callable until renamed.",
|
||||
group.group_name,
|
||||
)
|
||||
if group.group_name in seen_group_names:
|
||||
raise ValueError(
|
||||
f"routing_groups: group names must be unique, duplicate group_name '{group.group_name}'."
|
||||
)
|
||||
seen_group_names.add(group.group_name)
|
||||
|
||||
self._validate_routing_strategy(group.routing_strategy)
|
||||
|
||||
for model_name in group.models:
|
||||
if model_name in self._model_to_group:
|
||||
raise ValueError(
|
||||
f"routing_groups: model_name '{model_name}' appears in "
|
||||
f"both '{self._model_to_group[model_name]}' and "
|
||||
f"'{group.group_name}'. Each model may belong to at most one group."
|
||||
)
|
||||
if known_model_names and model_name not in known_model_names:
|
||||
verbose_router_logger.warning(
|
||||
"routing_groups: model_name '%s' (group '%s') is not in model_list; "
|
||||
"the group entry will only take effect once a deployment with that "
|
||||
"model_name is added.",
|
||||
model_name,
|
||||
group.group_name,
|
||||
)
|
||||
self._model_to_group[model_name] = group.group_name
|
||||
|
||||
self._routing_groups[group.group_name] = group
|
||||
|
||||
strategy_value = self._normalize_strategy(group.routing_strategy) or ""
|
||||
group_selector = self._build_strategy_selector(
|
||||
strategy=group.routing_strategy,
|
||||
routing_strategy_args=group.routing_strategy_args or {},
|
||||
built: Final = tuple(
|
||||
(
|
||||
group,
|
||||
self._build_strategy_selector(
|
||||
strategy=group.routing_strategy,
|
||||
routing_strategy_args=group.routing_strategy_args or {},
|
||||
register_callbacks=False,
|
||||
),
|
||||
)
|
||||
self._group_selectors[group.group_name] = (
|
||||
{strategy_value: group_selector} if group_selector is not None else {}
|
||||
for group in groups
|
||||
)
|
||||
self._replace_routing_groups(built)
|
||||
|
||||
def _replace_routing_groups(
|
||||
self,
|
||||
built: tuple[tuple[RoutingGroup, RouterStrategySelector | None], ...],
|
||||
) -> None:
|
||||
previous_selectors: Final[Mapping[str, Mapping[str, RouterStrategySelector]]] = getattr(
|
||||
self, "_group_selectors", {}
|
||||
)
|
||||
self._unregister_router_selectors(
|
||||
tuple(sel for selectors in previous_selectors.values() for sel in selectors.values())
|
||||
)
|
||||
for _, selector in built:
|
||||
if selector is not None:
|
||||
self._register_router_selector(selector)
|
||||
|
||||
self._routing_groups: dict[str, RoutingGroup] = {group.group_name: group for group, _ in built}
|
||||
self._model_to_group: dict[str, str] = {
|
||||
model_name: group.group_name for group, _ in built for model_name in group.models
|
||||
}
|
||||
self._group_selectors: dict[str, dict[str, RouterStrategySelector]] = {
|
||||
group.group_name: (
|
||||
{} if selector is None else {self._normalize_strategy(group.routing_strategy) or "": selector}
|
||||
)
|
||||
for group, selector in built
|
||||
}
|
||||
self._invalidate_model_group_info_cache()
|
||||
self._invalidate_access_groups_cache()
|
||||
|
||||
def get_routing_group(self, model_name: str) -> RoutingGroup | None:
|
||||
"""
|
||||
|
|
@ -12032,7 +12021,6 @@ class Router:
|
|||
_casted_value = int(kwargs[var])
|
||||
setattr(self, var, _casted_value)
|
||||
elif var == "routing_groups":
|
||||
self._routing_groups_input = kwargs[var]
|
||||
rebuild_routing_groups = True
|
||||
elif var == "optional_pre_call_checks":
|
||||
self.set_optional_pre_call_checks(kwargs[var])
|
||||
|
|
@ -12073,7 +12061,9 @@ class Router:
|
|||
self._apply_updated_routing_strategy_args()
|
||||
|
||||
if rebuild_routing_groups:
|
||||
self._init_routing_groups(self._routing_groups_input)
|
||||
routing_groups_input: Final = kwargs.get("routing_groups", self._routing_groups_input)
|
||||
self._init_routing_groups(routing_groups_input)
|
||||
self._routing_groups_input = routing_groups_input
|
||||
verbose_router_logger.debug("Updated Router settings: %s", self.get_settings())
|
||||
|
||||
def _get_client(self, deployment, kwargs, client_type=None):
|
||||
|
|
|
|||
95
litellm/router_utils/routing_groups.py
Normal file
95
litellm/router_utils/routing_groups.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
"""
|
||||
Validation for `router_settings.routing_groups`, shared by the Router and the
|
||||
proxy's config-update endpoint so a config the UI saves cannot be one the
|
||||
runtime refuses to load.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_router_logger
|
||||
from litellm.types.router import RoutingGroup, RoutingStrategy
|
||||
|
||||
|
||||
def validate_routing_strategy(routing_strategy: RoutingStrategy | str | None) -> None:
|
||||
"""
|
||||
Raises `ValueError` unless `routing_strategy` is a known strategy or None.
|
||||
|
||||
See: https://github.com/BerriAI/litellm/issues/11330
|
||||
"""
|
||||
if routing_strategy is None:
|
||||
return
|
||||
|
||||
valid_strategy_strings: Final = ("simple-shuffle", "lar1", *(s.value for s in RoutingStrategy))
|
||||
is_valid_string: Final = isinstance(routing_strategy, str) and routing_strategy in valid_strategy_strings
|
||||
is_valid_enum: Final = isinstance(routing_strategy, RoutingStrategy)
|
||||
if not is_valid_string and not is_valid_enum:
|
||||
raise ValueError(
|
||||
f"Invalid routing_strategy: '{routing_strategy}'. "
|
||||
f"Valid options: {list(valid_strategy_strings)}. "
|
||||
f"Check 'router_settings.routing_strategy' in your config.yaml "
|
||||
f"or the 'routing_strategy' parameter if using the Router SDK directly."
|
||||
)
|
||||
|
||||
|
||||
def parse_routing_groups(
|
||||
groups_input: Sequence[RoutingGroup | dict] | None,
|
||||
known_model_names: frozenset[str] = frozenset(),
|
||||
) -> tuple[RoutingGroup, ...]:
|
||||
"""
|
||||
Parses and validates `routing_groups`, raising `ValueError` on the first
|
||||
problem found. Every check runs before the caller mutates any state, so an
|
||||
invalid update can never leave a router holding a half-applied set of
|
||||
groups.
|
||||
"""
|
||||
if not groups_input:
|
||||
return ()
|
||||
|
||||
groups: Final = tuple(raw if isinstance(raw, RoutingGroup) else RoutingGroup(**raw) for raw in groups_input)
|
||||
|
||||
if any(not group.group_name for group in groups):
|
||||
raise ValueError("routing_groups: group_name must be non-empty.")
|
||||
|
||||
if any(group.group_name == "default" for group in groups):
|
||||
raise ValueError("routing_groups: 'default' is reserved for the implicit fallback group.")
|
||||
|
||||
names: Final = tuple(group.group_name for group in groups)
|
||||
duplicate_names: Final = frozenset(name for name in names if names.count(name) > 1)
|
||||
if duplicate_names:
|
||||
raise ValueError(f"routing_groups: group names must be unique, duplicate group_name '{min(duplicate_names)}'.")
|
||||
|
||||
for group in groups:
|
||||
validate_routing_strategy(group.routing_strategy)
|
||||
|
||||
owners_by_model: Final = tuple(
|
||||
(model_name, tuple(group.group_name for group in groups if model_name in group.models))
|
||||
for model_name in dict.fromkeys(model_name for group in groups for model_name in group.models)
|
||||
)
|
||||
conflicts: Final = tuple(
|
||||
f"model_name '{model_name}' appears in {' and '.join(repr(owner) for owner in owners)}"
|
||||
for model_name, owners in owners_by_model
|
||||
if len(owners) > 1
|
||||
)
|
||||
if conflicts:
|
||||
raise ValueError(f"routing_groups: {'; '.join(conflicts)}. Each model may belong to at most one group.")
|
||||
|
||||
unknown_models: Final = (
|
||||
tuple(
|
||||
(model_name, group.group_name)
|
||||
for group in groups
|
||||
for model_name in group.models
|
||||
if model_name not in known_model_names
|
||||
)
|
||||
if known_model_names
|
||||
else ()
|
||||
)
|
||||
for model_name, group_name in unknown_models:
|
||||
verbose_router_logger.warning(
|
||||
"routing_groups: model_name '%s' (group '%s') is not in model_list; "
|
||||
"the group entry will only take effect once a deployment with that "
|
||||
"model_name is added.",
|
||||
model_name,
|
||||
group_name,
|
||||
)
|
||||
|
||||
return groups
|
||||
|
|
@ -4959,6 +4959,71 @@ async def test_add_router_settings_from_db_config_merge_logic():
|
|||
assert combined_settings["nested_config"] == expected_nested
|
||||
|
||||
|
||||
def _routing_groups_router():
|
||||
from litellm import Router
|
||||
|
||||
return Router(
|
||||
model_list=[
|
||||
{"model_name": "m1", "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test"}},
|
||||
{"model_name": "m2", "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"}},
|
||||
],
|
||||
routing_groups=[{"group_name": "g1", "models": ["m1"], "routing_strategy": "latency-based-routing"}],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_db_routing_groups_do_not_abort_other_router_settings():
|
||||
"""Regression: an overlapping routing_groups value persisted in DB used to raise out of
|
||||
_add_router_settings_from_db_config, which skipped SSO / guardrail loading downstream."""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
router = _routing_groups_router()
|
||||
mock_db_config = MagicMock()
|
||||
mock_db_config.param_value = {
|
||||
"num_retries": 7,
|
||||
"routing_groups": [
|
||||
{"group_name": "g1", "models": ["m1"], "routing_strategy": "latency-based-routing"},
|
||||
{"group_name": "g2", "models": ["m1"], "routing_strategy": "least-busy"},
|
||||
],
|
||||
}
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config)
|
||||
|
||||
await ProxyConfig()._add_router_settings_from_db_config(
|
||||
config_data={}, llm_router=router, prisma_client=mock_prisma_client
|
||||
)
|
||||
|
||||
assert router.num_retries == 7
|
||||
assert router._model_to_group == {"m1": "g1"}
|
||||
assert router._get_routing_context("m1", None)[0] == "latency-based-routing"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_valid_db_routing_groups_still_replace_router_groups():
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
router = _routing_groups_router()
|
||||
mock_db_config = MagicMock()
|
||||
mock_db_config.param_value = {
|
||||
"num_retries": 7,
|
||||
"routing_groups": [{"group_name": "g2", "models": ["m2"], "routing_strategy": "least-busy"}],
|
||||
}
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config)
|
||||
|
||||
await ProxyConfig()._add_router_settings_from_db_config(
|
||||
config_data={}, llm_router=router, prisma_client=mock_prisma_client
|
||||
)
|
||||
|
||||
assert router.num_retries == 7
|
||||
assert router._model_to_group == {"m2": "g2"}
|
||||
assert router._get_routing_context("m2", None)[0] == "least-busy"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_router_settings_from_db_config_empty_db_lists_do_not_clobber_config_fallbacks():
|
||||
"""
|
||||
|
|
@ -9224,6 +9289,52 @@ def test_update_config_writes_only_sent_section(_update_config_setup):
|
|||
restore()
|
||||
|
||||
|
||||
def test_update_config_rejects_overlapping_routing_groups_before_writing(_update_config_setup):
|
||||
"""Regression: overlapping groups were persisted and only failed at router reload, where the
|
||||
failure took SSO and the other DB-backed settings down with it."""
|
||||
existing_groups = [{"group_name": "g1", "models": ["m1"], "routing_strategy": "least-busy"}]
|
||||
client, prisma, restore = _update_config_setup(
|
||||
initial_rows={"router_settings": {"num_retries": 2, "routing_groups": existing_groups}}
|
||||
)
|
||||
try:
|
||||
resp = client.post(
|
||||
"/config/update",
|
||||
json={
|
||||
"router_settings": {
|
||||
"routing_groups": [
|
||||
*existing_groups,
|
||||
{"group_name": "g2", "models": ["m1"], "routing_strategy": "latency-based-routing"},
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
assert "'m1' appears in 'g1' and 'g2'" in resp.text
|
||||
assert prisma.db.litellm_config.upsert_calls == []
|
||||
assert prisma.db.litellm_config.rows["router_settings"]["routing_groups"] == existing_groups
|
||||
finally:
|
||||
restore()
|
||||
|
||||
|
||||
def test_update_config_accepts_disjoint_routing_groups(_update_config_setup):
|
||||
client, prisma, restore = _update_config_setup(initial_rows={"router_settings": {"num_retries": 2}})
|
||||
groups = [
|
||||
{"group_name": "g1", "models": ["m1"], "routing_strategy": "least-busy"},
|
||||
{"group_name": "g2", "models": ["m2"], "routing_strategy": "latency-based-routing"},
|
||||
]
|
||||
try:
|
||||
resp = client.post("/config/update", json={"router_settings": {"routing_groups": groups}})
|
||||
assert resp.status_code == 200
|
||||
stored = prisma.db.litellm_config.rows["router_settings"]
|
||||
assert stored["num_retries"] == 2
|
||||
assert [(g["group_name"], g["models"]) for g in stored["routing_groups"]] == [
|
||||
("g1", ["m1"]),
|
||||
("g2", ["m2"]),
|
||||
]
|
||||
finally:
|
||||
restore()
|
||||
|
||||
|
||||
def test_update_config_env_var_round_trip_not_double_encrypted(_update_config_setup, monkeypatch):
|
||||
"""Endpoint-level regression for the /config/update double-encryption bug.
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from collections.abc import Callable
|
|||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
|
@ -806,6 +807,118 @@ def test_strategy_reinit_unregisters_override_selectors():
|
|||
assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger
|
||||
|
||||
|
||||
def _single_latency_group():
|
||||
return [{"group_name": "g1", "models": ["filtered-model"], "routing_strategy": "latency-based-routing"}]
|
||||
|
||||
|
||||
def _assert_still_routes_with_original_group(router, selector):
|
||||
assert list(router._routing_groups) == ["g1"]
|
||||
assert router._model_to_group == {"filtered-model": "g1"}
|
||||
assert router._group_selectors["g1"]["latency-based-routing"] is selector
|
||||
assert router._get_routing_context("filtered-model", None) == ("latency-based-routing", selector)
|
||||
assert sum(1 for cb in litellm.callbacks if cb is selector) == 1
|
||||
|
||||
|
||||
def test_failed_routing_groups_update_keeps_previous_groups(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
router = _build_router(routing_groups=_single_latency_group())
|
||||
selector = router._group_selectors["g1"]["latency-based-routing"]
|
||||
|
||||
with pytest.raises(ValueError, match="appears in"):
|
||||
router.update_settings(
|
||||
routing_groups=[
|
||||
*_single_latency_group(),
|
||||
{"group_name": "g2", "models": ["filtered-model"], "routing_strategy": "least-busy"},
|
||||
],
|
||||
)
|
||||
|
||||
_assert_still_routes_with_original_group(router, selector)
|
||||
assert sum(1 for cb in litellm.callbacks if type(cb) is not type(selector)) == 0
|
||||
assert litellm.input_callback == []
|
||||
|
||||
|
||||
def test_failed_routing_groups_update_does_not_poison_later_strategy_changes(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
router = _build_router(routing_groups=_single_latency_group())
|
||||
|
||||
with pytest.raises(ValueError, match="appears in"):
|
||||
router.update_settings(
|
||||
routing_groups=[
|
||||
*_single_latency_group(),
|
||||
{"group_name": "g2", "models": ["filtered-model"], "routing_strategy": "least-busy"},
|
||||
],
|
||||
)
|
||||
|
||||
router.update_settings(routing_strategy="least-busy")
|
||||
|
||||
assert list(router._routing_groups) == ["g1"]
|
||||
assert [g["group_name"] for g in router.get_settings()["routing_groups"]] == ["g1"]
|
||||
|
||||
|
||||
def test_overlap_error_names_every_conflicting_model():
|
||||
with pytest.raises(ValueError, match="appears in") as exc_info:
|
||||
_build_router(
|
||||
routing_groups=[
|
||||
{
|
||||
"group_name": "g1",
|
||||
"models": ["filtered-model", "other-model"],
|
||||
"routing_strategy": "latency-based-routing",
|
||||
},
|
||||
{
|
||||
"group_name": "g2",
|
||||
"models": ["filtered-model", "other-model"],
|
||||
"routing_strategy": "least-busy",
|
||||
},
|
||||
],
|
||||
)
|
||||
message = str(exc_info.value)
|
||||
assert "'filtered-model' appears in 'g1' and 'g2'" in message
|
||||
assert "'other-model' appears in 'g1' and 'g2'" in message
|
||||
|
||||
|
||||
def test_invalid_group_strategy_keeps_previous_groups(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
router = _build_router(routing_groups=_single_latency_group())
|
||||
selector = router._group_selectors["g1"]["latency-based-routing"]
|
||||
|
||||
with pytest.raises(ValueError, match="Invalid routing_strategy"):
|
||||
router.update_settings(
|
||||
routing_groups=[
|
||||
{"group_name": "g2", "models": ["other-model"], "routing_strategy": "not-a-real-strategy"},
|
||||
],
|
||||
)
|
||||
|
||||
_assert_still_routes_with_original_group(router, selector)
|
||||
|
||||
|
||||
def test_unbuildable_group_selector_keeps_previous_groups(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
router = _build_router(routing_groups=_single_latency_group())
|
||||
selector = router._group_selectors["g1"]["latency-based-routing"]
|
||||
|
||||
with pytest.raises(ValidationError, match="ttl"):
|
||||
router.update_settings(
|
||||
routing_groups=[
|
||||
{"group_name": "g0", "models": ["other-model"], "routing_strategy": "least-busy"},
|
||||
*_single_latency_group(),
|
||||
{
|
||||
"group_name": "g2",
|
||||
"models": ["other-model-2"],
|
||||
"routing_strategy": "latency-based-routing",
|
||||
"routing_strategy_args": {"ttl": "not-a-number"},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
_assert_still_routes_with_original_group(router, selector)
|
||||
assert litellm.callbacks == [selector]
|
||||
assert litellm.input_callback == []
|
||||
|
||||
|
||||
def test_override_selectors_are_not_registered_process_wide(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
monkeypatch.setattr(litellm, "input_callback", [])
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ const renderModal = (overrides: Partial<React.ComponentProps<typeof RoutingGroup
|
|||
strategyDescriptions={STRATEGY_DESCRIPTIONS}
|
||||
modelOptions={MODEL_OPTIONS}
|
||||
existingGroupNames={["already-taken", "other-group"]}
|
||||
groupNameByModel={{}}
|
||||
onClose={onClose}
|
||||
onSubmit={onSubmit}
|
||||
{...overrides}
|
||||
|
|
@ -307,6 +308,19 @@ describe("RoutingGroupModal", () => {
|
|||
expect(onSubmit.mock.calls[0][0]).toStrictEqual(expected);
|
||||
});
|
||||
|
||||
it("blocks a model another group already claims", async () => {
|
||||
const user = userEvent.setup();
|
||||
const { onSubmit } = renderModal({ groupNameByModel: { "gpt-4o": "cheap" } });
|
||||
|
||||
await typeName(user, "security");
|
||||
await pickModels(user, "gpt-4o");
|
||||
await pickStrategy(user, "latency-based-routing");
|
||||
await save(user, "Create Group");
|
||||
|
||||
expect(await screen.findByText(/Already claimed: gpt-4o/)).toBeInTheDocument();
|
||||
expect(onSubmit).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("describes the selected strategy", async () => {
|
||||
renderModal();
|
||||
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ import {
|
|||
toRoutingGroupFormValues,
|
||||
} from "./routingGroupPayload";
|
||||
import type { RoutingGroup } from "./types";
|
||||
import { modelConflictError } from "./modelOwnership";
|
||||
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
import { Button } from "@/components/ui/button";
|
||||
|
||||
|
|
@ -40,6 +41,7 @@ interface RoutingGroupModalProps {
|
|||
strategyDescriptions: Record<string, string>;
|
||||
modelOptions: string[];
|
||||
existingGroupNames: string[];
|
||||
groupNameByModel: Record<string, string>;
|
||||
onClose: () => void;
|
||||
onSubmit: (group: RoutingGroup) => Promise<void> | void;
|
||||
saving?: boolean;
|
||||
|
|
@ -57,6 +59,7 @@ const RoutingGroupModal: React.FC<RoutingGroupModalProps> = ({
|
|||
strategyDescriptions,
|
||||
modelOptions,
|
||||
existingGroupNames,
|
||||
groupNameByModel,
|
||||
onClose,
|
||||
onSubmit,
|
||||
saving,
|
||||
|
|
@ -77,12 +80,20 @@ const RoutingGroupModal: React.FC<RoutingGroupModalProps> = ({
|
|||
.min(1, "Group name is required")
|
||||
.max(GROUP_NAME_MAX_LENGTH, `Must be ${GROUP_NAME_MAX_LENGTH} characters or fewer`)
|
||||
.refine((value) => !reservedNames.has(value.toLowerCase()), "A group with this name already exists"),
|
||||
models: z.array(z.string()).min(1, "Select at least one model"),
|
||||
models: z
|
||||
.array(z.string())
|
||||
.min(1, "Select at least one model")
|
||||
.superRefine((models, ctx) => {
|
||||
const conflict = modelConflictError(models, groupNameByModel);
|
||||
if (conflict !== null) {
|
||||
ctx.addIssue({ code: "custom", message: conflict });
|
||||
}
|
||||
}),
|
||||
routing_strategy: z.string().min(1, "Strategy is required"),
|
||||
routing_strategy_args: z.string(),
|
||||
};
|
||||
return z.object(shape);
|
||||
}, [reservedNames]);
|
||||
}, [reservedNames, groupNameByModel]);
|
||||
|
||||
const form = useZodForm(schema, { defaultValues: toRoutingGroupFormValues(initialValue, availableStrategies) });
|
||||
|
||||
|
|
@ -124,7 +135,7 @@ const RoutingGroupModal: React.FC<RoutingGroupModalProps> = ({
|
|||
control={form.control}
|
||||
name="models"
|
||||
label="Models"
|
||||
description="Models from your model list that this group routes between."
|
||||
description="Models from your model list that this group routes between. A model can only be in one group."
|
||||
>
|
||||
{({ id, value, onChange, "aria-invalid": ariaInvalid, "aria-describedby": ariaDescribedBy }) => (
|
||||
<Combobox multiple items={modelOptions} value={value} onValueChange={onChange}>
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ import RoutingGroupsTable from "./RoutingGroupsTable";
|
|||
import RoutingGroupModal from "./RoutingGroupModal";
|
||||
import { toast } from "@/lib/toast";
|
||||
import type { RoutingGroup } from "./types";
|
||||
import { groupNameByModel } from "./modelOwnership";
|
||||
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
|
||||
|
||||
const RoutingGroups: React.FC = () => {
|
||||
|
|
@ -30,7 +31,7 @@ const RoutingGroups: React.FC = () => {
|
|||
const [editingGroup, setEditingGroup] = useState<RoutingGroup | null>(null);
|
||||
const [deletingGroup, setDeletingGroup] = useState<RoutingGroup | null>(null);
|
||||
|
||||
const groups = data?.routingGroups ?? [];
|
||||
const groups = useMemo(() => data?.routingGroups ?? [], [data?.routingGroups]);
|
||||
|
||||
const filteredGroups = useMemo(() => {
|
||||
const q = searchQuery.trim().toLowerCase();
|
||||
|
|
@ -51,6 +52,11 @@ const RoutingGroups: React.FC = () => {
|
|||
|
||||
const strategyDescriptions = routerFields?.routing_strategy_descriptions ?? {};
|
||||
|
||||
const ownerByModel = useMemo(
|
||||
() => groupNameByModel(groups, drawerMode === "edit" ? editingGroup?.group_name : undefined),
|
||||
[groups, drawerMode, editingGroup],
|
||||
);
|
||||
|
||||
const modelOptions = useMemo<string[]>(() => {
|
||||
const records = (modelHub?.data ?? []) as Array<{ model_group?: string }>;
|
||||
const names = records.map((r) => r.model_group).filter((n): n is string => Boolean(n));
|
||||
|
|
@ -160,6 +166,7 @@ const RoutingGroups: React.FC = () => {
|
|||
strategyDescriptions={strategyDescriptions}
|
||||
modelOptions={modelOptions}
|
||||
existingGroupNames={groups.map((g) => g.group_name)}
|
||||
groupNameByModel={ownerByModel}
|
||||
onClose={() => setDrawerOpen(false)}
|
||||
onSubmit={handleSubmit}
|
||||
saving={saveMutation.isPending}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,33 @@
|
|||
import { describe, expect, it } from "vitest";
|
||||
|
||||
import { groupNameByModel, modelConflictError } from "./modelOwnership";
|
||||
import type { RoutingGroup } from "./types";
|
||||
|
||||
const groups: RoutingGroup[] = [
|
||||
{ group_name: "cheap", models: ["m1", "m2"], routing_strategy: "latency-based-routing" },
|
||||
{ group_name: "security", models: ["m3"], routing_strategy: "least-busy" },
|
||||
];
|
||||
|
||||
describe("groupNameByModel", () => {
|
||||
it("maps every claimed model to its owning group", () => {
|
||||
expect(groupNameByModel(groups)).toEqual({ m1: "cheap", m2: "cheap", m3: "security" });
|
||||
});
|
||||
|
||||
it("excludes the group being edited so its own models stay selectable", () => {
|
||||
expect(groupNameByModel(groups, "cheap")).toEqual({ m3: "security" });
|
||||
});
|
||||
});
|
||||
|
||||
describe("modelConflictError", () => {
|
||||
it("passes models that no other group claims", () => {
|
||||
expect(modelConflictError(["m4"], groupNameByModel(groups, "cheap"))).toBeNull();
|
||||
expect(modelConflictError(undefined, groupNameByModel(groups))).toBeNull();
|
||||
});
|
||||
|
||||
it("names every model already claimed by another group", () => {
|
||||
const error = modelConflictError(["m1", "m3", "m4"], groupNameByModel(groups));
|
||||
expect(error).toBe(
|
||||
'Each model may belong to at most one group. Already claimed: m1 (in "cheap"), m3 (in "security")',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,18 @@
|
|||
import type { RoutingGroup } from "./types";
|
||||
|
||||
export const groupNameByModel = (groups: RoutingGroup[], excludeGroupName?: string): Record<string, string> =>
|
||||
Object.fromEntries(
|
||||
groups
|
||||
.filter((group) => group.group_name !== excludeGroupName)
|
||||
.flatMap((group) => group.models.map((model) => [model, group.group_name] as const)),
|
||||
);
|
||||
|
||||
export const modelConflictError = (
|
||||
models: string[] | undefined,
|
||||
ownerByModel: Record<string, string>,
|
||||
): string | null => {
|
||||
const conflicts = (models ?? []).filter((model) => ownerByModel[model] !== undefined);
|
||||
if (conflicts.length === 0) return null;
|
||||
const detail = conflicts.map((model) => `${model} (in "${ownerByModel[model]}")`).join(", ");
|
||||
return `Each model may belong to at most one group. Already claimed: ${detail}`;
|
||||
};
|
||||
Loading…
Add table
Reference in a new issue