mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(router): keep team wildcard routers fresh and prioritize them over global patterns
team_pattern_routers retained deleted/replaced deployments, so team users could keep resolving stale credentials; now set_model_list resets the registry and deployment removal prunes it. Also consult the team wildcard router before the global pattern_router in get_deployment_credentials_with_provider so a global pattern like "openai/*" no longer shadows the team's own entry Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
bea11ddedd
commit
836bf0807b
3 changed files with 121 additions and 8 deletions
|
|
@ -7876,6 +7876,7 @@ class Router:
|
|||
self.model_id_to_deployment_index_map = {} # Reset the index
|
||||
self.model_name_to_deployment_indices = {} # Reset the model_name index
|
||||
self.team_model_to_deployment_indices = {} # Reset the team_model index
|
||||
self.team_pattern_routers = {}
|
||||
self.team_public_model_names = frozenset()
|
||||
# Reset per-strategy router registries so hot-reload doesn't leave
|
||||
# stale routers pointing at the old model_list.
|
||||
|
|
@ -8232,6 +8233,12 @@ class Router:
|
|||
public_model_name for _, public_model_name in self.team_model_to_deployment_indices
|
||||
)
|
||||
|
||||
for team_id in list(self.team_pattern_routers.keys()):
|
||||
team_pattern_router = self.team_pattern_routers[team_id]
|
||||
team_pattern_router.remove_deployment(model_id)
|
||||
if not team_pattern_router.patterns:
|
||||
del self.team_pattern_routers[team_id]
|
||||
|
||||
def _update_team_model_index(self, model: dict, idx: int) -> None:
|
||||
"""
|
||||
Helper to update team_model_to_deployment_indices for a single deployment.
|
||||
|
|
@ -8460,8 +8467,8 @@ class Router:
|
|||
return None
|
||||
|
||||
def get_deployment_credentials_with_provider(
|
||||
self, model_id: str, team_id: Optional[str] = None
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
self, model_id: str, team_id: str | None = None
|
||||
) -> dict[str, Any] | None:
|
||||
"""
|
||||
Get API credentials and provider info from a model name in model_list.
|
||||
Useful for passthrough endpoints (files, batches, etc.) that need credentials.
|
||||
|
|
@ -8492,19 +8499,22 @@ class Router:
|
|||
if deployment is None:
|
||||
deployment = self.get_deployment_by_model_group_name(model_group_name=model_id)
|
||||
|
||||
# If not found, check team-scoped deployments (team public model names,
|
||||
# e.g. team wildcard models like "openai/*", live in a separate index).
|
||||
# If not found, check team-scoped deployments whose team public model
|
||||
# name exactly matches model_id (wildcard team names are matched via
|
||||
# team_pattern_routers below).
|
||||
if deployment is None and team_id is not None:
|
||||
team_indices = self.team_model_to_deployment_indices.get((team_id, model_id), [])
|
||||
if team_indices:
|
||||
team_model = self.model_list[team_indices[0]]
|
||||
deployment = Deployment(**team_model) if isinstance(team_model, dict) else team_model
|
||||
|
||||
# If still not found, check for wildcard pattern matches
|
||||
# If still not found, check for wildcard pattern matches. Team wildcard
|
||||
# matches take priority so a global pattern (e.g. "openai/*") doesn't
|
||||
# shadow the team's own entry.
|
||||
if deployment is None:
|
||||
potential_wildcard_models = self.pattern_router.route(model_id) or []
|
||||
if not potential_wildcard_models and team_id is not None and team_id in self.team_pattern_routers:
|
||||
potential_wildcard_models = self.team_pattern_routers[team_id].route(model_id) or []
|
||||
team_pattern_router = self.team_pattern_routers.get(team_id) if team_id is not None else None
|
||||
team_wildcard_models = (team_pattern_router.route(model_id) or []) if team_pattern_router else []
|
||||
potential_wildcard_models = team_wildcard_models or self.pattern_router.route(model_id) or []
|
||||
if potential_wildcard_models:
|
||||
# Use the first matching wildcard deployment
|
||||
deployment_dict = potential_wildcard_models[0]
|
||||
|
|
|
|||
|
|
@ -73,6 +73,17 @@ class PatternMatchRouter:
|
|||
self.patterns[regex] = []
|
||||
self.patterns[regex].append(llm_deployment)
|
||||
|
||||
def remove_deployment(self, model_id: str) -> None:
|
||||
"""
|
||||
Remove every deployment with the given model id from the pattern registry,
|
||||
dropping any pattern whose deployment list becomes empty.
|
||||
"""
|
||||
self.patterns = {
|
||||
regex: remaining
|
||||
for regex, deployments in self.patterns.items()
|
||||
if (remaining := [d for d in deployments if (d.get("model_info") or {}).get("id") != model_id])
|
||||
}
|
||||
|
||||
def _pattern_to_regex(self, pattern: str) -> str:
|
||||
"""
|
||||
Convert a wildcard pattern to a regex pattern
|
||||
|
|
|
|||
|
|
@ -3535,6 +3535,98 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name():
|
|||
litellm.credential_list = []
|
||||
|
||||
|
||||
def _team_wildcard_model(api_key: str, model_id: str = "team-wildcard-id") -> dict:
|
||||
return {
|
||||
"model_name": f"model_name_team-1_{model_id}",
|
||||
"litellm_params": {"model": "openai/*", "api_key": api_key},
|
||||
"model_info": {
|
||||
"id": model_id,
|
||||
"team_id": "team-1",
|
||||
"team_public_model_name": "openai/*",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_get_deployment_credentials_with_provider_team_wildcard_priority():
|
||||
"""
|
||||
Regression: a global wildcard pattern (e.g. "openai/*") must not shadow a
|
||||
team's own wildcard entry. When team_id is provided, the team wildcard
|
||||
deployment's credentials win; without team_id the global one is used.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "global-key"},
|
||||
},
|
||||
_team_wildcard_model(api_key="team-key"),
|
||||
],
|
||||
)
|
||||
|
||||
team_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-1"
|
||||
)
|
||||
assert team_credentials is not None
|
||||
assert team_credentials["api_key"] == "team-key"
|
||||
|
||||
global_credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2"
|
||||
)
|
||||
assert global_credentials is not None
|
||||
assert global_credentials["api_key"] == "global-key"
|
||||
|
||||
|
||||
def test_team_wildcard_credentials_not_usable_after_delete_deployment():
|
||||
"""
|
||||
Regression: team_pattern_routers retained deleted deployments, so a team
|
||||
user could keep resolving credentials of a deleted wildcard deployment.
|
||||
"""
|
||||
router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")])
|
||||
|
||||
assert (
|
||||
router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-1"
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
router.delete_deployment(id="team-wildcard-id")
|
||||
|
||||
assert (
|
||||
router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-1"
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list():
|
||||
"""
|
||||
Regression: replacing a team wildcard deployment (upsert or model list
|
||||
reload) must serve the new credentials, not the stale cached ones.
|
||||
"""
|
||||
from litellm.types.router import Deployment
|
||||
|
||||
router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")])
|
||||
|
||||
router.upsert_deployment(
|
||||
deployment=Deployment(**_team_wildcard_model(api_key="new-key"))
|
||||
)
|
||||
credentials = router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-1"
|
||||
)
|
||||
assert credentials is not None
|
||||
assert credentials["api_key"] == "new-key"
|
||||
|
||||
router.set_model_list(model_list=[])
|
||||
assert (
|
||||
router.get_deployment_credentials_with_provider(
|
||||
model_id="openai/gpt-5.2", team_id="team-1"
|
||||
)
|
||||
is None
|
||||
)
|
||||
|
||||
|
||||
def test_get_available_guardrail_single_deployment():
|
||||
"""
|
||||
Test get_available_guardrail returns the single guardrail when only one exists.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue