(UI) - Fix show correct count of internal user keys on Users Page (#9082)

* get_user_key_counts

* fix get_user_key_counts

* fix get_user_key_counts

* test_get_users_filters_dashboard_keys

* remove unused func
This commit is contained in:
Ishaan Jaff 2025-03-08 16:13:18 -08:00 committed by GitHub
parent 73df319f4e
commit b41311bb21
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 125 additions and 7 deletions

View file

@ -739,6 +739,43 @@ async def user_update(
)
async def get_user_key_counts(
prisma_client,
user_ids: Optional[List[str]] = None,
):
"""
Helper function to get the count of keys for each user using Prisma's count method.
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
"""
from litellm.constants import UI_SESSION_TOKEN_TEAM_ID
if not user_ids or len(user_ids) == 0:
return {}
result = {}
# Get count for each user_id individually
for user_id in user_ids:
count = await prisma_client.db.litellm_verificationtoken.count(
where={
"user_id": user_id,
"OR": [
{"team_id": None},
{"team_id": {"not": UI_SESSION_TOKEN_TEAM_ID}},
],
}
)
result[user_id] = count
return result
@router.get(
"/user/get_users",
tags=["Internal User management"],
@ -830,14 +867,9 @@ async def get_users(
# Get key count for each user
if users is not None:
user_keys = await prisma_client.db.litellm_verificationtoken.group_by(
by=["user_id"],
count={"user_id": True},
where={"user_id": {"in": [user.user_id for user in users]}},
user_key_counts = await get_user_key_counts(
prisma_client, [user.user_id for user in users]
)
user_key_counts = {
item["user_id"]: item["_count"]["user_id"] for item in user_keys
}
else:
user_key_counts = {}

View file

@ -378,6 +378,92 @@ async def test_get_users(prisma_client):
await prisma_client.db.litellm_usertable.delete(where={"user_id": user.user_id})
@pytest.mark.asyncio
async def test_get_users_filters_dashboard_keys(prisma_client):
"""
Tests that /users/list endpoint doesn't return keys with team_id='litellm-dashboard'
The dashboard keys should be filtered out from the response
"""
litellm.set_verbose = True
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
# Create a test user
new_user_id = f"test_user_with_keys-{uuid.uuid4()}"
test_user = NewUserRequest(
user_id=new_user_id,
user_role=LitellmUserRoles.INTERNAL_USER.value,
auto_create_key=False,
)
await new_user(
test_user,
UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN,
api_key="sk-1234",
user_id="admin",
),
)
# Create two keys for the user - one with team_id="litellm-dashboard" and one without
regular_key = await generate_key_helper_fn(
user_id=test_user.user_id,
request_type="key",
team_id="litellm-dashboard", # This key should be included in the response
models=[],
aliases={},
config={},
spend=0,
duration=None,
)
regular_key = await generate_key_helper_fn(
user_id=test_user.user_id,
request_type="key",
team_id="NEW_TEAM", # This key should be included in the response
models=[],
aliases={},
config={},
spend=0,
duration=None,
)
regular_key = await generate_key_helper_fn(
user_id=test_user.user_id,
request_type="key",
team_id=None, # This key should be included in the response
models=[],
aliases={},
config={},
spend=0,
duration=None,
)
# Test get_users for the specific user
result = await get_users(
user_ids=test_user.user_id,
role=None,
page=1,
page_size=20,
)
print("get users result", result)
assert "users" in result
assert len(result["users"]) == 1
# Verify the key count is correct (should be 1, not counting dashboard keys)
user = result["users"][0]
assert user.user_id == test_user.user_id
assert user.key_count == 2 # Only count the regular keys, not the UI dashboard key
# Clean up test user and keys
await prisma_client.db.litellm_usertable.delete(
where={"user_id": test_user.user_id}
)
@pytest.mark.asyncio
async def test_get_users_key_count(prisma_client):
"""