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 9

fix(team-routing): use deterministic team model group names

Browse source 
Use a deterministic internal model_name for team-scoped deployments so sibling deployments with the same public model share a routing group. This makes team alias writes idempotent and preserves multi-deployment failover/load balancing behavior.

Made-with: Cursor
This commit is contained in:
Sameer Kankute 2026-03-23 15:31:43 +05:30 • committed by shivam
parent 61409275c8
commit 415f53e24e
2 changed files with 172 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -13,13 +13,13 @@ model/{model_id}/update - PATCH endpoint for model update.
import asyncio
import datetime
import json
from litellm._uuid import uuid
from typing import Dict, List, Literal, Optional, Tuple, Union, cast
from fastapi import APIRouter, Depends, HTTPException, Request, status
from pydantic import BaseModel, ConfigDict, Field
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._types import (
CommonProxyErrors,
@ -322,9 +322,13 @@ async def _add_team_model_to_db(
"""
If 'team_id' is provided,
- generate a unique 'model_name' for the model (e.g. 'model_name_{team_id}_{uuid})
- store the model in the db with the unique 'model_name'
- store a team model alias mapping {"model_name": "model_name_{team_id}_{uuid}"}
- generate a deterministic 'model_name' for the model (e.g. 'model_name_{team_id}_{public_name}')
- store the model in the db with this shared group name
- store a team model alias mapping {"public_name": "model_name_{team_id}_{public_name}"}
Using a deterministic name (not UUID) ensures sibling deployments for the
same public model share a model_name, so the router treats them as a single
candidate pool for load balancing and failover.
"""
_team_id = model_params.model_info.team_id
if _team_id is None:
@ -333,9 +337,9 @@ async def _add_team_model_to_db(
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()}"
group_model_name = f"model_name_{_team_id}_{original_model_name}"
model_params.model_name = unique_model_name
model_params.model_name = group_model_name
## CREATE MODEL IN DB ##
model_response = await _add_model_to_db(
@ -348,7 +352,7 @@ async def _add_team_model_to_db(
await update_team(
data=UpdateTeamRequest(
team_id=_team_id,
model_aliases={original_model_name: unique_model_name},
model_aliases={original_model_name: group_model_name},
),
user_api_key_dict=user_api_key_dict,
http_request=Request(scope={"type": "http"}),
@ -453,14 +457,14 @@ async def _setup_new_team_model_assignment(
patch_data: updateDeployment,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Set up a new team model with unique name, alias, and team membership."""
unique_model_name = f"model_name_{team_id}_{uuid.uuid4()}"
patch_data.model_name = unique_model_name
"""Set up a new team model with deterministic name, alias, and team membership."""
group_model_name = f"model_name_{team_id}_{public_model_name}"
patch_data.model_name = group_model_name
await update_team(
data=UpdateTeamRequest(
team_id=team_id,
model_aliases={public_model_name: unique_model_name},
model_aliases={public_model_name: group_model_name},
),
user_api_key_dict=user_api_key_dict,
http_request=Request(scope={"type": "http"}),

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

@ -1,13 +1,14 @@
import json
import os
import sys
from litellm._uuid import uuid
from typing import Dict, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi.testclient import TestClient
from litellm._uuid import uuid
sys.path.insert(
0, os.path.abspath("../../../..")
) # Adds the parent directory to the system path
@ -399,7 +400,9 @@ class TestClearCache:
"""
Test that clear_cache clears DB models and preserves config models.
"""
from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache
from litellm.proxy.management_endpoints.model_management_endpoints import (
clear_cache,
)
# Create mock router with mixed DB and config models
mock_router = MagicMock()
@ -407,18 +410,18 @@ class TestClearCache:
{
"model_name": "gpt-4",
"model_info": {"id": "db-model-1", "db_model": True},
"litellm_params": {"model": "gpt-4"}
"litellm_params": {"model": "gpt-4"},
},
{
"model_name": "gpt-3.5-turbo",
"model_name": "gpt-3.5-turbo",
"model_info": {"id": "config-model-1", "db_model": False},
"litellm_params": {"model": "gpt-3.5-turbo"}
"litellm_params": {"model": "gpt-3.5-turbo"},
},
{
"model_name": "claude-3",
"model_info": {"id": "db-model-2", "db_model": True},
"litellm_params": {"model": "claude-3"}
}
"litellm_params": {"model": "claude-3"},
},
]
mock_router.delete_deployment = MagicMock(return_value=True)
mock_router.auto_routers = MagicMock()
@ -466,8 +469,8 @@ class TestUpdatePublicModelGroups:
"""
import litellm
from litellm.proxy.management_endpoints.model_management_endpoints import (
update_public_model_groups,
UpdatePublicModelGroupsRequest,
update_public_model_groups,
)
old_db_models = ["db-model-1", "db-model-2"]
@ -525,7 +528,10 @@ class TestUpdatePublicModelGroups:
)
old_links = {"Old Doc": "https://old.example.com"}
new_links = {"New Doc": "https://new.example.com", "API Ref": "https://api.example.com"}
new_links = {
"New Doc": "https://new.example.com",
"API Ref": "https://api.example.com",
}
async def mock_get_config(*args, **kwargs):
litellm.public_model_groups_links = old_links
@ -558,6 +564,100 @@ class TestUpdatePublicModelGroups:
litellm.public_model_groups_links = original_value
class TestTeamModelAliasSiblingOverwrite:
"""
Verify that two sibling team deployments for the same public model name
produce the same deterministic internal model_name, so the alias write
is idempotent and the router groups both deployments together.
"""
@pytest.mark.asyncio
async def test_sibling_team_models_share_deterministic_name(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_team_model_to_db,
)
from litellm.types.router import ModelInfo
team_id = "team_alias_overwrite"
public_name = "gpt-4.1-mini"
captured_alias_calls = []
async def mock_update_team(data, user_api_key_dict, http_request):
if data.model_aliases:
captured_alias_calls.append(dict(data.model_aliases))
async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client):
return MagicMock(model_id=str(uuid.uuid4()))
async def mock_team_model_add(data, http_request, user_api_key_dict):
pass
user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
prisma_client = MockPrismaClient(team_exists=True)
deployment_1 = Deployment(
model_name=public_name,
litellm_params=LiteLLM_Params(
model="azure/gpt-4o-mini",
api_key="key-1",
api_base="https://eastus.example.openai.azure.com",
),
model_info=ModelInfo(team_id=team_id),
)
deployment_2 = Deployment(
model_name=public_name,
litellm_params=LiteLLM_Params(
model="azure/gpt-4o-mini",
api_key="key-2",
api_base="https://westus.example.openai.azure.com",
),
model_info=ModelInfo(team_id=team_id),
)
with patch(
"litellm.proxy.management_endpoints.model_management_endpoints.update_team",
side_effect=mock_update_team,
), patch(
"litellm.proxy.management_endpoints.model_management_endpoints._add_model_to_db",
side_effect=mock_add_model_to_db,
), patch(
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
side_effect=mock_team_model_add,
):
await _add_team_model_to_db(
model_params=deployment_1,
user_api_key_dict=user,
prisma_client=prisma_client,
)
await _add_team_model_to_db(
model_params=deployment_2,
user_api_key_dict=user,
prisma_client=prisma_client,
)
assert len(captured_alias_calls) == 2
internal_name_1 = captured_alias_calls[0][public_name]
internal_name_2 = captured_alias_calls[1][public_name]
expected_group_name = f"model_name_{team_id}_{public_name}"
# Both sibling deployments get the same deterministic group name
assert internal_name_1 == expected_group_name
assert internal_name_2 == expected_group_name
assert internal_name_1 == internal_name_2, (
"Sibling deployments must share the same model_name so the "
"router treats them as a single candidate pool"
)
# The second alias write is idempotent — same key, same value
final_aliases = {}
for alias_call in captured_alias_calls:
final_aliases.update(alias_call)
assert final_aliases == {public_name: expected_group_name}
class TestTeamModelUpdate:
"""Test team model update handles team_id consistently with model creation"""
@ -657,27 +757,37 @@ class TestModelInfoEndpoint:
user_id="test_user",
api_key="test_key",
models=["gpt-4", "claude-3"],
team_models=["gpt-3.5-turbo"]
team_models=["gpt-3.5-turbo"],
)
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \
patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, \
patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, \
patch("litellm.proxy.proxy_server.get_complete_model_list") as mock_get_complete_models, \
patch("litellm.get_llm_provider") as mock_get_provider:
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch(
"litellm.proxy.proxy_server.get_key_models"
) as mock_get_key_models, patch(
"litellm.proxy.proxy_server.get_team_models"
) as mock_get_team_models, patch(
"litellm.proxy.proxy_server.get_complete_model_list"
) as mock_get_complete_models, patch(
"litellm.get_llm_provider"
) as mock_get_provider:
# Setup mocks
mock_router.get_model_names.return_value = ["gpt-4", "claude-3", "gpt-3.5-turbo"]
mock_router.get_model_names.return_value = [
"gpt-4",
"claude-3",
"gpt-3.5-turbo",
]
mock_router.get_model_access_groups.return_value = {}
mock_get_key_models.return_value = ["gpt-4", "claude-3"]
mock_get_team_models.return_value = ["gpt-3.5-turbo"]
mock_get_complete_models.return_value = ["gpt-4", "claude-3", "gpt-3.5-turbo"]
mock_get_complete_models.return_value = [
"gpt-4",
"claude-3",
"gpt-3.5-turbo",
]
mock_get_provider.return_value = (None, "openai", None, None)
# Test accessible model
result = await model_info(
model_id="gpt-4",
user_api_key_dict=user_api_key_dict
model_id="gpt-4", user_api_key_dict=user_api_key_dict
)
assert result["id"] == "gpt-4"
@ -688,22 +798,25 @@ class TestModelInfoEndpoint:
@pytest.mark.asyncio
async def test_model_info_inaccessible_model_returns_404(self):
"""Test model_info returns 404 for inaccessible models"""
from litellm.proxy.proxy_server import model_info
from fastapi import HTTPException
from litellm.proxy.proxy_server import model_info
# Mock user with limited access
user_api_key_dict = UserAPIKeyAuth(
user_id="test_user",
api_key="test_key",
models=["gpt-4"], # Only has access to gpt-4
team_models=[]
team_models=[],
)
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \
patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, \
patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, \
patch("litellm.proxy.proxy_server.get_complete_model_list") as mock_get_complete_models:
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch(
"litellm.proxy.proxy_server.get_key_models"
) as mock_get_key_models, patch(
"litellm.proxy.proxy_server.get_team_models"
) as mock_get_team_models, patch(
"litellm.proxy.proxy_server.get_complete_model_list"
) as mock_get_complete_models:
# Setup mocks - user only has access to gpt-4
mock_router.get_model_names.return_value = ["gpt-4", "claude-3"]
mock_router.get_model_access_groups.return_value = {}
@ -715,32 +828,35 @@ class TestModelInfoEndpoint:
with pytest.raises(HTTPException) as exc_info:
await model_info(
model_id="claude-3", # Not in user's accessible models
user_api_key_dict=user_api_key_dict
user_api_key_dict=user_api_key_dict,
)
assert exc_info.value.status_code == 404
assert "does not exist or is not accessible" in exc_info.value.detail
@pytest.mark.asyncio
@pytest.mark.asyncio
async def test_model_info_team_model_access(self):
"""Test model_info works with team model access"""
from litellm.proxy.proxy_server import model_info
# Mock user with team access
user_api_key_dict = UserAPIKeyAuth(
user_id="test_user",
api_key="test_key",
api_key="test_key",
team_id="test_team",
models=[], # No direct key models
team_models=["team-model-1"]
team_models=["team-model-1"],
)
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, \
patch("litellm.proxy.proxy_server.get_key_models") as mock_get_key_models, \
patch("litellm.proxy.proxy_server.get_team_models") as mock_get_team_models, \
patch("litellm.proxy.proxy_server.get_complete_model_list") as mock_get_complete_models, \
patch("litellm.get_llm_provider") as mock_get_provider:
with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch(
"litellm.proxy.proxy_server.get_key_models"
) as mock_get_key_models, patch(
"litellm.proxy.proxy_server.get_team_models"
) as mock_get_team_models, patch(
"litellm.proxy.proxy_server.get_complete_model_list"
) as mock_get_complete_models, patch(
"litellm.get_llm_provider"
) as mock_get_provider:
# Setup mocks
mock_router.get_model_names.return_value = ["team-model-1"]
mock_router.get_model_access_groups.return_value = {}
@ -751,10 +867,9 @@ class TestModelInfoEndpoint:
# Test team model access
result = await model_info(
model_id="team-model-1",
user_api_key_dict=user_api_key_dict
model_id="team-model-1", user_api_key_dict=user_api_key_dict
)
assert result["id"] == "team-model-1"
assert result["object"] == "model"
assert result["object"] == "model"
assert result["owned_by"] == "custom"