diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 8f9c6edd2b..9e9301dbbe 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -218,8 +218,14 @@ async def patch_model( prisma_client=prisma_client, premium_user=premium_user, ) - # Create update dictionary only for provided fields - update_data = update_db_model(db_model=db_model, updated_patch=patch_data) + + # Handle team model updates with proper alias management + update_data = await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) # Add metadata about update update_data["updated_by"] = ( @@ -361,6 +367,143 @@ async def _add_team_model_to_db( return model_response +async def _update_team_model_in_db( + db_model: Deployment, + patch_data: updateDeployment, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: PrismaClient, +) -> PrismaCompatibleUpdateDBModel: + """ + Handle team model updates with proper alias management. + + If patch_data contains a team_id: + - Creates unique internal model_name and team alias + - Adds model to team object + - Preserves team_public_model_name for external reference + """ + # Validate team_id if present in patch_data + from litellm.proxy.proxy_server import premium_user + + await ModelManagementAuthChecks.allow_team_model_action( + model_params=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + premium_user=premium_user, + ) + + patch_team_id = patch_data.model_info.team_id if patch_data.model_info else None + + # No team_id in patch, proceed with standard update + if patch_team_id is None: + return update_db_model(db_model=db_model, updated_patch=patch_data) + + # Determine public model name + public_model_name = _get_public_model_name( + patch_data=patch_data, + db_model=db_model, + ) + + # Ensure model_info exists and set team_public_model_name + if patch_data.model_info is None: + from litellm.types.router import ModelInfo + patch_data.model_info = ModelInfo() + patch_data.model_info.team_public_model_name = public_model_name + + # Check if team assignment is new or changed + db_team_id = db_model.model_info.team_id if db_model.model_info else None + is_new_team_assignment = db_team_id != patch_team_id + + if is_new_team_assignment: + await _setup_new_team_model_assignment( + team_id=patch_team_id, + public_model_name=public_model_name, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + ) + else: + await _update_existing_team_model_assignment( + team_id=patch_team_id, + public_model_name=public_model_name, + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + ) + + return update_db_model(db_model=db_model, updated_patch=patch_data) + + +def _get_public_model_name( + patch_data: updateDeployment, + db_model: Deployment, +) -> str: + """Determine the public model name from patch or existing model.""" + if patch_data.model_name: + return patch_data.model_name + + if db_model.model_info and db_model.model_info.team_public_model_name: + return db_model.model_info.team_public_model_name + + return db_model.model_name + + +async def _setup_new_team_model_assignment( + team_id: str, + public_model_name: str, + 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 + + await update_team( + data=UpdateTeamRequest( + team_id=team_id, + model_aliases={public_model_name: unique_model_name}, + ), + user_api_key_dict=user_api_key_dict, + http_request=Request(scope={"type": "http"}), + ) + + await team_model_add( + data=TeamModelAddRequest( + team_id=team_id, + models=[public_model_name], + ), + http_request=Request(scope={"type": "http"}), + user_api_key_dict=user_api_key_dict, + ) + + +async def _update_existing_team_model_assignment( + team_id: str, + public_model_name: str, + db_model: Deployment, + patch_data: updateDeployment, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Update an existing team model if the public name changed.""" + old_public_name = ( + db_model.model_info.team_public_model_name + if db_model.model_info + else None + ) + + # Update alias only if public name changed + if old_public_name and public_model_name != old_public_name: + await update_team( + data=UpdateTeamRequest( + team_id=team_id, + model_aliases={public_model_name: db_model.model_name}, + ), + user_api_key_dict=user_api_key_dict, + http_request=Request(scope={"type": "http"}), + ) + + # Keep existing unique model_name + patch_data.model_name = None + + class ModelManagementAuthChecks: """ Common auth checks for model management endpoints 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 2f736e4e1d..e70bc57e59 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 @@ -453,6 +453,92 @@ class TestClearCache: ) +class TestTeamModelUpdate: + """Test team model update handles team_id consistently with model creation""" + + @pytest.mark.asyncio + async def test_patch_model_with_team_id_creates_proper_setup(self): + """Test PATCH with team_id creates unique model name, alias, and team membership like POST does""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_team_model_in_db, + ) + from litellm.types.router import ModelInfo + + patch_data = updateDeployment( + model_name="tenant-azure-gpt4", + model_info=ModelInfo( + team_id="test_team_123", + base_model="azure/gpt-4", + ), + ) + db_model = Deployment( + model_name="original-model", + litellm_params=LiteLLM_Params(model="test_model"), + model_info=ModelInfo(), + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + prisma_client = MockPrismaClient(team_exists=True) + + with patch( + "litellm.proxy.proxy_server.premium_user", + True, + ), patch( + "litellm.proxy.management_endpoints.model_management_endpoints.update_team" + ) as mock_update_team, patch( + "litellm.proxy.management_endpoints.model_management_endpoints.team_model_add" + ) as mock_team_model_add: + result = await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + + assert result.get("model_name", "").startswith("model_name_test_team_123_") + assert "team_public_model_name" in str(result.get("model_info", "")) + mock_update_team.assert_called_once() + mock_team_model_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""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + _update_team_model_in_db, + ) + from litellm.types.router import ModelInfo + + patch_data = updateDeployment( + model_name="tenant-azure-gpt4", + model_info=ModelInfo(team_id="test_team_123"), + ) + db_model = Deployment( + model_name="original-model", + litellm_params=LiteLLM_Params(model="test_model"), + model_info=ModelInfo(), + ) + user_api_key_dict = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER, + ) + prisma_client = MockPrismaClient(team_exists=True, user_admin=False) + + with patch( + "litellm.proxy.proxy_server.premium_user", + True, + ): + with pytest.raises(Exception) as exc_info: + await _update_team_model_in_db( + db_model=db_model, + patch_data=patch_data, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, # type: ignore + ) + assert "403" in str(exc_info.value) + + class TestModelInfoEndpoint: """Test the model_info endpoint for retrieving individual model information"""