diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fec7150cc2..d66cdd4ad6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -784,9 +784,11 @@ class KeyRequest(LiteLLMPydanticObjectBase): class LiteLLM_ModelTable(LiteLLMPydanticObjectBase): + id: Optional[int] = None model_aliases: Optional[Union[str, dict]] = None # json dump the dict created_by: str updated_by: str + team: Optional["LiteLLM_TeamTable"] = None model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 5a939b3e97..b6a4465cab 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -14,7 +14,7 @@ import asyncio import datetime import json import uuid -from typing import Dict, List, Literal, Optional, Union, cast +from typing import Dict, List, Literal, Optional, Tuple, Union, cast from fastapi import APIRouter, Depends, HTTPException, Request, status from pydantic import BaseModel @@ -508,6 +508,35 @@ async def delete_model( premium_user=premium_user, ) + # delete team model alias + if model_params.model_info.team_id is not None: + removed_model_aliases = await delete_team_model_alias( + public_model_name=model_params.model_name, + prisma_client=prisma_client, + ) + + valid_team_model_aliases = [ + model + for team_id, model in removed_model_aliases + if team_id == model_params.model_info.team_id + ] + + ## UPDATE TEAM TO NOT LIST MODEL ## + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": model_params.model_info.team_id} + ) + if existing_team_row is not None: + existing_team_row.models = [ + model + for model in existing_team_row.models + if model not in valid_team_model_aliases + ] + + await prisma_client.db.litellm_teamtable.update( + where={"team_id": model_params.model_info.team_id}, + data={"models": existing_team_row.models}, + ) + # update DB if store_model_in_db is True: """ @@ -572,6 +601,45 @@ async def delete_model( ) +async def delete_team_model_alias( + public_model_name: str, + prisma_client: PrismaClient, +) -> List[Tuple[str, str]]: + """ + Delete a team model alias + + Iterate through all team model aliases and delete the one that matches the model_id + + Returns: + - List of team id + model alias pairs that were removed + """ + team_model_aliases = await prisma_client.db.litellm_modeltable.find_many( + include={"team": True} + ) + tasks = [] + removed_model_aliases = [] + for team_model_alias in team_model_aliases: + model_aliases = team_model_alias.model_aliases # {"alias": "public model name"} + id = team_model_alias.id + + if public_model_name in model_aliases.values(): + key = list(model_aliases.keys())[ + list(model_aliases.values()).index(public_model_name) + ] + if team_model_alias.team is not None: + removed_model_aliases.append((team_model_alias.team.team_id, key)) + del model_aliases[key] + tasks.append( + prisma_client.db.litellm_modeltable.update( + where={"id": id}, + data={"model_aliases": json.dumps(model_aliases)}, + ) + ) + await asyncio.gather(*tasks) + + return removed_model_aliases + + #### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 @router.post( "/model/new", diff --git a/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py index ebe99892cc..129d0d205a 100644 --- a/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1,7 +1,8 @@ import json import os import sys -from typing import Optional +import uuid +from typing import Dict, Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -11,6 +12,7 @@ sys.path.insert( 0, os.path.abspath("../../../..") ) # Adds the parent directory to the system path from litellm.proxy._types import ( + LiteLLM_ModelTable, LiteLLM_TeamTable, LitellmUserRoles, Member, @@ -230,6 +232,135 @@ class TestModelManagementAuthChecks: assert "403" in str(exc_info.value) +class MockModelTable: + def __init__(self, model_aliases: Dict[str, str], include: Optional[dict] = None): + for alias, model in model_aliases.items(): + setattr(self, alias, model) + self.id = str(uuid.uuid4()) + self.model_aliases = model_aliases + + +class MockPrismaDB: + def __init__(self, model_aliases_list): + self.litellm_modeltable = self + self.model_aliases_list = model_aliases_list + self.update_calls = [] + + async def find_many(self, include=None): + print(f"self.model_aliases_list: {self.model_aliases_list}") + return [LiteLLM_ModelTable(**aliases) for aliases in self.model_aliases_list] + + async def update(self, where, data): + self.update_calls.append({"where": where, "data": data}) + return None + + +class MockPrismaWrapper: + def __init__(self, model_aliases_list): + self.litellm_modeltable = MockPrismaDB(model_aliases_list) + + +class TestDeleteTeamModelAlias: + @pytest.mark.asyncio + async def test_delete_team_model_alias_success(self): + """Test successful deletion of a team model alias""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + delete_team_model_alias, + ) + + # Setup test data + model_aliases_list = [ + { + "id": 1, + "model_aliases": { + "alias1": "public_model_1", + "alias2": "public_model_2", + }, + "updated_by": "test_user", + "created_by": "test_user", + }, + { + "id": 2, + "model_aliases": { + "alias3": "public_model_3", + "alias4": "public_model_1", + }, + "updated_by": "test_user", + "created_by": "test_user", + }, # public_model_1 appears twice + ] + + # Create mock prisma client + mock_prisma = MockPrismaClient(team_exists=True) + mock_prisma.db = MockPrismaWrapper(model_aliases_list) + + # Call the function + await delete_team_model_alias( + public_model_name="public_model_1", prisma_client=mock_prisma + ) + + # Verify results + mock_db = mock_prisma.db.litellm_modeltable + assert ( + len(mock_db.update_calls) == 2 + ) # Should have 2 update calls since public_model_1 appears twice + + # Verify first update + first_update = mock_db.update_calls[0] + assert first_update["where"] == {"id": 1} + assert json.loads(first_update["data"]["model_aliases"]) == { + "alias2": "public_model_2" + } + + # Verify second update + second_update = mock_db.update_calls[1] + assert second_update["where"] == {"id": 2} + assert json.loads(second_update["data"]["model_aliases"]) == { + "alias3": "public_model_3" + } + + @pytest.mark.asyncio + async def test_delete_team_model_alias_no_matches(self): + """Test deletion when no matching model alias exists""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + delete_team_model_alias, + ) + + # Setup test data with no matching model + model_aliases_list = [ + { + "id": 1, + "model_aliases": { + "alias1": "public_model_1", + "alias2": "public_model_2", + }, + "updated_by": "test_user", + "created_by": "test_user", + }, + { + "id": 2, + "model_aliases": { + "alias3": "public_model_3", + "alias4": "public_model_4", + }, + "updated_by": "test_user", + "created_by": "test_user", + }, + ] + + # Create mock prisma client + mock_prisma = MockPrismaClient(team_exists=True) + mock_prisma.db = MockPrismaWrapper(model_aliases_list) + + # Call the function with non-existent model + await delete_team_model_alias( + public_model_name="non_existent_model", prisma_client=mock_prisma + ) + + # Verify no updates were made + mock_db = mock_prisma.db.litellm_modeltable + assert len(mock_db.update_calls) == 0 + class TestClearCache: """ Tests for the clear_cache function in model_management_endpoints.py