diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 633de86aa6e..8a112188d9a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3523,6 +3523,15 @@ async def list_team( prisma_client, _team_ids, user_id=user_id ) + # Batch-fetch all keys for all teams in a single query instead of + # issuing one query per team (N+1 problem). + _all_keys = await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": _team_ids}} + ) + _keys_by_team: dict[str, list] = {} + for _key in _all_keys: + _keys_by_team.setdefault(_key.team_id, []).append(_key) + returned_responses: List[TeamListResponseObject] = [] for team in filtered_response: _team_memberships: List[LiteLLM_TeamMembership] = [] @@ -3530,10 +3539,7 @@ async def list_team( if tm.team_id == team.team_id: _team_memberships.append(tm) - # add all keys that belong to the team - keys = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": team.team_id} - ) + keys = _keys_by_team.get(team.team_id, []) try: returned_responses.append( 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 4d949cfbe69..1aee1d49658 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -6091,3 +6091,69 @@ async def test_list_available_teams_returns_empty_list_when_none_configured(): assert result == [] litellm.default_internal_user_params = original + + +@pytest.mark.asyncio +async def test_list_team_v1_batches_key_queries(): + """ + Test that list_team fetches all keys in a single batched query + instead of issuing one query per team (N+1). + """ + from unittest.mock import AsyncMock, MagicMock, Mock, patch + + from fastapi import Request + + from litellm.proxy._types import ( + LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LitellmUserRoles, + TeamListResponseObject, + UserAPIKeyAuth, + ) + from litellm.proxy.management_endpoints.team_endpoints import list_team + + mock_request = Mock(spec=Request) + + mock_user_api_key_dict = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + + # Two teams + team1 = LiteLLM_TeamTable(team_id="team-1", team_alias="Team One") + team2 = LiteLLM_TeamTable(team_id="team-2", team_alias="Team Two") + + # Mock keys belonging to different teams + key1 = MagicMock() + key1.team_id = "team-1" + key2 = MagicMock() + key2.team_id = "team-1" + key3 = MagicMock() + key3.team_id = "team-2" + + with patch( + "litellm.proxy.proxy_server.prisma_client" + ) as mock_prisma_client, patch( + "litellm.proxy.management_endpoints.team_endpoints._authorize_and_filter_teams", + new_callable=AsyncMock, + return_value=[team1, team2], + ), patch( + "litellm.proxy.management_endpoints.team_endpoints.get_all_team_memberships", + new_callable=AsyncMock, + return_value=[], + ): + mock_find_many = AsyncMock(return_value=[key1, key2, key3]) + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many + + result = await list_team( + http_request=mock_request, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify keys are correctly distributed + assert len(result) == 2 + # Results are sorted by team_alias + assert result[0].team_id == "team-1" + assert result[0].keys == [key1, key2] + assert result[1].team_id == "team-2" + assert result[1].keys == [key3]