mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +00:00
(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:
parent
73df319f4e
commit
b41311bb21
2 changed files with 125 additions and 7 deletions
|
@ -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 = {}
|
||||
|
||||
|
|
|
@ -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):
|
||||
"""
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue