scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #34564 from BerriAI/litellm_fix_router_registry_leak_on_edit

Browse source 
fix(router): release the pre-routing strategy slot when a deployment is replaced or deleted
This commit is contained in:
tin-berri 2026-07-27 16:46:40 -07:00 • committed by GitHub
commit 8f86c87f8e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 577 additions and 55 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -1040,22 +1040,6 @@ class ModelManagementAuthChecks:
return True
def _deployment_name_and_model(deployment: Optional[Union[Deployment, Dict[str, object]]]) -> Tuple[Optional[str], str]:
"""Return (model_name, litellm_params.model) for a deployment.
delete_deployment is annotated to return a Deployment but hands back the raw
model_list dict at runtime, so both shapes are handled; the model defaults to "".
"""
if deployment is None:
return None, ""
if isinstance(deployment, dict):
name = deployment.get("model_name")
params = deployment.get("litellm_params")
model = params.get("model") if isinstance(params, dict) else None
return (name if isinstance(name, str) else None), (model if isinstance(model, str) else "")
return deployment.model_name, str(getattr(deployment.litellm_params, "model", "") or "")
#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964
@router.post(
"/model/delete",
@ -1127,19 +1111,7 @@ async def delete_model(
## DELETE FROM ROUTER ##
if llm_router is not None:
deleted_deployment = llm_router.delete_deployment(id=model_info.id)
# delete_deployment only drops the deployment from model_list; the auto/
# complexity router registries are keyed by model_name and would otherwise
# retain a stale (now unbacked) entry, so evict it here too. Guard on the
# auto_router/ prefix (as clear_cache does): a regular DB model that merely
# shares a model_name with a config-defined router must not evict that router,
# since add_deployment never restores config-defined routers.
deleted_name, deleted_model = _deployment_name_and_model(deleted_deployment)
if deleted_name is not None and deleted_model.startswith("auto_router/"):
llm_router.auto_routers.pop(deleted_name, None)
llm_router.complexity_routers.pop(deleted_name, None)
llm_router.adaptive_routers.pop(deleted_name, None)
llm_router.quality_routers.pop(deleted_name, None)
llm_router.delete_deployment(id=model_info.id)
# Runs after the row delete so the sibling check sees post-delete state.
if model_params.model_info.team_id is not None:

109
litellm/router.py
View file

@ -7702,6 +7702,21 @@ class Router:
"""True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
return litellm_params.model.startswith("auto_router/adaptive_router")
def _deployment_participates_in_adaptive_routing(self, litellm_params: LiteLLM_Params) -> bool:
"""True when this deployment owns an `adaptive_routers` entry once finalized:
a dedicated adaptive router, or a complexity router whose config enables the
adaptive companion. Mirrors the two arms of
`_finalize_adaptive_router_if_configured`, which is the registry's only writer."""
if self._is_adaptive_router_deployment(litellm_params=litellm_params):
return True
if not self._is_complexity_router_deployment(litellm_params=litellm_params):
return False
config = litellm_params.complexity_router_config
if not config:
return False
adaptive_flag: object = config.get("adaptive")
return bool(adaptive_flag)
@staticmethod
def _has_registered_strategy(
registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]],
@ -7734,20 +7749,56 @@ class Router:
TaggedPreRoutingStrategy(tags=tags, strategy=strategy),
]
@staticmethod
def _unregister_pre_routing_strategy(
registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]],
model_name: str,
tags: tuple[str, ...],
) -> bool:
"""Drop the strategy registered for this exact (model_name, tags) pair, leaving
strategies registered under the same name with different tags in place. Returns
whether anything was actually dropped."""
existing = registry.get(model_name, [])
remaining = [entry for entry in existing if entry.tags != tags]
if len(remaining) == len(existing):
return False
if remaining:
registry[model_name] = remaining
else:
registry.pop(model_name, None)
return True
def _unregister_pre_routing_strategy_for_deployment(self, deployment: Deployment) -> None:
"""
Release the pre-routing strategy a deployment holds, so removing it from the
model_list also frees its (model_name, tags) slot.
Without this, re-adding the deployment (an edit arriving via upsert_deployment,
or a router recreated under a name that was deleted earlier) hits the
"already exists" guard in `_register_pre_routing_strategy`, which
`ignore_invalid_deployments` swallows - the deployment then silently never
makes it back into the model_list.
Released from every registry rather than the first match, because registration is
one-to-many: a complexity router configured with `adaptive` is also registered in
`adaptive_routers` under the same (model_name, tags) by the deferred finalize pass.
Guarded on the auto_router/ prefix so removing a *regular* deployment can't evict a
router that merely shares its model_name.
"""
if not deployment.litellm_params.model.startswith("auto_router/"):
return
model_name = deployment.model_name
tags = self._deployment_tags(deployment)
for registry in (self.auto_routers, self.complexity_routers, self.quality_routers):
self._unregister_pre_routing_strategy(registry, model_name, tags)
if self._unregister_pre_routing_strategy(self.adaptive_routers, model_name, tags):
self._sync_adaptive_router_hooks()
def _finalize_adaptive_router_if_configured(self) -> None:
"""Locate every adaptive-router deployment in the finalized model_list and
build an AdaptiveRouter for each. Safe no-op when none are configured.
Idempotent: skips any deployment whose (model_name, tags) pair is already
initialized, so hot-reloads don't rebuild routers that would lose state."""
# Drop any adaptive-router hooks left over from a previous Router
# instance (e.g. after `/config/reload` replaced `llm_router`). Without
# this, stale AdaptiveRouterPostCallHook callbacks from the old Router
# remain wired up in `litellm.callbacks` and double-fire signal
# recording for every request.
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
for entry in self.model_list or []:
lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params
lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None
@ -7779,6 +7830,16 @@ class Router:
TaggedPreRoutingStrategy(tags=tagged.tags, strategy=adaptive_router),
]
self._sync_adaptive_router_hooks()
def _sync_adaptive_router_hooks(self) -> None:
"""Rebuild the AdaptiveRouterPostCallHook set so it is exactly one hook per
currently registered adaptive router. Run at every point the adaptive registry
changes, otherwise a released router keeps recording turns through its hook."""
from litellm.router_strategy.adaptive_router.hooks import (
AdaptiveRouterPostCallHook,
)
for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook):
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
for tagged_adaptive_routers in self.adaptive_routers.values():
@ -8401,13 +8462,29 @@ class Router:
self._invalidate_access_groups_cache()
self._update_deployment_indices_after_removal(model_id=deployment_id, removal_idx=removal_idx)
# Free the outgoing deployment's pre-routing strategy slot (keyed by the
# OLD model_name/tags) before the re-add below re-registers it.
self._unregister_pre_routing_strategy_for_deployment(deployment=_deployment_on_router)
# if the model_id is not in router
self.add_deployment(deployment=deployment)
# add_deployment() builds every strategy EXCEPT the adaptive one, which
# set_model_list() defers until the whole model_list is visible. Re-run that
# deferred pass so an adaptive router whose slot was just released above is
# rebuilt rather than left unregistered.
if self._deployment_participates_in_adaptive_routing(litellm_params=deployment.litellm_params) or (
_deployment_on_router is not None
and self._deployment_participates_in_adaptive_routing(
litellm_params=_deployment_on_router.litellm_params
)
):
self._finalize_adaptive_router_if_configured()
return deployment
except Exception as e:
if self.ignore_invalid_deployments:
verbose_router_logger.debug(
f"Error upserting deployment: {e}, ignoring and continuing with other deployments."
verbose_router_logger.warning(
f"Error upserting deployment {deployment.model_name} (id={deployment.model_info.id}): {e}. "
"Dropping it and continuing with other deployments."
)
return None
else:
@ -8436,6 +8513,16 @@ class Router:
_budget_limiter = self._get_router_deployment_budget_limiter()
if _budget_limiter is not None:
_budget_limiter.unregister_deployment_budget(model_id=id)
try:
self._unregister_pre_routing_strategy_for_deployment(
deployment=item if isinstance(item, Deployment) else Deployment(**item)
)
except Exception:
verbose_router_logger.exception(
"delete_deployment: could not release pre-routing strategies for model_id=%s; "
"the deployment is out of the model_list and its indices are repaired",
id,
)
return item
else:
return None

50
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -636,14 +636,33 @@ class TestDeleteModelClearsRouterRegistry:
not just from model_list, or a stale (now unbacked) router entry lingers until restart.
"""
@staticmethod
def _complexity_router_deployment(model_id: str, tags: list | None = None) -> dict:
return {
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}},
"complexity_router_default_model": "gpt-4o",
**({"tags": tags} if tags else {}),
},
"model_info": {"id": model_id, "db_model": True},
}
@pytest.mark.asyncio
async def test_delete_model_pops_router_registries(self):
async def test_delete_model_releases_only_the_deleted_routers_slot(self):
"""Deleting one tagged router must release its own slot and leave a sibling
sharing the model_name registered. A blanket pop(model_name) here would take
both down, and nothing reloads on the delete path to restore the survivor.
"""
import litellm
from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete
from litellm.proxy.management_endpoints.model_management_endpoints import (
delete_model as delete_model_endpoint,
)
from litellm.proxy.management_endpoints.model_management_endpoints import ModelInfoDelete
model_id = "router-del-1"
surviving_id = "router-del-2"
admin_user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
db_row = LiteLLM_ProxyModelTable(
model_id=model_id,
@ -660,16 +679,16 @@ class TestDeleteModelClearsRouterRegistry:
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row)
mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row)
mock_router = MagicMock()
mock_router.delete_deployment = MagicMock(
return_value={
"model_name": "smart-router",
"litellm_params": {"model": "auto_router/complexity_router"},
"model_info": {"id": model_id},
}
real_router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
self._complexity_router_deployment(model_id, tags=["team-a"]),
self._complexity_router_deployment(surviving_id, tags=["team-b"]),
],
ignore_invalid_deployments=True,
)
mock_router.auto_routers = {"smart-router": MagicMock()}
mock_router.complexity_routers = {"smart-router": MagicMock()}
assert len(real_router.complexity_routers["smart-router"]) == 2
_PS = "litellm.proxy.proxy_server"
with (
@ -679,16 +698,17 @@ class TestDeleteModelClearsRouterRegistry:
patch(f"{_PS}.proxy_logging_obj", MagicMock()),
patch(f"{_PS}.general_settings", {}),
patch(f"{_PS}.premium_user", True),
patch(f"{_PS}.llm_router", mock_router),
patch(f"{_PS}.llm_router", real_router),
):
await delete_model_endpoint(
model_info=ModelInfoDelete(id=model_id),
user_api_key_dict=admin_user,
)
mock_router.delete_deployment.assert_called_once_with(id=model_id)
assert "smart-router" not in mock_router.auto_routers
assert "smart-router" not in mock_router.complexity_routers
assert model_id not in [m["model_info"]["id"] for m in real_router.model_list]
surviving = real_router.complexity_routers["smart-router"]
assert len(surviving) == 1
assert surviving[0].tags == ("team-b",)
@pytest.mark.asyncio
async def test_delete_regular_model_preserves_config_router_sharing_name(self):

443
tests/test_litellm/test_router.py
View file

@ -5936,3 +5936,446 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags():
bedrock_tags=request_tags,
)
assert mock_sign.call_args.kwargs["data"]["tags"] == request_tags
class TestPreRoutingStrategyRegistryLifecycle:
"""
Regression tests: a deployment leaving the model_list must release the
pre-routing strategy slot it holds in `auto_routers` / `complexity_routers` /
`adaptive_routers` / `quality_routers`.
Before this fix, editing an auto-router-family model (a UI save, which reaches
every other pod as an `upsert_deployment` from the periodic DB reload) popped
the deployment out of the model_list and then failed to re-add it: registration
raised "already exists" against the stale registry entry, and
`ignore_invalid_deployments=True` swallowed the error. The router vanished from
the Models page and stayed gone until a proxy restart, while the DB row and the
"saved successfully" response both looked fine.
"""
@staticmethod
def _complexity_router_params(default_model: str, tags=None) -> dict:
return {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}
},
"complexity_router_default_model": default_model,
**({"tags": tags} if tags else {}),
}
@classmethod
def _router_with_complexity_router(cls, default_model: str = "gpt-4o") -> "litellm.Router":
return litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
{
"model_name": "smart-router",
"litellm_params": cls._complexity_router_params(default_model),
"model_info": {"id": "router-1", "db_model": True},
},
],
ignore_invalid_deployments=True,
)
@staticmethod
def _model_names(router: "litellm.Router") -> list:
return [model["model_name"] for model in router.model_list]
def test_upsert_of_edited_router_keeps_it_routable(self):
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = self._router_with_complexity_router()
router.upsert_deployment(
deployment=Deployment(
model_name="smart-router",
litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")),
model_info=ModelInfo(id="router-1", db_model=True),
)
)
assert "smart-router" in self._model_names(router)
registered = router.complexity_routers["smart-router"]
assert len(registered) == 1
# the surviving strategy is the edited one, not the pre-edit leftover
assert registered[0].strategy.config.default_model == "gpt-4o-mini"
def test_unchanged_upsert_leaves_router_untouched(self):
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = self._router_with_complexity_router()
strategy_before = router.complexity_routers["smart-router"][0].strategy
for _ in range(3):
router.upsert_deployment(
deployment=Deployment(
model_name="smart-router",
litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o")),
model_info=ModelInfo(id="router-1", db_model=True),
)
)
assert "smart-router" in self._model_names(router)
assert router.complexity_routers["smart-router"][0].strategy is strategy_before
def test_delete_frees_the_name_for_a_new_router(self):
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = self._router_with_complexity_router()
router.delete_deployment(id="router-1")
assert "smart-router" not in router.complexity_routers
router.add_deployment(
deployment=Deployment(
model_name="smart-router",
litellm_params=LiteLLM_Params(**self._complexity_router_params("gpt-4o-mini")),
model_info=ModelInfo(id="router-2", db_model=True),
)
)
assert "smart-router" in self._model_names(router)
assert router.complexity_routers["smart-router"][0].strategy.config.default_model == "gpt-4o-mini"
def test_delete_only_frees_the_matching_tag_slot(self):
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
{
"model_name": "shared-router",
"litellm_params": self._complexity_router_params("gpt-4o", tags=["team-a"]),
"model_info": {"id": "router-a"},
},
{
"model_name": "shared-router",
"litellm_params": self._complexity_router_params("gpt-4o-mini", tags=["team-b"]),
"model_info": {"id": "router-b"},
},
],
ignore_invalid_deployments=True,
)
assert len(router.complexity_routers["shared-router"]) == 2
router.delete_deployment(id="router-a")
remaining = router.complexity_routers["shared-router"]
assert len(remaining) == 1
assert remaining[0].tags == ("team-b",)
def test_delete_of_regular_model_preserves_router_sharing_its_name(self):
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "gpt-4o-mini"}},
{
"model_name": "shared-name",
"litellm_params": self._complexity_router_params("gpt-4o"),
"model_info": {"id": "router-1"},
},
{
"model_name": "shared-name",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": "regular-1"},
},
],
ignore_invalid_deployments=True,
)
strategy = router.complexity_routers["shared-name"][0].strategy
router.delete_deployment(id="regular-1")
assert router.complexity_routers["shared-name"][0].strategy is strategy
def test_upsert_of_edited_adaptive_router_rebuilds_it(self):
"""Adaptive routers are built by set_model_list()'s deferred pass, not by
add_deployment(), so releasing the slot on edit must be paired with a rebuild -
otherwise the edit silently turns adaptive routing off."""
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
def adaptive_params(available_models: list) -> dict:
return {
"model": "auto_router/adaptive_router",
"adaptive_router_config": {"available_models": available_models},
}
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
{
"model_name": "adaptive-router",
"litellm_params": adaptive_params(["gpt-4o-mini"]),
"model_info": {"id": "router-1", "db_model": True},
},
],
ignore_invalid_deployments=True,
)
assert "adaptive-router" in router.adaptive_routers
router.upsert_deployment(
deployment=Deployment(
model_name="adaptive-router",
litellm_params=LiteLLM_Params(**adaptive_params(["gpt-4o", "gpt-4o-mini"])),
model_info=ModelInfo(id="router-1", db_model=True),
)
)
assert "adaptive-router" in self._model_names(router)
registered = router.adaptive_routers["adaptive-router"]
assert len(registered) == 1
assert set(registered[0].strategy.config.available_models) == {"gpt-4o", "gpt-4o-mini"}
def test_delete_repairs_indices_even_when_strategy_release_fails(self):
"""Structural removal and strategy release are not equally critical. Once the entry
leaves model_list the index maps must be repaired no matter what, so releasing the
registry slot runs after that repair and cannot abandon the router half-updated."""
router = self._router_with_complexity_router()
idx = router.model_id_to_deployment_index_map["router-1"]
router.model_list[idx] = {"model_name": "smart-router", "litellm_params": None}
returned = router.delete_deployment(id="router-1")
assert returned is not None
assert "router-1" not in router.model_id_to_deployment_index_map
assert all(entry.get("model_info", {}).get("id") != "router-1" for entry in router.model_list)
assert router.get_deployment(model_id="router-1") is None
assert "gpt-4o" in self._model_names(router)
def test_delete_of_adaptive_enabled_complexity_router_frees_both_registries(self):
"""A complexity router with adaptive set is registered in BOTH complexity_routers
and adaptive_routers under the same (model_name, tags). Releasing only the first
match leaves the adaptive strategy live, so a deleted alias stays routable and its
post-call hook keeps recording."""
import litellm as litellm_module
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
params = {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"},
"adaptive": True,
},
"complexity_router_default_model": "gpt-4o",
}
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
{
"model_name": "hybrid-router",
"litellm_params": params,
"model_info": {"id": "router-1", "db_model": True},
},
],
ignore_invalid_deployments=True,
)
assert "hybrid-router" in router.complexity_routers
assert "hybrid-router" in router.adaptive_routers
hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
assert len(hooks) == 1
router.delete_deployment(id="router-1")
assert "hybrid-router" not in router.complexity_routers
assert "hybrid-router" not in router.adaptive_routers
remaining_hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(
AdaptiveRouterPostCallHook
)
assert remaining_hooks == []
def test_upsert_of_edited_quality_router_keeps_it_routable(self):
"""_unregister_pre_routing_strategy_for_deployment dispatches on four prefixes;
quality_router is one of them and would otherwise go unexercised."""
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
def quality_params(default_model: str) -> dict:
return {
"model": "auto_router/quality_router",
"quality_router_default_model": default_model,
}
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
{
"model_name": "quality-router",
"litellm_params": quality_params("gpt-4o"),
"model_info": {"id": "router-1", "db_model": True},
},
],
ignore_invalid_deployments=True,
)
assert "quality-router" in router.quality_routers
router.upsert_deployment(
deployment=Deployment(
model_name="quality-router",
litellm_params=LiteLLM_Params(**quality_params("gpt-4o-mini")),
model_info=ModelInfo(id="router-1", db_model=True),
)
)
assert "quality-router" in self._model_names(router)
registered = router.quality_routers["quality-router"]
assert len(registered) == 1
assert registered[0].strategy.config.default_model == "gpt-4o-mini"
@staticmethod
def _hybrid_router_params(tiers: dict) -> dict:
return {
"model": "auto_router/complexity_router",
"complexity_router_config": {"tiers": tiers, "adaptive": True},
"complexity_router_default_model": "gpt-4o",
}
@classmethod
def _router_with_hybrid_router(cls) -> "litellm.Router":
return litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "openai/gpt-4o"}},
{"model_name": "gpt-4o-mini", "litellm_params": {"model": "openai/gpt-4o-mini"}},
{
"model_name": "hybrid-router",
"litellm_params": cls._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o"}),
"model_info": {"id": "router-1", "db_model": True},
},
],
ignore_invalid_deployments=True,
)
def test_upsert_of_edited_hybrid_complexity_router_relinks_adaptive(self):
"""Editing an adaptive-enabled complexity router releases its adaptive companion
along with the complexity slot; the finalize re-run must fire for it (not just for
`auto_router/adaptive_router` deployments) or the rebuilt complexity router keeps
routing while bandit recording, DB persistence and /adaptive_router/state all
silently stop until the next full reload."""
import litellm as litellm_module
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = self._router_with_hybrid_router()
assert "hybrid-router" in router.adaptive_routers
router.upsert_deployment(
deployment=Deployment(
model_name="hybrid-router",
litellm_params=LiteLLM_Params(
**self._hybrid_router_params(
{"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}
)
),
model_info=ModelInfo(id="router-1", db_model=True),
)
)
assert "hybrid-router" in self._model_names(router)
assert "hybrid-router" in router.complexity_routers
assert "hybrid-router" in router.adaptive_routers
rebuilt = router.complexity_routers["hybrid-router"][0].strategy
assert router.adaptive_routers["hybrid-router"][0].strategy is rebuilt.adaptive_router
hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
assert len(hooks) == 1
def test_upsert_turning_adaptive_on_builds_the_companion(self):
"""An edit that flips `adaptive: true` on an existing complexity router must
register the companion immediately; neither side of the old prefix-only gate
matches a complexity deployment, so the flip was a silent no-op until restart."""
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = self._router_with_complexity_router()
assert "smart-router" not in router.adaptive_routers
params = self._complexity_router_params("gpt-4o")
params["complexity_router_config"] = {**params["complexity_router_config"], "adaptive": True}
router.upsert_deployment(
deployment=Deployment(
model_name="smart-router",
litellm_params=LiteLLM_Params(**params),
model_info=ModelInfo(id="router-1", db_model=True),
)
)
assert "smart-router" in router.adaptive_routers
def test_unregister_pre_routing_strategy_scopes_the_drop_by_tags(self):
"""The bool return drives the hook re-sync; a tag mismatch must report False and
leave the registry untouched, and dropping the last entry must free the key."""
from litellm.types.router import TaggedPreRoutingStrategy
registry = {
"m": [
TaggedPreRoutingStrategy(tags=("team-a",), strategy=object()),
TaggedPreRoutingStrategy(tags=(), strategy=object()),
]
}
assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-b",)) is False
assert len(registry["m"]) == 2
assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-a",)) is True
assert [entry.tags for entry in registry["m"]] == [()]
assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ()) is True
assert "m" not in registry
def test_unregister_for_deployment_ignores_non_router_deployments(self):
"""Direct twin of the endpoint-level test: a regular deployment that shares a
router's model_name must not evict the router's registry slot."""
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
router = self._router_with_complexity_router()
router._unregister_pre_routing_strategy_for_deployment(
deployment=Deployment(
model_name="smart-router",
litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
model_info=ModelInfo(id="plain-1", db_model=True),
)
)
assert "smart-router" in router.complexity_routers
def test_sync_adaptive_router_hooks_keeps_one_hook_per_registered_router(self):
"""Re-syncing must replace, not accumulate: a duplicated hook double-fires
bandit signal recording for every request."""
import litellm as litellm_module
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
router = self._router_with_hybrid_router()
router._sync_adaptive_router_hooks()
router._sync_adaptive_router_hooks()
hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook)
assert len(hooks) == 1
def test_deployment_participates_in_adaptive_routing_matrix(self):
"""The upsert finalize re-run keys off this predicate for both the incoming and
outgoing deployment; a false negative silently strands the adaptive companion."""
from litellm.types.router import LiteLLM_Params
router = self._router_with_complexity_router()
cases = [
({"model": "auto_router/adaptive_router", "adaptive_router_config": {}}, True),
(self._hybrid_router_params({"SIMPLE": "gpt-4o-mini"}), True),
(self._complexity_router_params("gpt-4o"), False),
(
{
"model": "auto_router/complexity_router",
"complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "adaptive": False},
"complexity_router_default_model": "gpt-4o",
},
False,
),
({"model": "openai/gpt-4o"}, False),
]
for params, expected in cases:
actual = router._deployment_participates_in_adaptive_routing(
litellm_params=LiteLLM_Params(**params)
)
assert actual is expected, params["model"]