diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 4d4d1ef2774..993a271fb9e 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -3577,6 +3577,9 @@ async def team_info( team_id: str = fastapi.Query( default=None, description="Team ID in the request parameters" ), + key_limit: Optional[int] = fastapi.Query( + default=None, description="Limit the number of keys returned", gt=0 + ), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -3632,6 +3635,7 @@ async def team_info( table_name="key", query_type="find_all", expires=datetime.now(), + limit=key_limit, ) if keys is None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0423ba7d3f5..cfbabfc3626 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3431,6 +3431,7 @@ class PrismaClient: r.expires = r.expires.isoformat() elif query_type == "find_all" and team_id is not None: response = await VerificationTokenRepository(self).table.find_many( + take=limit, where={"team_id": team_id}, include={"litellm_budget_table": True}, ) 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 b81807ee19e..a649bc7225e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -9313,3 +9313,34 @@ async def test_clear_team_member_budget_fields_no_budget_row_skips_update(): mock_update_budget.assert_not_awaited() assert "team_member_budget" not in result assert "team_member_rpm_limit" not in result + + +@pytest.mark.asyncio +async def test_team_info_forwards_key_limit_to_get_data(): + """/team/info must thread its ``key_limit`` query param into the key + lookup so the database caps how many keys are returned for the team. + """ + from fastapi import Request + + from litellm.proxy.management_endpoints import team_endpoints + + mock_prisma = MagicMock() + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=LiteLLM_TeamTable(team_id="team-1") + ) + mock_prisma.get_data = AsyncMock(return_value=[]) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + team_endpoints, "get_all_team_memberships", AsyncMock(return_value=[]) + ), + ): + await team_endpoints.team_info( + http_request=MagicMock(spec=Request), + team_id="team-1", + key_limit=7, + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert mock_prisma.get_data.await_args.kwargs["limit"] == 7 diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index 08d1ef619a7..d517c1c346f 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -514,3 +514,26 @@ async def test_get_data_combined_view_returns_view_for_deprecated_key( assert isinstance(response, LiteLLM_VerificationTokenView) assert response.token == active_hash + + +@pytest.mark.asyncio +@pytest.mark.parametrize("limit", [5, None]) +async def test_get_data_team_keys_forward_limit_as_take( + prisma_client: PrismaClient, limit: Any +) -> None: + """The /team/info ``key_limit`` must reach Prisma as ``take`` so the + database caps how many of a team's keys come back. + ``limit=None`` leaves ``take`` unset so every key is returned. + """ + prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + await prisma_client.get_data( + team_id="team-1", + table_name="key", + query_type="find_all", + limit=limit, + ) + assert prisma_client.db.litellm_verificationtoken.find_many.await_args.kwargs == { + "take": limit, + "where": {"team_id": "team-1"}, + "include": {"litellm_budget_table": True}, + } diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 100b7523830..4eb2cfea410 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -48838,6 +48838,8 @@ export interface operations { query?: { /** @description Team ID in the request parameters */ team_id?: string; + /** @description Limit the number of keys returned */ + key_limit?: number | null; }; header?: never; path?: never;