From cec7e4941afb6fc5e16f3958fd7722cb1617842e Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 28 Jun 2025 22:04:49 -0700 Subject: [PATCH] UI QA Fixes - prevent team model reset on model add + return team-only models on /v2/model/info + render team member budget correctly (#12144) * fix(team_endpoints.py): prevent overwriting current list of team models on new model add * fix(networking.tsx): fix default proxy base url * fix(proxy_server.py): include team only models when retrieving all deployments on `/v2/model/info` helper util ensures team only models are shown to user * fix(router.py): check model name by team public model name when team id given Fixes issue where team member could not see team only models when clicking into that team on `Models + Endpoints` * fix(team_member_view.tsx): fix rendering team member budget, when budget is set * test: update tests * test: update unit test --- .../management_endpoints/team_endpoints.py | 40 +++--- litellm/proxy/proxy_server.py | 12 +- litellm/router.py | 38 +++++- .../test_team_endpoints.py | 94 +++++++++----- tests/test_litellm/proxy/test_proxy_server.py | 16 +-- tests/test_litellm/test_router.py | 121 ++++++++++++++++++ .../src/components/networking.tsx | 2 +- .../src/components/team/team_member_view.tsx | 6 +- 8 files changed, 261 insertions(+), 68 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index a1935255ef..3f5431c3ef 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -947,6 +947,7 @@ def team_member_add_duplication_check( This check is done BEFORE we create/fetch the user, so it only prevents obvious duplicates where both user_id and user_email match exactly. """ + def _check_member_duplication(member: Member): # Check by user_id if provided if member.user_id is not None: @@ -958,7 +959,7 @@ def team_member_add_duplication_check( param="user_id", code="400", ) - + # Check by user_email if provided if member.user_email is not None: for existing_member in existing_team_row.members_with_roles: @@ -1014,13 +1015,13 @@ async def _process_team_members( """Process and add new team members.""" updated_users: List[LiteLLM_UserTable] = [] updated_team_memberships: List[LiteLLM_TeamMembership] = [] - + default_team_budget_id = ( complete_team_data.metadata.get("team_member_budget_id") if complete_team_data.metadata is not None else None ) - + if isinstance(data.member, Member): try: updated_user, updated_tm = await add_new_member( @@ -1068,7 +1069,7 @@ async def _process_team_members( updated_users.append(updated_user) if updated_tm is not None: updated_team_memberships.append(updated_tm) - + return updated_users, updated_team_memberships @@ -1080,7 +1081,7 @@ async def _update_team_members_list( """Update the team's members_with_roles list.""" if isinstance(data.member, Member): new_member = data.member.model_copy() - + # get user id if new_member.user_id is None and new_member.user_email is not None: for user in updated_users: @@ -1089,33 +1090,42 @@ async def _update_team_members_list( and user.user_email == new_member.user_email ): new_member.user_id = user.user_id - + # Check if member already exists in team before adding member_already_exists = False for existing_member in complete_team_data.members_with_roles: - if (new_member.user_id is not None and existing_member.user_id == new_member.user_id) or \ - (new_member.user_email is not None and existing_member.user_email == new_member.user_email): + if ( + new_member.user_id is not None + and existing_member.user_id == new_member.user_id + ) or ( + new_member.user_email is not None + and existing_member.user_email == new_member.user_email + ): member_already_exists = True break - + if not member_already_exists: complete_team_data.members_with_roles.append(new_member) - + elif isinstance(data.member, List): for nm in data.member: if nm.user_id is None and nm.user_email is not None: for user in updated_users: if user.user_email is not None and user.user_email == nm.user_email: nm.user_id = user.user_id - + # Check if member already exists in team before adding member_already_exists = False for existing_member in complete_team_data.members_with_roles: - if (nm.user_id is not None and existing_member.user_id == nm.user_id) or \ - (nm.user_email is not None and existing_member.user_email == nm.user_email): + if ( + nm.user_id is not None and existing_member.user_id == nm.user_id + ) or ( + nm.user_email is not None + and existing_member.user_email == nm.user_email + ): member_already_exists = True break - + if not member_already_exists: complete_team_data.members_with_roles.append(nm) @@ -2381,7 +2391,7 @@ def add_new_models_to_team( ): # implies all model access current_models = [SpecialModelNames.all_proxy_models.value] else: - current_models = [] + current_models = team_obj.models updated_models = list(set(current_models + new_models)) return updated_models diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4234633c80..a30112bcbd 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5608,7 +5608,9 @@ def _add_team_models_to_all_models( team_models.setdefault(model_id, set()).add(team_object.team_id) else: for model_name in team_object.models: - _models = llm_router.get_model_list(model_name=model_name) + _models = llm_router.get_model_list( + model_name=model_name, team_id=team_object.team_id + ) if _models is not None: for model in _models: model_id = model.get("model_info", {}).get("id", None) @@ -5644,6 +5646,7 @@ async def get_all_team_models( team_db_objects = await prisma_client.db.litellm_teamtable.find_many( where={"team_id": {"in": user_teams}} ) + team_db_objects_typed = [ LiteLLM_TeamTable(**team_db_object.model_dump()) for team_db_object in team_db_objects @@ -5714,6 +5717,7 @@ async def get_all_team_and_direct_access_models( prisma_client=prisma_client, llm_router=llm_router, ) + for _model in all_models: model_id = _model.get("model_info", {}).get("id", None) team_only_model_id = _model.get("model_info", {}).get("team_id", None) @@ -5729,9 +5733,11 @@ async def get_all_team_and_direct_access_models( ) ## ADD DIRECT_ACCESS TO RELEVANT MODELS + for _model in all_models: model_id = _model.get("model_info", {}).get("id", None) if model_id is not None and model_id in direct_access_models: + _model["model_info"]["direct_access"] = True ## FILTER OUT MODELS THAT ARE NOT IN DIRECT_ACCESS_MODELS OR ACCESS_VIA_TEAM_IDS - only show user models they can call @@ -7523,6 +7529,7 @@ async def new_invitation( from litellm.proxy.management_helpers.user_invitation import ( create_invitation_for_user, ) + global prisma_client if prisma_client is None: @@ -7541,7 +7548,7 @@ async def new_invitation( ) }, ) - + response = await create_invitation_for_user( data=data, user_api_key_dict=user_api_key_dict, @@ -7551,7 +7558,6 @@ async def new_invitation( raise handle_exception_on_proxy(e) - @router.get( "/invitation/info", tags=["Invite Links"], diff --git a/litellm/router.py b/litellm/router.py index 77f99ea84b..dc12b924ff 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3643,7 +3643,10 @@ class Router: litellm.ContentPolicyViolationError: when `mock_testing_content_policy_fallbacks=True` passed in request params """ mock_testing_params = MockRouterTestingParams.from_kwargs(kwargs) - if mock_testing_params.mock_testing_fallbacks is not None and mock_testing_params.mock_testing_fallbacks is True: + if ( + mock_testing_params.mock_testing_fallbacks is not None + and mock_testing_params.mock_testing_fallbacks is True + ): raise litellm.InternalServerError( model=model_group, llm_provider="", @@ -5584,8 +5587,27 @@ class Router: return model["model_name"] return None + def should_include_deployment( + self, model_name: str, model: dict, team_id: Optional[str] = None + ) -> bool: + """ + Get the team-specific model name if team_id matches the deployment. + """ + if ( + team_id is not None + and model["model_info"].get("team_id") == team_id + and model_name == model["model_info"].get("team_public_model_name") + ): + return True + elif model_name is not None and model["model_name"] == model_name: + return True + return False + def _get_all_deployments( - self, model_name: str, model_alias: Optional[str] = None + self, + model_name: str, + model_alias: Optional[str] = None, + team_id: Optional[str] = None, ) -> List[DeploymentTypedDict]: """ Return all deployments of a model name @@ -5594,7 +5616,9 @@ class Router: """ returned_models: List[DeploymentTypedDict] = [] for model in self.model_list: - if model_name is not None and model["model_name"] == model_name: + if self.should_include_deployment( + model_name=model_name, model=model, team_id=team_id + ): if model_alias is not None: alias_model = copy.deepcopy(model) alias_model["model_name"] = model_alias @@ -5692,16 +5716,20 @@ class Router: return returned_models def get_model_list( - self, model_name: Optional[str] = None + self, model_name: Optional[str] = None, team_id: Optional[str] = None ) -> Optional[List[DeploymentTypedDict]]: """ Includes router model_group_alias'es as well + + if team_id specified, returns matching team-specific models """ if hasattr(self, "model_list"): returned_models: List[DeploymentTypedDict] = [] if model_name is not None: - returned_models.extend(self._get_all_deployments(model_name=model_name)) + returned_models.extend( + self._get_all_deployments(model_name=model_name, team_id=team_id) + ) if hasattr(self, "model_group_alias"): returned_models.extend( diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 67e4966f98..ec455a5417 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -751,11 +751,11 @@ async def test_validate_team_member_add_permissions_admin(): # Create admin user admin_user = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN.value) - + # Create mock team team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = "test-team-123" - + # Should not raise any exception for admin await _validate_team_member_add_permissions( user_api_key_dict=admin_user, @@ -776,21 +776,21 @@ async def test_validate_team_member_add_permissions_non_admin(): regular_user = UserAPIKeyAuth( user_id="regular-user", user_role=LitellmUserRoles.INTERNAL_USER.value, - team_id="different-team" + team_id="different-team", ) - + # Create mock team team = MagicMock(spec=LiteLLM_TeamTable) team.team_id = "test-team-123" team.members_with_roles = [] - + # Mock the helper functions to return False with patch( "litellm.proxy.management_endpoints.team_endpoints._is_user_team_admin", - return_value=False + return_value=False, ), patch( "litellm.proxy.management_endpoints.team_endpoints._is_available_team", - return_value=False + return_value=False, ): # Should raise HTTPException for non-admin with pytest.raises(HTTPException) as exc_info: @@ -798,7 +798,7 @@ async def test_validate_team_member_add_permissions_non_admin(): user_api_key_dict=regular_user, complete_team_data=team, ) - + assert exc_info.value.status_code == 403 assert "not proxy admin OR team admin" in str(exc_info.value.detail) @@ -808,30 +808,30 @@ async def test_process_team_members_single_member(): """ Test _process_team_members with a single member """ + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_UserTable from litellm.proxy.management_endpoints.team_endpoints import _process_team_members - from litellm.proxy._types import LiteLLM_UserTable, LiteLLM_TeamMembership # Mock dependencies mock_prisma_client = MagicMock() mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.metadata = {"team_member_budget_id": "budget-123"} - + # Mock user and membership objects mock_user = MagicMock(spec=LiteLLM_UserTable) mock_user.user_id = "new-user-123" mock_membership = MagicMock(spec=LiteLLM_TeamMembership) - + # Create request with single member single_member = Member(user_email="new@example.com", role="user") request_data = TeamMemberAddRequest( team_id="test-team-123", member=single_member, ) - + with patch( "litellm.proxy.management_endpoints.team_endpoints.add_new_member", new_callable=AsyncMock, - return_value=(mock_user, mock_membership) + return_value=(mock_user, mock_membership), ) as mock_add_member: users, memberships = await _process_team_members( data=request_data, @@ -840,13 +840,13 @@ async def test_process_team_members_single_member(): user_api_key_dict=UserAPIKeyAuth(), litellm_proxy_admin_name="admin", ) - + # Verify results assert len(users) == 1 assert len(memberships) == 1 assert users[0] == mock_user assert memberships[0] == mock_membership - + # Verify add_new_member was called correctly mock_add_member.assert_called_once_with( new_member=single_member, @@ -864,14 +864,14 @@ async def test_process_team_members_multiple_members(): """ Test _process_team_members with multiple members """ + from litellm.proxy._types import LiteLLM_TeamMembership, LiteLLM_UserTable from litellm.proxy.management_endpoints.team_endpoints import _process_team_members - from litellm.proxy._types import LiteLLM_UserTable, LiteLLM_TeamMembership # Mock dependencies mock_prisma_client = MagicMock() mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.metadata = None - + # Create multiple members as dictionaries (they will be converted to Member objects) members = [ {"user_email": "user1@example.com", "role": "user"}, @@ -882,15 +882,18 @@ async def test_process_team_members_multiple_members(): member=members, max_budget_in_team=100.0, ) - + # Mock different users and memberships for each call mock_users = [MagicMock(spec=LiteLLM_UserTable) for _ in range(2)] mock_memberships = [MagicMock(spec=LiteLLM_TeamMembership) for _ in range(2)] - + with patch( "litellm.proxy.management_endpoints.team_endpoints.add_new_member", new_callable=AsyncMock, - side_effect=[(mock_users[0], mock_memberships[0]), (mock_users[1], mock_memberships[1])] + side_effect=[ + (mock_users[0], mock_memberships[0]), + (mock_users[1], mock_memberships[1]), + ], ) as mock_add_member: users, memberships = await _process_team_members( data=request_data, @@ -899,13 +902,13 @@ async def test_process_team_members_multiple_members(): user_api_key_dict=UserAPIKeyAuth(), litellm_proxy_admin_name="admin", ) - + # Verify results assert len(users) == 2 assert len(memberships) == 2 assert users == mock_users assert memberships == mock_memberships - + # Verify add_new_member was called for each member assert mock_add_member.call_count == 2 @@ -915,33 +918,33 @@ async def test_update_team_members_list_single_member(): """ Test _update_team_members_list with a single member """ - from litellm.proxy.management_endpoints.team_endpoints import _update_team_members_list from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) # Create mock team with existing members mock_team = MagicMock(spec=LiteLLM_TeamTable) - mock_team.members_with_roles = [ - Member(user_id="existing-user", role="admin") - ] - + mock_team.members_with_roles = [Member(user_id="existing-user", role="admin")] + # Create new member without user_id new_member = Member(user_email="new@example.com", role="user") request_data = TeamMemberAddRequest( team_id="test-team-123", member=new_member, ) - + # Create mock user with matching email mock_user = MagicMock(spec=LiteLLM_UserTable) mock_user.user_id = "new-user-123" mock_user.user_email = "new@example.com" - + await _update_team_members_list( data=request_data, complete_team_data=mock_team, updated_users=[mock_user], ) - + # Verify member was added assert len(mock_team.members_with_roles) == 2 added_member = mock_team.members_with_roles[1] @@ -955,32 +958,53 @@ async def test_update_team_members_list_duplicate_prevention(): """ Test _update_team_members_list prevents duplicate members """ - from litellm.proxy.management_endpoints.team_endpoints import _update_team_members_list from litellm.proxy._types import LiteLLM_UserTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _update_team_members_list, + ) # Create mock team with existing members mock_team = MagicMock(spec=LiteLLM_TeamTable) mock_team.members_with_roles = [ Member(user_id="existing-user", user_email="existing@example.com", role="admin") ] - + # Try to add the same member again duplicate_member = Member(user_id="existing-user", role="user") request_data = TeamMemberAddRequest( team_id="test-team-123", member=duplicate_member, ) - + # Create mock user mock_user = MagicMock(spec=LiteLLM_UserTable) mock_user.user_id = "existing-user" mock_user.user_email = "existing@example.com" - + await _update_team_members_list( data=request_data, complete_team_data=mock_team, updated_users=[mock_user], ) - + # Verify member was NOT added (still only 1 member) assert len(mock_team.members_with_roles) == 1 + + +def test_add_new_models_to_team_with_existing_models(): + """ + Test add_new_models_to_team function with existing models + """ + from litellm.proxy._types import SpecialModelNames + from litellm.proxy.management_endpoints.team_endpoints import add_new_models_to_team + + team_obj = MagicMock(spec=LiteLLM_TeamTable) + team_obj.models = ["model1", "model2"] + new_models = ["model3", "model4"] + + updated_models = add_new_models_to_team( + team_obj=team_obj, + new_models=new_models, + ) + + assert updated_models.sort() == ["model1", "model2", "model3", "model4"].sort() diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8688d08ee4..056228ac0a 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -320,7 +320,7 @@ async def test_get_all_team_models(): # Mock router mock_router = MagicMock() - def mock_get_model_list(model_name): + def mock_get_model_list(model_name, team_id=None): if model_name == "gpt-4": return mock_models_gpt4 elif model_name == "gpt-3.5-turbo": @@ -355,10 +355,10 @@ async def test_get_all_team_models(): # Verify router.get_model_list was called for each model expected_calls = [ - mock.call(model_name="gpt-4"), - mock.call(model_name="gpt-3.5-turbo"), - mock.call(model_name="claude-3"), - mock.call(model_name="gpt-4"), # Called again for team2 + mock.call(model_name="gpt-4", team_id="team1"), + mock.call(model_name="gpt-3.5-turbo", team_id="team1"), + mock.call(model_name="claude-3", team_id="team2"), + mock.call(model_name="gpt-4", team_id="team2"), ] mock_router.get_model_list.assert_has_calls(expected_calls, any_order=True) @@ -386,8 +386,8 @@ async def test_get_all_team_models(): # Verify router.get_model_list was called only for team1 models expected_calls = [ - mock.call(model_name="gpt-4"), - mock.call(model_name="gpt-3.5-turbo"), + mock.call(model_name="gpt-4", team_id="team1"), + mock.call(model_name="gpt-3.5-turbo", team_id="team1"), ] mock_router.get_model_list.assert_has_calls(expected_calls, any_order=True) @@ -413,7 +413,7 @@ async def test_get_all_team_models(): mock_router.reset_mock() mock_litellm_teamtable.find_many.return_value = [mock_team1] - def mock_get_model_list_with_none(model_name): + def mock_get_model_list_with_none(model_name, team_id=None): if model_name == "gpt-4": return mock_models_gpt4 # Return None for gpt-3.5-turbo to test None handling diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6752d79224..ee637e810b 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -481,3 +481,124 @@ async def test_router_filter_team_based_models(): ) assert result is not None + + +def test_router_should_include_deployment(): + """ + Test the should_include_deployment method with various scenarios + + The method logic: + 1. Returns True if: team_id matches AND model_name matches team_public_model_name + 2. Returns True if: model_name matches AND deployment has no team_id + 3. Otherwise returns False + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + "model_info": { + "team_id": "test-team", + }, + }, + ], + ) + + # Test deployment structures + deployment_with_team_and_public_name = { + "model_name": "gpt-3.5-turbo", + "model_info": { + "team_id": "test-team", + "team_public_model_name": "team-gpt-model", + }, + } + + deployment_with_team_no_public_name = { + "model_name": "gpt-3.5-turbo", + "model_info": { + "team_id": "test-team", + }, + } + + deployment_without_team = { + "model_name": "gpt-4", + "model_info": {}, + } + + deployment_different_team = { + "model_name": "claude-3", + "model_info": { + "team_id": "other-team", + "team_public_model_name": "team-claude-model", + }, + } + + # Test Case 1: Team-specific deployment - team_id and team_public_model_name match + result = router.should_include_deployment( + model_name="team-gpt-model", + model=deployment_with_team_and_public_name, + team_id="test-team", + ) + assert ( + result is True + ), "Should return True when team_id and team_public_model_name match" + + # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name + result = router.should_include_deployment( + model_name="different-model", + model=deployment_with_team_and_public_name, + team_id="test-team", + ) + assert ( + result is False + ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + + # Test Case 3: Team-specific deployment - team_id doesn't match + result = router.should_include_deployment( + model_name="team-gpt-model", + model=deployment_with_team_and_public_name, + team_id="different-team", + ) + assert result is False, "Should return False when team_id doesn't match" + + # Test Case 4: Team-specific deployment with no team_public_model_name - should fail + result = router.should_include_deployment( + model_name="gpt-3.5-turbo", + model=deployment_with_team_no_public_name, + team_id="test-team", + ) + assert ( + result is True + ), "Should return True when team deployment has no team_public_model_name to match" + + # Test Case 5: Non-team deployment - model_name matches and no team_id + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id=None + ) + assert ( + result is True + ), "Should return True when model_name matches and deployment has no team_id" + + # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) + result = router.should_include_deployment( + model_name="gpt-4", model=deployment_without_team, team_id="any-team" + ) + assert ( + result is True + ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + + # Test Case 7: Non-team deployment - model_name doesn't match + result = router.should_include_deployment( + model_name="different-model", model=deployment_without_team, team_id=None + ) + assert result is False, "Should return False when model_name doesn't match" + + # Test Case 8: Team deployment accessed without matching team_id + result = router.should_include_deployment( + model_name="gpt-3.5-turbo", + model=deployment_with_team_and_public_name, + team_id=None, + ) + assert ( + result is True + ), "Should return True when matching model with exact model_name" diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 085fe5cfe9..21c2ea02ef 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -19,7 +19,7 @@ import { } from "./email_events/types"; const isLocal = process.env.NODE_ENV === "development"; -export const defaultProxyBaseUrl = isLocal ? "http://localhost:43845" : null; +export const defaultProxyBaseUrl = isLocal ? "http://localhost:4000" : null; const defaultServerRootPath = "/"; export let serverRootPath = defaultServerRootPath; export let proxyBaseUrl = defaultProxyBaseUrl; diff --git a/ui/litellm-dashboard/src/components/team/team_member_view.tsx b/ui/litellm-dashboard/src/components/team/team_member_view.tsx index e435411aa7..5699edb3e1 100644 --- a/ui/litellm-dashboard/src/components/team/team_member_view.tsx +++ b/ui/litellm-dashboard/src/components/team/team_member_view.tsx @@ -67,7 +67,11 @@ const TeamMembersComponent: React.FC = ({ if (!userId) return null; const membership = teamData.team_memberships.find(tm => tm.user_id === userId); console.log(`membership=${membership}`); - return formatNumber(membership?.litellm_budget_table?.max_budget || null); + const maxBudget = membership?.litellm_budget_table?.max_budget; + if (maxBudget === null || maxBudget === undefined) { + return null; + } + return formatNumber(maxBudget); }; return (