mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
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 <noreply@anthropic.com>
This commit is contained in:
parent
6e96885629
commit
eb45a088d3
7 changed files with 417 additions and 99 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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": {
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue