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(
|
@router.get(
|
||||||
"/user/get_users",
|
"/user/get_users",
|
||||||
tags=["Internal User management"],
|
tags=["Internal User management"],
|
||||||
|
@ -830,14 +867,9 @@ async def get_users(
|
||||||
|
|
||||||
# Get key count for each user
|
# Get key count for each user
|
||||||
if users is not None:
|
if users is not None:
|
||||||
user_keys = await prisma_client.db.litellm_verificationtoken.group_by(
|
user_key_counts = await get_user_key_counts(
|
||||||
by=["user_id"],
|
prisma_client, [user.user_id for user in users]
|
||||||
count={"user_id": True},
|
|
||||||
where={"user_id": {"in": [user.user_id for user in users]}},
|
|
||||||
)
|
)
|
||||||
user_key_counts = {
|
|
||||||
item["user_id"]: item["_count"]["user_id"] for item in user_keys
|
|
||||||
}
|
|
||||||
else:
|
else:
|
||||||
user_key_counts = {}
|
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})
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_get_users_key_count(prisma_client):
|
async def test_get_users_key_count(prisma_client):
|
||||||
"""
|
"""
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue