From eb45a088d3d3c255118d0effb41e4613af0907c0 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Wed, 9 Sep 2026 15:53:05 -0700 Subject: [PATCH] fix(router): resolve team-scoped auto-routers by their public name (#40432) A team-scoped auto-router is stored under an internal model_name_{team_id}_{uuid} with the caller-facing name in model_info.team_public_model_name, and the four pre-routing strategy registries key on that internal name. A team key asks for the public name, so the strategy lookup missed, the team early-resolve exit handed back the marker deployment itself, and every call 400'd with "Unmapped LLM provider". The strategy lookup now resolves the requested name through the same team-first, then global, then admin-across-teams deployment resolution the deployment path uses, and looks the registries up under the model_name of whatever that resolves to. Both exits of _common_checks_available_deployment drop strategy markers through one helper, so a marker-only resolution is rejected as uncallable on every path. The request team id has one reader. Resolves LIT-7363 Claude-Session: https://claude.ai/code/session_01NU97S7d2FUDDvTk59k53Wp Co-authored-by: Claude Fable 5.1 --- .../guardrails/auto_router_compression.py | 25 +- litellm/router.py | 144 +++++---- litellm/router_utils/common_utils.py | 14 +- .../test_auto_router_compression.py | 26 +- .../proxy/test_common_request_processing.py | 2 +- .../router_strategy/test_complexity_router.py | 4 +- tests/test_litellm/test_router.py | 301 +++++++++++++++++- 7 files changed, 417 insertions(+), 99 deletions(-) diff --git a/litellm/proxy/guardrails/auto_router_compression.py b/litellm/proxy/guardrails/auto_router_compression.py index 98707e7ddca..c37b9fff1f0 100644 --- a/litellm/proxy/guardrails/auto_router_compression.py +++ b/litellm/proxy/guardrails/auto_router_compression.py @@ -80,17 +80,19 @@ def policy_from_litellm_params(litellm_params: Mapping[str, object]) -> AutoRout def policy_for_model( llm_router: "Router | None", model_alias: str, - team_id: str | None, + request_kwargs: Mapping[str, object], request_tags: Sequence[str], ) -> AutoRouterCompressionPolicy | None: - """The compression policy of the auto router marker `model_alias` resolves to. + """The compression policy of the auto router marker `model_alias` resolves to for this caller. - Pre-call arming and the routing hook both resolve through here, so an alias with - several tag-scoped markers cannot suppress under one and then route under another. + Pre-call arming and the routing hook both resolve through here, and here resolves through the + router's own request-scoped deployment lookup, so an alias with several tag-scoped markers + cannot suppress under one and then route under another, and a team router reached by its + public name carries its policy for every principal that can reach it. """ if llm_router is None: return None - deployments: Final = llm_router.get_model_list(model_name=model_alias, team_id=team_id) or () + deployments: Final = llm_router.deployments_for_request(model_alias, request_kwargs) markers: Final = tuple( litellm_params for deployment in deployments @@ -108,17 +110,6 @@ def policy_for_model( return next((policy for policy in candidates if policy is not None), None) -def team_id_from_request(request_kwargs: Mapping[str, object]) -> str | None: - """The caller's team id, from whichever metadata bucket this surface writes to.""" - for meta_key in ("metadata", "litellm_metadata"): - meta = request_kwargs.get(meta_key) - if isinstance(meta, Mapping): - team_id = meta.get("user_api_key_team_id") - if isinstance(team_id, str): - return team_id - return None - - def _compression_guardrail_classes() -> tuple[type, ...]: """The registered guardrail classes whose provider compresses prompts.""" from litellm.proxy.guardrails.guardrail_registry import guardrail_class_registry @@ -172,7 +163,7 @@ async def arm_pre_call( policy: Final = policy_for_model( llm_router=llm_router, model_alias=model_alias, - team_id=team_id_from_request(data), + request_kwargs=data, request_tags=_get_tags_from_request_kwargs(data), ) if policy is None: diff --git a/litellm/router.py b/litellm/router.py index 68ace283949..88601e5b97b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -148,6 +148,7 @@ from litellm.router_utils.common_utils import ( _is_proxy_admin_request, filter_team_based_models, filter_web_search_deployments, + get_request_team_id, resolve_model_group_alias, truncate_fallback_error_detail, warn_on_provider_credential_mismatch, @@ -12337,27 +12338,7 @@ class Router: if team_deployments: return model, team_deployments elif include_team_models: - team_deployments = [ - self.model_list[index] - for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] - team_ids: Final = { - team_id - for deployment in team_deployments - for team_id in [(deployment.get("model_info") or {}).get("team_id")] - if team_id is not None - } - if len(team_ids) > 1: - raise litellm.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) + team_deployments = self._team_deployments_across_teams(model) if team_deployments: return model, team_deployments @@ -12384,6 +12365,45 @@ class Router: return None + def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: + """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" + team_deployments: Final = [ + self.model_list[index] + for (_, public_model_name), indices in self.team_model_to_deployment_indices.items() + if public_model_name == model + for index in indices + ] + team_ids: Final = { + team_id + for deployment in team_deployments + for team_id in [(deployment.get("model_info") or {}).get("team_id")] + if team_id is not None + } + if len(team_ids) > 1: + raise litellm.BadRequestError( + message=( + f"Model name '{model}' matches deployments from multiple teams. " + "Specify the deployment ID directly to disambiguate." + ), + model=model, + llm_provider="", + ) + return team_deployments + + def deployments_for_request( + self, model: str, request_kwargs: Mapping[str, object] + ) -> Sequence[DeploymentTypedDict]: + """The deployments `model` names for this caller, through the same alias, then team-first, then + global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so + strategy selection and compression policy can never disagree with deployment selection about + which marker a name means.""" + registered_name: Final = self._get_model_from_alias(model=model) or model + team_id: Final = get_request_team_id(request_kwargs) + deployments: Final = self._get_all_deployments(model_name=registered_name, team_id=team_id) + if deployments or team_id is not None or not _is_proxy_admin_request(request_kwargs): + return deployments + return self._team_deployments_across_teams(registered_name) + @staticmethod def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool: litellm_params: Final = deployment.get("litellm_params") @@ -12411,11 +12431,7 @@ class Router: - Dict, if specific model chosen """ - request_team_id: str | None = None - if request_kwargs is not None: - metadata: Final = request_kwargs.get("metadata") or {} - litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) # check if aliases set on litellm model alias map if specific_deployment is True: return model, self._get_deployment_by_litellm_model(model=model) @@ -12440,7 +12456,9 @@ class Router: include_team_models=_is_proxy_admin_request(request_kwargs), ) if early is not None: - return early + if not isinstance(early[1], list): + return early + return early[0], self._drop_strategy_markers(early[0], early[1]) ## get healthy deployments ### get all deployments @@ -12517,19 +12535,22 @@ class Router: model ] # update the model to the actual value if an alias has been passed in - marker_flags: Final = tuple(self._is_strategy_marker_deployment(d) for d in healthy_deployments) - if not any(marker_flags): - return model, healthy_deployments - selectable: Final = [ # mutable-ok: matches this function's list contract expected by downstream filters - d for d, is_marker in zip(healthy_deployments, marker_flags, strict=True) if not is_marker + return model, self._drop_strategy_markers(model, healthy_deployments) + + def _drop_strategy_markers( + self, model: str, deployments: Sequence[DeploymentTypedDict] + ) -> list[DeploymentTypedDict]: + """A strategy marker is never a callable deployment, whichever resolution arm produced it.""" + selectable: Final = [ # mutable-ok: matches _common_checks_available_deployment's list contract + d for d in deployments if not self._is_strategy_marker_deployment(d) ] - if not selectable: + if deployments and not selectable: raise litellm.BadRequestError( message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}", model=model, llm_provider="", ) - return model, selectable + return selectable def _filter_deployments_by_model_access_groups( self, @@ -13219,12 +13240,8 @@ class Router: return filtered - def _model_name_has_plain_deployments(self, model: str) -> bool: - indices: Final = self.model_name_to_deployment_indices.get(model) or () - return any(not self._is_strategy_marker_deployment(self.model_list[idx]) for idx in indices) - def _select_pre_routing_strategy( - self, model: str, request_kwargs: dict + self, model: str, request_kwargs: Mapping[str, object] ) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None": """ Resolve the pre-routing strategy for `model`, disambiguating deployments @@ -13235,6 +13252,12 @@ class Router: deployment the strategy was registered from via its (model_name, tags) pair. + The registries are keyed by the marker deployment's own `model_name`, which + for a team-scoped router is the internal `model_name_{team}_{uuid}` while + the caller sends the team's public name. So the names looked up are the + `model_name`s of whatever deployments this caller's request resolves `model` + to, and `model` itself when it resolves to none. + With tag filtering enabled, router-wide or by the request's enable_tag_filtering (which the proxy sets from key/team router_settings), strategies that all carry real tags matching none of @@ -13242,12 +13265,14 @@ class Router: deployments: returning None hands the request to ordinary tag-aware deployment selection. """ - candidates: Final[list[TaggedPreRoutingStrategy[PreRoutingStrategy]]] = [ - *self.auto_routers.get(model, []), - *self.complexity_routers.get(model, []), - *self.adaptive_routers.get(model, []), - *self.quality_routers.get(model, []), - ] + registries: Final = (self.auto_routers, self.complexity_routers, self.adaptive_routers, self.quality_routers) + if not any(registries): + return None + deployments: Final = self.deployments_for_request(model, request_kwargs) + registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,) + candidates: Final = tuple( + tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) + ) if not candidates: return None @@ -13265,7 +13290,7 @@ class Router: if ( (self.enable_tag_filtering or request_scoped_filtering) and all(tagged.tags for tagged in candidates) - and self._model_name_has_plain_deployments(model) + and any(not self._is_strategy_marker_deployment(d) for d in deployments) ): return None return candidates[0] @@ -13377,11 +13402,12 @@ class Router: Used for the litellm auto-router to modify the request before the routing decision is made. - `model` is whatever the caller asked for, which may be a `model_group_alias` key, while the - strategy registries and the marker deployment are keyed by the marker's own `model_name`, so - every lookup below resolves the alias first. Only the lookups: the caller-facing name stays - the alias, since spend metadata is stamped before routing and the response carries the tier - group the strategy picked. + `model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's + public model name, while the strategy registries and the marker deployment are keyed by the + marker's own `model_name`, so every lookup below resolves the alias first and the team name + through the deployment path. Only the lookups: the caller-facing name stays the alias, since + spend metadata is stamped before routing and the response carries the tier group the + strategy picked. """ requested_registered_model_name: Final = self._get_model_from_alias(model=model) or model registered_model_name: Final = await self._resolve_claude_code_session_router( @@ -13418,7 +13444,6 @@ class Router: messages_for_routing, model_hop_compression_armed, policy_for_model, - team_id_from_request, ) # Same tag-aware lookup the proxy's pre-call arming used, so an alias with @@ -13426,7 +13451,7 @@ class Router: compression_policy: Final = policy_for_model( llm_router=self, model_alias=registered_model_name, - team_id=team_id_from_request(request_kwargs), + request_kwargs=request_kwargs, request_tags=_get_tags_from_request_kwargs(request_kwargs), ) # Shared compression already ran in the pre-call hook, so reuse it rather than @@ -13495,7 +13520,9 @@ class Router: # Per-tier `litellm_params` on the hook response are deliberate overrides # the caller applies on top, so those keys are never forwarded here. marker_params: Final = ( - self._forwardable_alias_marker_params(model=registered_model_name, strategy_tags=selected_strategy.tags) + self._forwardable_alias_marker_params( + model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs + ) if pre_routing_hook_response is not None else () ) @@ -13513,13 +13540,14 @@ class Router: return pre_routing_hook_response def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...] + self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] ) -> tuple[tuple[str, object], ...]: marker_params: Final = tuple( litellm_params - for idx in self.model_name_to_deployment_indices.get(model, ()) - if isinstance(litellm_params := self.model_list[idx].get("litellm_params", {}), dict) - and str(litellm_params.get("model", "")).startswith(AUTO_ROUTER_MODEL_PREFIX) + for deployment in self.deployments_for_request(model, request_kwargs) + if str((litellm_params := deployment["litellm_params"]).get("model", "")).startswith( + AUTO_ROUTER_MODEL_PREFIX + ) ) tag_matched: Final = tuple( params for params in marker_params if tuple(params.get("tags") or ()) == strategy_tags diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 280a7defcf8..c1c6d25beca 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -26,6 +26,18 @@ def _is_proxy_admin_request(request_kwargs: Mapping[str, object] | None) -> bool return getattr(user_api_key_auth, "user_role", None) == "proxy_admin" +def get_request_team_id(request_kwargs: Mapping[str, object] | None) -> str | None: + """The caller's team id, from whichever metadata bucket this surface writes to.""" + if request_kwargs is None: + return None + for bucket_name in ("metadata", "litellm_metadata"): + bucket = request_kwargs.get(bucket_name) + team_id = bucket.get("user_api_key_team_id") if isinstance(bucket, Mapping) else None + if isinstance(team_id, str) and team_id: + return team_id + return None + + def resolve_model_group_alias(model_group_alias: object, model: str) -> str | None: """ Resolve ``model`` through a ``model_group_alias`` map. @@ -110,7 +122,7 @@ def filter_team_based_models( metadata: Final = request_kwargs.get("metadata") or {} litellm_metadata: Final = request_kwargs.get("litellm_metadata") or {} - request_team_id: Final = metadata.get("user_api_key_team_id") or litellm_metadata.get("user_api_key_team_id") + request_team_id: Final = get_request_team_id(request_kwargs) if request_team_id is None and _is_proxy_admin_request(request_kwargs) and isinstance(healthy_deployments, list): requested_model: Final = ( request_kwargs.get("model") or metadata.get("model_group") or litellm_metadata.get("model_group") diff --git a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py index db2f94306fb..0d723d6671f 100644 --- a/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py +++ b/tests/test_litellm/proxy/guardrails/test_auto_router_compression.py @@ -56,13 +56,13 @@ class TestPolicyFromLitellmParams: class _FakeRouter: - """Minimal stand-in for litellm.Router.get_model_list, for policy_for_model.""" + """Minimal stand-in for litellm.Router.deployments_for_request, for policy_for_model.""" def __init__(self, deployments: list[dict[str, Any]]): self._deployments = deployments - def get_model_list(self, model_name, team_id=None): - return [d for d in self._deployments if d.get("model_name") == model_name] + def deployments_for_request(self, model, request_kwargs): + return [d for d in self._deployments if d.get("model_name") == model] def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[str, Any]: @@ -78,23 +78,23 @@ def _marker(compression: dict[str, str], tags: list[str] | None = None) -> dict[ class TestPolicyForModel: def test_no_router_returns_none(self): - assert policy_for_model(llm_router=None, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=None, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_no_marker_deployment_returns_none(self): router = _FakeRouter([{"model_name": "smart-router", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_without_policy_returns_none(self): router = _FakeRouter( [{"model_name": "smart-router", "litellm_params": {"model": "auto_router/complexity_router"}}] ) - assert policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) is None + assert policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) is None def test_marker_deployment_with_policy_is_found(self): router = _FakeRouter( [_marker({"auto_router_routing_compression": "headroom-a", "auto_router_model_compression": "none"})] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=()) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=()) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) def test_picks_the_marker_whose_tags_the_request_carries(self): @@ -107,8 +107,8 @@ class TestPolicyForModel: ] ) - eu = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) - us = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + eu = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) + us = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert eu == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) assert us == AutoRouterCompressionPolicy(routing="headroom-us", model=None) @@ -116,7 +116,7 @@ class TestPolicyForModel: def test_untagged_marker_matches_any_request(self): router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-a"})]) policy = policy_for_model( - llm_router=router, model_alias="smart-router", team_id=None, request_tags=("anything",) + llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("anything",) ) assert policy == AutoRouterCompressionPolicy(routing="headroom-a", model=None) @@ -128,14 +128,14 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-default"}), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-default", model=None) def test_no_untagged_fallback_means_no_policy(self): """No matching marker means no policy, not an unrelated slice's compression.""" router = _FakeRouter([_marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"])]) assert ( - policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("us",)) is None + policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("us",)) is None ) def test_tag_scoped_marker_takes_precedence_over_untagged(self): @@ -147,7 +147,7 @@ class TestPolicyForModel: _marker({"auto_router_routing_compression": "headroom-eu"}, tags=["eu"]), ] ) - policy = policy_for_model(llm_router=router, model_alias="smart-router", team_id=None, request_tags=("eu",)) + policy = policy_for_model(llm_router=router, model_alias="smart-router", request_kwargs={}, request_tags=("eu",)) assert policy == AutoRouterCompressionPolicy(routing="headroom-eu", model=None) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index d2ba6527826..50b26577e5c 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -417,7 +417,7 @@ class TestProxyBaseLLMRequestProcessing: ) fake_llm_router = MagicMock() - fake_llm_router.get_model_list.return_value = [ + fake_llm_router.deployments_for_request.return_value = [ { "model_name": "smart-router", "litellm_params": { diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 9f925499f1a..51103297c58 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3846,11 +3846,11 @@ class TestRouterPreRoutingSharedAliasName: def test_forwardable_alias_marker_params_reads_the_marker_entry_only(self): router = Router(model_list=[self._plain_entry(), self._marker_entry(), self._tier_entry()]) - forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=())) + forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=(), request_kwargs={})) assert forwarded["drop_params"] is True assert "api_key" not in forwarded and "api_base" not in forwarded - assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=()) == () + assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=(), request_kwargs={}) == () @staticmethod def _region_marker_entry() -> dict: diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 80da922724d..1e0327e5ac2 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -10003,13 +10003,6 @@ class TestTaggedAutoRouterOnSharedModelName: def test_deployment_without_litellm_params_mapping_is_not_a_marker(self): assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False - def test_model_name_has_plain_deployments_reflects_the_pool(self): - mixed = self._router(marker_tags=["route"], include_plain_sibling=True, enable_tag_filtering=True) - marker_only = self._router(marker_tags=["route"], include_plain_sibling=False, enable_tag_filtering=True) - - assert mixed._model_name_has_plain_deployments("gpt4o") is True - assert marker_only._model_name_has_plain_deployments("gpt4o") is False - class TestAutoRouterSharedModelNameConnectionParams: """A plain deployment sharing its model_name with an `auto_router/` marker must not have @@ -10618,6 +10611,300 @@ class TestModelGroupAliasReachesPreRoutingStrategies: ) +class TestTeamPublicNameReachesPreRoutingStrategies: + """A team-scoped strategy router is stored under an internal `model_name_{team}_{uuid}` with the + caller-facing name in `model_info.team_public_model_name`, and the four registries key on that + internal name. A team key asks for the public name, so the hook has to resolve it to the team's + marker through the same team-first resolution the deployment path uses, and a resolution that + yields only markers is not callable on any path (LIT-7363).""" + + MARKER_TIMEOUT = 42.0 + REGISTRY_NAMES = ("auto_routers", "complexity_routers", "adaptive_routers", "quality_routers") + TEAM = "team-a" + OTHER_TEAM = "team-b" + PUBLIC_NAME = "smart-route" + INTERNAL_NAME = "model_name_team-a_0b3c" + SIBLING_INTERNAL_NAME = "model_name_team-a_9e1d" + + class _RewriteStrategy: + def __init__(self, rewrite_to: str = "gemini-flash"): + self.rewrite_to = rewrite_to + + async def async_pre_routing_hook( + self, model, request_kwargs, messages=None, input=None, specific_deployment=False + ): + from litellm.types.router import PreRoutingHookResponse + + return PreRoutingHookResponse(model=self.rewrite_to, messages=messages) + + @classmethod + def _team_marker(cls, internal_name: str, tags: list[str] | None = None) -> dict: + tiers = dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), "gemini-flash") + return { + "model_name": internal_name, + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": {"tiers": tiers}, + "complexity_router_default_model": "gemini-flash", + "timeout": cls.MARKER_TIMEOUT, + **({"tags": tags} if tags else {}), + }, + "model_info": {"team_id": cls.TEAM, "team_public_model_name": cls.PUBLIC_NAME}, + } + + @classmethod + def _router( + cls, + registrations: dict[str, "TestTeamPublicNameReachesPreRoutingStrategies._RewriteStrategy"], + registry_name: str = "complexity_routers", + extra_deployments: tuple[dict, ...] = (), + markers: tuple[dict, ...] | None = None, + enable_tag_filtering: bool = False, + ) -> "litellm.Router": + from litellm.types.router import TaggedPreRoutingStrategy + + markers = markers if markers is not None else (cls._team_marker(cls.INTERNAL_NAME),) + tier = { + "model_name": "gemini-flash", + "litellm_params": {"model": "gemini/gemini-3.6-flash", "mock_response": "routed by the tier"}, + } + router = litellm.Router( + model_list=[*markers, tier, *extra_deployments], + enable_tag_filtering=enable_tag_filtering, + ) + tags_by_name = {m["model_name"]: tuple(m["litellm_params"].get("tags") or ()) for m in markers} + for name in cls.REGISTRY_NAMES: + setattr(router, name, {}) + setattr( + router, + registry_name, + { + name: [TaggedPreRoutingStrategy(tags=tags_by_name[name], strategy=strategy)] + for name, strategy in registrations.items() + }, + ) + return router + + @staticmethod + def _messages() -> list[dict[str, str]]: + return [{"role": "user", "content": "What is the capital of France?"}] + + @classmethod + def _team_request(cls, team_id: str | None = "team-a", tags: list[str] | None = None) -> dict: + metadata = {**({"user_api_key_team_id": team_id} if team_id else {}), **({"tags": tags} if tags else {})} + return {"metadata": metadata} + + @pytest.mark.parametrize("registry_name", REGISTRY_NAMES) + @pytest.mark.asyncio + async def test_team_key_dispatches_to_the_strategy_registered_under_the_internal_name(self, registry_name): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}, registry_name=registry_name) + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + + @pytest.mark.asyncio + async def test_team_key_deployment_selection_lands_on_the_tier_and_forwards_the_marker_params(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = self._team_request() + + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert deployment["litellm_params"]["model"] == "gemini/gemini-3.6-flash" + assert request_kwargs["timeout"] == self.MARKER_TIMEOUT + + @pytest.mark.asyncio + async def test_another_team_never_reaches_the_strategy_or_the_marker(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + + assert ( + await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + is None + ) + with pytest.raises(litellm.BadRequestError): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(self.OTHER_TEAM), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_sibling_team_markers_select_by_request_tag_then_default(self): + router = self._router( + { + self.INTERNAL_NAME: self._RewriteStrategy("cn-model"), + self.SIBLING_INTERNAL_NAME: self._RewriteStrategy("us-model"), + }, + markers=( + self._team_marker(self.INTERNAL_NAME, tags=["cn"]), + self._team_marker(self.SIBLING_INTERNAL_NAME, tags=["us", "default"]), + ), + ) + + async def routed(tags: list[str] | None) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=tags), messages=self._messages() + ) + return response.model if response else None + + assert await routed(["cn"]) == "cn-model" + assert await routed(["us"]) == "us-model" + assert await routed(None) == "us-model" + + @pytest.mark.asyncio + async def test_team_public_name_shadows_a_global_model_for_that_team_only(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": self.PUBLIC_NAME, "litellm_params": {"model": "openai/gpt-4o"}},), + ) + + async def routed(request_kwargs: dict) -> str | None: + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return response.model if response else None + + async def selected(request_kwargs: dict) -> str: + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + return deployment["litellm_params"]["model"] + + assert await routed(self._team_request()) == "gemini-flash" + assert await selected(self._team_request()) == "gemini/gemini-3.6-flash" + for request_kwargs in (self._team_request(None), self._team_request(self.OTHER_TEAM)): + assert await routed(request_kwargs) is None + assert await selected(request_kwargs) == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_tag_filtering_hands_untagged_team_requests_to_the_team_plain_sibling(self): + plain_sibling = { + "model_name": self.SIBLING_INTERNAL_NAME, + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": {"team_id": self.TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + markers=(self._team_marker(self.INTERNAL_NAME, tags=["route"]),), + extra_deployments=(plain_sibling,), + enable_tag_filtering=True, + ) + + tagged = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(tags=["route"]), messages=self._messages() + ) + assert tagged is not None and tagged.model == "gemini-flash" + for _ in range(20): + deployment = await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + assert deployment["litellm_params"]["model"] == "openai/gpt-4o" + + @pytest.mark.asyncio + async def test_marker_only_team_resolution_is_rejected_as_uncallable(self): + import re + + from litellm.types.router import RouterErrors + + router = self._router({}) + + with pytest.raises( + litellm.BadRequestError, match=re.escape(RouterErrors.only_strategy_marker_deployments.value) + ): + await router.async_get_available_deployment( + model=self.PUBLIC_NAME, request_kwargs=self._team_request(), messages=self._messages() + ) + + @pytest.mark.asyncio + async def test_proxy_admin_without_a_team_reaches_the_team_strategy_by_public_name(self): + router = self._router({self.INTERNAL_NAME: self._RewriteStrategy()}) + request_kwargs = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + + response = await router.async_pre_routing_hook( + model=self.PUBLIC_NAME, request_kwargs=request_kwargs, messages=self._messages() + ) + + assert response is not None + assert response.model == "gemini-flash" + @pytest.mark.asyncio + async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): + router = self._router( + {self.INTERNAL_NAME: self._RewriteStrategy()}, + extra_deployments=({"model_name": "shared-name", "litellm_params": {"model": "openai/gpt-4o"}},), + ) + principals = { + "team": self._team_request(), + "other-team": self._team_request(self.OTHER_TEAM), + "no-team": self._team_request(None), + "admin": {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}}, + } + for principal, request_kwargs in principals.items(): + for model in (self.PUBLIC_NAME, "shared-name", "gemini-flash", "missing"): + resolved = [d["model_name"] for d in router.deployments_for_request(model, request_kwargs)] + callable_names = [ + name + for name, deployment in zip(resolved, router.deployments_for_request(model, request_kwargs)) + if not router._is_strategy_marker_deployment(deployment) + ] + if resolved and not callable_names: + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + elif not resolved: + with pytest.raises(litellm.BadRequestError): + router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) + else: + _, deployments = router._common_checks_available_deployment( + model=model, request_kwargs=request_kwargs + ) + assert [d["model_name"] for d in deployments] == callable_names, (principal, model) + + def test_drop_strategy_markers_keeps_plain_deployments_and_rejects_marker_only_sets(self): + router = self._router({}) + marker = router.model_list[0] + plain = {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o"}} + + assert router._drop_strategy_markers("x", [marker, plain]) == [plain] + assert router._drop_strategy_markers("x", [plain]) == [plain] + assert router._drop_strategy_markers("x", []) == [] + with pytest.raises(litellm.BadRequestError, match="strategy router marker"): + router._drop_strategy_markers("x", [marker]) + + def test_team_deployments_across_teams_unions_one_team_and_rejects_two(self): + other_team_marker = { + **self._team_marker(self.SIBLING_INTERNAL_NAME), + "model_info": {"team_id": self.OTHER_TEAM, "team_public_model_name": self.PUBLIC_NAME}, + } + one_team = self._router({}) + two_teams = self._router({}, markers=(self._team_marker(self.INTERNAL_NAME), other_team_marker)) + + assert [d["model_name"] for d in one_team._team_deployments_across_teams(self.PUBLIC_NAME)] == [ + self.INTERNAL_NAME + ] + assert one_team._team_deployments_across_teams("missing") == [] + with pytest.raises(litellm.BadRequestError, match="multiple teams"): + two_teams._team_deployments_across_teams(self.PUBLIC_NAME) + + + def test_compression_policy_follows_the_same_resolution_for_every_principal(self): + from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model + + marker = self._team_marker(self.INTERNAL_NAME) + marker["litellm_params"]["auto_router_routing_compression"] = "headroom-team" + router = self._router({}, markers=(marker,)) + admin = {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}} + expected = AutoRouterCompressionPolicy(routing="headroom-team", model=None) + + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(), ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, admin, ()) == expected + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(self.OTHER_TEAM), ()) is None + assert policy_for_model(router, self.PUBLIC_NAME, self._team_request(None), ()) is None + + class TestAutoRouterCompressionDecoupling: """An auto router's `auto_router_routing_compression` / `auto_router_model_compression` decouple what the routing decision sees from what the model call sees. The one