From 173695f5e0ed6e2e8933fbec341cfdc6162843dc Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 23 Mar 2026 18:26:43 +0530 Subject: [PATCH] Fix greptile comments --- litellm/proxy/litellm_pre_call_utils.py | 12 +++- .../model_management_endpoints.py | 20 +++++- tests/proxy_unit_tests/test_proxy_utils.py | 43 ++++++++++++ .../test_model_management_endpoints.py | 70 ++++++++++++++++++- 4 files changed, 140 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 96a271dd02..4a12a0a577 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -26,6 +26,7 @@ _SPECIAL_HEADERS_CACHE = frozenset( v.value.lower() for v in SpecialHeaders._member_map_.values() ) from litellm.router import Router +from litellm.secret_managers.main import get_secret_bool from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes from litellm.types.utils import ( @@ -1313,9 +1314,16 @@ def _update_model_if_team_alias_exists( # (team models use team_public_model_name, not model_aliases) aliased_target = user_api_key_dict.team_model_aliases[_model] - # Check if the alias points to a stale team-scoped UUID name + # Optional bypass for stale aliases from pre-PR deployments: + # only enabled via feature flag to preserve backwards compatibility. + enable_stale_alias_bypass = get_secret_bool( + "LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", False + ) + # Check if the alias points to a team-scoped UUID name # (format: "model_name_{team_id}_{uuid}") - if aliased_target.startswith(f"model_name_{user_api_key_dict.team_id}_"): + if enable_stale_alias_bypass and aliased_target.startswith( + f"model_name_{user_api_key_dict.team_id}_" + ): # This is a stale alias from pre-PR deployments. # Check if current team deployments exist for the public name. if llm_router: diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 7d0181a368..40f4d722dc 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -469,6 +469,23 @@ async def _update_existing_team_model_assignment( prisma_client: Optional[PrismaClient], ) -> None: """Update an existing team model if the public name changed.""" + + def _get_team_public_model_name( + model_info: Optional[Union[dict, str]] + ) -> Optional[str]: + if isinstance(model_info, dict): + value = model_info.get("team_public_model_name") + return value if isinstance(value, str) else None + if isinstance(model_info, str): + try: + parsed = json.loads(model_info) + except (TypeError, ValueError): + return None + if isinstance(parsed, dict): + value = parsed.get("team_public_model_name") + return value if isinstance(value, str) else None + return None + old_public_name = ( db_model.model_info.team_public_model_name if db_model.model_info else None ) @@ -495,8 +512,7 @@ async def _update_existing_team_model_assignment( d for d in response if d.model_name != db_model.model_name - and (d.model_info or {}).get("team_public_model_name") - == old_public_name + and _get_team_public_model_name(d.model_info) == old_public_name ] if not other_deployments_with_old_name: diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 00d4cd24e4..5e75890388 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -2044,6 +2044,49 @@ 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(): + from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists + + class _MockRouter: + team_model_to_deployment_indices = {("team-1", "gpt-4o"): [0]} + + test_data = {"model": "gpt-4o"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_key", + team_id="team-1", + team_model_aliases={"gpt-4o": "model_name_team-1_legacy-uuid"}, + ) + + with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): + _update_model_if_team_alias_exists( + data=test_data, user_api_key_dict=user_api_key_dict + ) + + assert test_data.get("model") == "model_name_team-1_legacy-uuid" + + +def test_team_alias_stale_bypass_enabled_by_flag(monkeypatch): + from litellm.proxy.litellm_pre_call_utils import _update_model_if_team_alias_exists + + class _MockRouter: + team_model_to_deployment_indices = {("team-1", "gpt-4o"): [0]} + + test_data = {"model": "gpt-4o"} + user_api_key_dict = UserAPIKeyAuth( + api_key="test_key", + team_id="team-1", + team_model_aliases={"gpt-4o": "model_name_team-1_legacy-uuid"}, + ) + monkeypatch.setenv("LITELLM_ENABLE_TEAM_STALE_ALIAS_BYPASS", "true") + + with patch("litellm.proxy.proxy_server.llm_router", _MockRouter()): + _update_model_if_team_alias_exists( + data=test_data, user_api_key_dict=user_api_key_dict + ) + + assert test_data.get("model") == "gpt-4o" + + @pytest.fixture def mock_prisma_client(): client = MagicMock() diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 09410c19d3..83e6b0c93a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -70,10 +70,23 @@ class MockPrismaClient: # Filter deployments by team_id if specified if team_id_filter: + + def _get_team_id(model_info): + if isinstance(model_info, dict): + return model_info.get("team_id") + if isinstance(model_info, str): + try: + parsed = json.loads(model_info) + except (TypeError, ValueError): + return None + if isinstance(parsed, dict): + return parsed.get("team_id") + return None + return [ d for d in self.sibling_deployments - if d.model_info.get("team_id") == team_id_filter + if _get_team_id(d.model_info) == team_id_filter ] return self.sibling_deployments @@ -831,6 +844,61 @@ class TestTeamModelUpdate: # team_model_add should be called to add new public name mock_add.assert_called_once() + @pytest.mark.asyncio + async def test_rename_handles_legacy_string_model_info(self): + """Test rename path handles legacy string-encoded model_info rows without crashing.""" + from unittest.mock import MagicMock + + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_existing_team_model_assignment, + ) + from litellm.types.router import ModelInfo + + db_model = Deployment( + model_name="model_name_team_123_uuid1", + litellm_params=LiteLLM_Params(model="azure/gpt-4o-mini"), + model_info=ModelInfo( + team_id="team_123", team_public_model_name="old-public-name" + ), + ) + + sibling_deployment = MagicMock() + sibling_deployment.model_name = "model_name_team_123_uuid2" + sibling_deployment.model_info = ( + '{"team_id":"team_123","team_public_model_name":"old-public-name"}' + ) + + prisma_client = MockPrismaClient( + team_exists=True, sibling_deployments=[sibling_deployment] + ) + + patch_data = updateDeployment( + model_name="new-public-name", + model_info=ModelInfo(team_id="team_123"), + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + + with patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_delete" + ) as mock_delete, patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_add: + await _update_existing_team_model_assignment( + team_id="team_123", + public_model_name="new-public-name", + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + + mock_delete.assert_not_called() + mock_add.assert_called_once() + @pytest.mark.asyncio async def test_patch_model_with_team_id_validates_permissions(self): """Test PATCH with team_id runs same validation as POST for team permissions"""