mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
fix(router): address Greptile P1/P2 review comments
- Add deduplication guard in _update_team_model_index to prevent duplicate indices - Add wildcard comment in map_team_model for clarity - Add monkeypatch to test_team_alias_stale_bypass_disabled_by_default for determinism - Extract _get_team_deployments helper to centralize DB access pattern - Add clarifying comments for team_public_model_name assignment ordering Made-with: Cursor
This commit is contained in:
parent
08145c25c1
commit
cb466689a9
3 changed files with 40 additions and 21 deletions
|
|
@ -329,12 +329,17 @@ async def _add_team_model_to_db(
|
|||
_team_id = model_params.model_info.team_id
|
||||
if _team_id is None:
|
||||
return None
|
||||
# Capture the original public name before mutating model_params.model_name
|
||||
original_model_name = model_params.model_name
|
||||
|
||||
# Generate unique internal model_name for team-scoped deployment
|
||||
unique_model_name = f"model_name_{_team_id}_{uuid.uuid4()}"
|
||||
|
||||
# Store public name in model_info BEFORE overwriting model_name
|
||||
# so _add_model_to_db serializes the correct team_public_model_name
|
||||
if original_model_name:
|
||||
model_params.model_info.team_public_model_name = original_model_name
|
||||
|
||||
unique_model_name = f"model_name_{_team_id}_{uuid.uuid4()}"
|
||||
|
||||
model_params.model_name = unique_model_name
|
||||
|
||||
## CREATE MODEL IN DB ##
|
||||
|
|
@ -458,6 +463,25 @@ async def _setup_new_team_model_assignment(
|
|||
)
|
||||
|
||||
|
||||
async def _get_team_deployments(
|
||||
team_id: str, prisma_client: PrismaClient
|
||||
) -> List[LiteLLM_ProxyModelTable]:
|
||||
"""
|
||||
Fetch all deployments for a given team_id from the database.
|
||||
|
||||
Centralizes team deployment queries to ensure consistent filtering and error handling.
|
||||
"""
|
||||
response = await prisma_client.db.litellm_proxymodeltable.find_many(
|
||||
where={
|
||||
"model_info": {
|
||||
"path": ["team_id"],
|
||||
"equals": team_id,
|
||||
}
|
||||
}
|
||||
)
|
||||
return response if response else []
|
||||
|
||||
|
||||
async def _update_existing_team_model_assignment(
|
||||
team_id: str,
|
||||
public_model_name: str,
|
||||
|
|
@ -504,23 +528,14 @@ async def _update_existing_team_model_assignment(
|
|||
)
|
||||
return
|
||||
|
||||
response = await prisma_client.db.litellm_proxymodeltable.find_many(
|
||||
where={
|
||||
"model_info": {
|
||||
"path": ["team_id"],
|
||||
"equals": team_id,
|
||||
}
|
||||
}
|
||||
)
|
||||
if not response:
|
||||
other_deployments_with_old_name = []
|
||||
else:
|
||||
other_deployments_with_old_name = [
|
||||
d
|
||||
for d in response
|
||||
if d.model_name != db_model.model_name
|
||||
and _get_team_public_model_name(d.model_info) == old_public_name
|
||||
]
|
||||
# Query DB for all team deployments to check for sibling deployments
|
||||
team_deployments = await _get_team_deployments(team_id, prisma_client)
|
||||
other_deployments_with_old_name = [
|
||||
d
|
||||
for d in team_deployments
|
||||
if d.model_name != db_model.model_name
|
||||
and _get_team_public_model_name(d.model_info) == old_public_name
|
||||
]
|
||||
|
||||
# Add new name first, then delete old name to prevent access loss on partial failure
|
||||
await team_model_add(
|
||||
|
|
|
|||
|
|
@ -7191,7 +7191,8 @@ class Router:
|
|||
key = (team_id, team_public_model_name)
|
||||
if key not in self.team_model_to_deployment_indices:
|
||||
self.team_model_to_deployment_indices[key] = []
|
||||
self.team_model_to_deployment_indices[key].append(idx)
|
||||
if idx not in self.team_model_to_deployment_indices[key]:
|
||||
self.team_model_to_deployment_indices[key].append(idx)
|
||||
|
||||
def _add_model_to_list_and_index_map(
|
||||
self, model: dict, model_id: Optional[str] = None
|
||||
|
|
@ -8217,6 +8218,8 @@ class Router:
|
|||
if model.get("model_info", {}).get("team_id") == team_id:
|
||||
return team_model_name
|
||||
|
||||
# No team-scoped deployment found; wildcard/pattern routes are
|
||||
# handled downstream by the pattern_router in _common_checks_available_deployment.
|
||||
return None
|
||||
|
||||
def should_include_deployment(
|
||||
|
|
|
|||
|
|
@ -2044,7 +2044,8 @@ def test_update_model_if_team_alias_exists(data, user_api_key_dict, expected_mod
|
|||
assert test_data.get("model") == expected_model
|
||||
|
||||
|
||||
def test_team_alias_stale_bypass_disabled_by_default():
|
||||
def test_team_alias_stale_bypass_disabled_by_default(monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", raising=False)
|
||||
from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists
|
||||
|
||||
class _MockRouter:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue