Merge eba2685ece into 3746ba58d7

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -1889,34 +1889,42 @@ async def get_user_key_counts(
     user_ids: list[str] | None = None,
 ) -> Mapping[str, int]:
     """
-    Helper function to get the count of keys for each user using Prisma's count method.
+    Get the count of (non-UI-session) keys for each user via a single
+    Prisma ``group_by`` query, avoiding an N+1 per-user ``count`` loop.
 
     Args:
         prisma_client: The Prisma client instance
         user_ids: List of user IDs to get key counts for
 
     Returns:
-        Dictionary mapping user_id to key count
+        Dictionary mapping user_id to key count. Users in ``user_ids`` with no
+        keys are present with a value of ``0``.
     """
     from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
 
     if not user_ids or len(user_ids) == 0:
         return {}
 
-    result: Final[dict[str, int]] = {}
+    result = {user_id: 0 for user_id in user_ids}
 
-    # Get count for each user_id individually
-    for user_id in user_ids:
-        count = await _verification_token_table(prisma_client).count(
-            where={
-                "user_id": user_id,
-                "OR": [
-                    {"team_id": None},
-                    {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
-                ],
-            }
-        )
-        result[user_id] = count
+    grouped_counts = await prisma_client.db.litellm_verificationtoken.group_by(
+        by=["user_id"],
+        where={
+            "user_id": {"in": user_ids},
+            "OR": [
+                {"team_id": None},
+                {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
+            ],
+        },
+        count={"user_id": True},
+    )
+
+    for row in grouped_counts:
+        user_id = row.get("user_id")
+        if user_id is None:
+            continue
+        count_block = row.get("_count") or {}
+        result[user_id] = count_block.get("user_id", 0)
 
     return result
 

tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py

@@ -659,6 +659,56 @@ async def test_get_users_redacts_scim_enterprise_metadata(mocker):
     assert "scim_enterprise" not in (listed.metadata or {})
 
 
+@pytest.mark.asyncio
+async def test_get_user_key_counts_returns_empty_for_no_user_ids(mocker):
+    mock_prisma_client = mocker.MagicMock()
+
+    assert await get_user_key_counts(mock_prisma_client, None) == {}
+    assert await get_user_key_counts(mock_prisma_client, []) == {}
+    assert mock_prisma_client.db.litellm_verificationtoken.group_by.call_count == 0
+
+
+@pytest.mark.asyncio
+async def test_get_user_key_counts_uses_group_by(mocker):
+    from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
+
+    mock_prisma_client = mocker.MagicMock()
+    captured_kwargs = {}
+
+    async def mock_group_by(**kwargs):
+        captured_kwargs.update(kwargs)
+        return [
+            {"user_id": "user-a", "_count": {"user_id": 2}},
+            {"user_id": "user-c", "_count": {"user_id": 1}},
+            {"user_id": "user-d"},
+            {"user_id": None, "_count": {"user_id": 99}},
+        ]
+
+    async def mock_count(*args, **kwargs):
+        raise AssertionError("count() should not be called for user key counts")
+
+    mock_prisma_client.db.litellm_verificationtoken.group_by = mock_group_by
+    mock_prisma_client.db.litellm_verificationtoken.count = mock_count
+
+    result = await get_user_key_counts(
+        mock_prisma_client,
+        ["user-a", "user-b", "user-c", "user-d"],
+    )
+
+    assert result == {"user-a": 2, "user-b": 0, "user-c": 1, "user-d": 0}
+    assert captured_kwargs == {
+        "by": ["user_id"],
+        "where": {
+            "user_id": {"in": ["user-a", "user-b", "user-c", "user-d"]},
+            "OR": [
+                {"team_id": None},
+                {"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
+            ],
+        },
+        "count": {"user_id": True},
+    }
+
+
 def test_validate_sort_params():
     """
     Test that validate_sort_params returns None if sort_by is None