feat(team_endpoints): add query parameter key_limit to /team/info endpoint (#30006)

* feat(team_endpoints): Add query parameter key_limit to /team/info

* feat(team_endpoints): update schema.d.ts to include the new query parameter

* feat(team_endpoints): add tests for limitting key count in /team/info response

* feat(team_endpoints): Apply suggestions from greptile

* Set greater-than constraint on key-limit
* Fix type

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:

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},
                     )

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

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},
+    }

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;