Show 'user_email' on key table on UI (#8887)

* refactor(internal_user_endpoints.py): refactor `/user/list` to accept 'user_ids' and use prisma for db calls

enables bulk search from UI

* fix(internal_user_endpoints.py): fix linting errors

* fix(all_keys_table.tsx): show user email on create key table

make it easier for admin to know which key is associated to which user

* docs(internal_user_endpoints.py): improve docstring

* fix: sync schema with main

* fix(columns.tsx): display SSO ID on Internal User Table

make it easy to identify what the SSO ID for a user is

* fix(columns.tsx): add tooltip to header

help user understand what SSO ID means

* style: add more tooltips in the management flows

make it easier to understand what you're seeing

* style(all_keys_table.tsx): replace 'Not Set' with '-'

reduces words on table

* fix(internal_user_endpoints.py): fix user ids check

* test: fix test

* fix(internal_user_endpoints.py): maintain returning key count in `/user/list`
This commit is contained in:
Krish Dholakia 2025-02-27 21:56:14 -08:00 committed by GitHub
parent 475c1d0f99
commit bd2b6bdeb3
8 changed files with 199 additions and 67 deletions

View file

@ -1578,6 +1578,10 @@ class LiteLLM_UserTableFiltered(BaseModel): # done to avoid exposing sensitive
user_email: str user_email: str
class LiteLLM_UserTableWithKeyCount(LiteLLM_UserTable):
key_count: int = 0
class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase): class LiteLLM_EndUserTable(LiteLLMPydanticObjectBase):
user_id: str user_id: str
blocked: bool blocked: bool

View file

@ -753,6 +753,9 @@ async def get_users(
role: Optional[str] = fastapi.Query( role: Optional[str] = fastapi.Query(
default=None, description="Filter users by role" default=None, description="Filter users by role"
), ),
user_ids: Optional[str] = fastapi.Query(
default=None, description="Get list of users by user_ids"
),
page: int = fastapi.Query(default=1, ge=1, description="Page number"), page: int = fastapi.Query(default=1, ge=1, description="Page number"),
page_size: int = fastapi.Query( page_size: int = fastapi.Query(
default=25, ge=1, le=100, description="Number of items per page" default=25, ge=1, le=100, description="Number of items per page"
@ -770,12 +773,19 @@ async def get_users(
- proxy_admin_viewer - proxy_admin_viewer
- internal_user - internal_user
- internal_user_viewer - internal_user_viewer
user_ids: Optional[str]
Get list of users by user_ids. Comma separated list of user_ids.
page: int page: int
The page number to return The page number to return
page_size: int page_size: int
The number of items per page The number of items per page
Currently - admin-only endpoint. Currently - admin-only endpoint.
Example curl:
```
http://0.0.0.0:4000/user/list?user_ids=default_user_id,693c1a4a-1cc0-4c7c-afe8-b5d2c8d52e17
```
""" """
from litellm.proxy.proxy_server import prisma_client from litellm.proxy.proxy_server import prisma_client
@ -787,49 +797,69 @@ async def get_users(
# Calculate skip and take for pagination # Calculate skip and take for pagination
skip = (page - 1) * page_size skip = (page - 1) * page_size
take = page_size
# Prepare the query conditions # Prepare the query conditions
where_clause = "" # Build where conditions based on provided parameters
where_conditions: Dict[str, Any] = {}
if role: if role:
where_clause = f"""WHERE "user_role" = '{role}'""" where_conditions["user_role"] = {
"contains": role,
"mode": "insensitive", # Case-insensitive search
}
# Single optimized SQL query that gets both users and total count if user_ids and isinstance(user_ids, str):
sql_query = f""" user_id_list = [uid.strip() for uid in user_ids.split(",") if uid.strip()]
WITH total_users AS ( where_conditions["user_id"] = {
SELECT COUNT(*) AS total_number_internal_users "in": user_id_list, # Now passing a list of strings as required by Prisma
FROM "LiteLLM_UserTable" }
),
paginated_users AS ( users: Optional[List[LiteLLM_UserTable]] = (
SELECT await prisma_client.db.litellm_usertable.find_many(
u.*, where=where_conditions,
( skip=skip,
SELECT COUNT(*) take=page_size,
FROM "LiteLLM_VerificationToken" vt order={"created_at": "desc"},
WHERE vt."user_id" = u."user_id" )
) AS key_count
FROM "LiteLLM_UserTable" u
{where_clause}
LIMIT {take} OFFSET {skip}
) )
SELECT
(SELECT total_number_internal_users FROM total_users),
*
FROM paginated_users;
"""
# Execute the query # Get total count of user rows
results = await prisma_client.db.query_raw(sql_query) total_count = await prisma_client.db.litellm_usertable.count(
# Get total count from the first row (if results exist) where=where_conditions # type: ignore
total_count = 0 )
if len(results) > 0:
total_count = results[0].get("total_number_internal_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 = {
item["user_id"]: item["_count"]["user_id"] for item in user_keys
}
else:
user_key_counts = {}
verbose_proxy_logger.debug(f"Total count of users: {total_count}")
# Calculate total pages # Calculate total pages
total_pages = -(-total_count // page_size) # Ceiling division total_pages = -(-total_count // page_size) # Ceiling division
# Prepare response
user_list: List[LiteLLM_UserTableWithKeyCount] = []
if users is not None:
for user in users:
user_list.append(
LiteLLM_UserTableWithKeyCount(
**user.model_dump(), key_count=user_key_counts.get(user.user_id, 0)
)
) # Return full key object
else:
user_list = []
return { return {
"users": results, "users": user_list,
"total": total_count, "total": total_count,
"page": page, "page": page,
"page_size": page_size, "page_size": page_size,

View file

@ -370,11 +370,7 @@ async def test_get_users(prisma_client):
assert "users" in result assert "users" in result
for user in result["users"]: for user in result["users"]:
assert "user_id" in user assert isinstance(user, LiteLLM_UserTable)
assert "spend" in user
assert "user_email" in user
assert "user_role" in user
assert "key_count" in user
# Clean up test users # Clean up test users
for user in test_users: for user in test_users:
@ -397,12 +393,12 @@ async def test_get_users_key_count(prisma_client):
assert len(initial_users["users"]) > 0, "No users found to test with" assert len(initial_users["users"]) > 0, "No users found to test with"
test_user = initial_users["users"][0] test_user = initial_users["users"][0]
initial_key_count = test_user["key_count"] initial_key_count = test_user.key_count
# Create a new key for the selected user # Create a new key for the selected user
new_key = await generate_key_fn( new_key = await generate_key_fn(
data=GenerateKeyRequest( data=GenerateKeyRequest(
user_id=test_user["user_id"], user_id=test_user.user_id,
key_alias=f"test_key_{uuid.uuid4()}", key_alias=f"test_key_{uuid.uuid4()}",
models=["fake-model"], models=["fake-model"],
), ),
@ -418,8 +414,8 @@ async def test_get_users_key_count(prisma_client):
print("updated_users", updated_users) print("updated_users", updated_users)
updated_key_count = None updated_key_count = None
for user in updated_users["users"]: for user in updated_users["users"]:
if user["user_id"] == test_user["user_id"]: if user.user_id == test_user.user_id:
updated_key_count = user["key_count"] updated_key_count = user.key_count
break break
assert updated_key_count is not None, "Test user not found in updated users list" assert updated_key_count is not None, "Test user not found in updated users list"

View file

@ -1,5 +1,5 @@
"use client"; "use client";
import React, { useState } from "react"; import React, { useEffect, useState } from "react";
import { ColumnDef, Row } from "@tanstack/react-table"; import { ColumnDef, Row } from "@tanstack/react-table";
import { DataTable } from "./view_logs/table"; import { DataTable } from "./view_logs/table";
import { Select, SelectItem } from "@tremor/react" import { Select, SelectItem } from "@tremor/react"
@ -9,7 +9,7 @@ import { Tooltip } from "antd";
import { Team, KeyResponse } from "./key_team_helpers/key_list"; import { Team, KeyResponse } from "./key_team_helpers/key_list";
import FilterComponent from "./common_components/filter"; import FilterComponent from "./common_components/filter";
import { FilterOption } from "./common_components/filter"; import { FilterOption } from "./common_components/filter";
import { Organization } from "./networking"; import { Organization, userListCall } from "./networking";
import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn"; import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn";
import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn"; import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn";
interface AllKeysTableProps { interface AllKeysTableProps {
@ -34,6 +34,12 @@ interface AllKeysTableProps {
// Define columns similar to our logs table // Define columns similar to our logs table
interface UserResponse {
user_id: string;
user_email: string;
user_role: string;
}
const TeamFilter = ({ const TeamFilter = ({
teams, teams,
selectedTeam, selectedTeam,
@ -99,6 +105,18 @@ export function AllKeysTable({
'Team ID': '', 'Team ID': '',
'Organization ID': '' 'Organization ID': ''
}); });
const [userList, setUserList] = useState<UserResponse[]>([]);
useEffect(() => {
if (accessToken) {
const user_IDs = keys.map(key => key.user_id).filter(id => id !== null);
const fetchUserList = async () => {
const userListData = await userListCall(accessToken, user_IDs, 1, 100);
setUserList(userListData.users);
};
fetchUserList();
}
}, [accessToken, keys]);
const handleFilterChange = (newFilters: Record<string, string>) => { const handleFilterChange = (newFilters: Record<string, string>) => {
// Update filters state // Update filters state
@ -163,7 +181,7 @@ export function AllKeysTable({
className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]" className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]"
onClick={() => setSelectedKeyId(info.getValue() as string)} onClick={() => setSelectedKeyId(info.getValue() as string)}
> >
{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "Not Set"} {info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}
</Button> </Button>
</Tooltip> </Tooltip>
</div> </div>
@ -186,31 +204,28 @@ export function AllKeysTable({
{ {
header: "Team ID", header: "Team ID",
accessorKey: "team_id", accessorKey: "team_id",
cell: (info) => <Tooltip title={info.getValue() as string}>{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "Not Set"}</Tooltip> cell: (info) => <Tooltip title={info.getValue() as string}>{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}</Tooltip>
}, },
{ {
header: "Key Alias", header: "Key Alias",
accessorKey: "key_alias", accessorKey: "key_alias",
cell: (info) => <Tooltip title={info.getValue() as string}>{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "Not Set"}</Tooltip> cell: (info) => <Tooltip title={info.getValue() as string}>{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}</Tooltip>
}, },
{ {
header: "Organization ID", header: "Organization ID",
accessorKey: "organization_id", accessorKey: "organization_id",
cell: (info) => <Tooltip title={info.getValue() as string}>{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "Not Set"}</Tooltip> cell: (info) => info.getValue() ? info.renderValue() : "-",
},
{
header: "User Email",
accessorKey: "user_id",
cell: (info) => {
const userId = info.getValue() as string;
const user = userList.find(u => u.user_id === userId);
return user?.user_email ? user.user_email : "-";
},
}, },
// {
// header: "User Email",
// accessorKey: "user_id",
// cell: (info) => {
// const userId = info.getValue() as string;
// return userId ? (
// <Tooltip title={userId}>
// <span>{userId.slice(0, 5)}...</span>
// </Tooltip>
// ) : "Not Set";
// },
// },
{ {
header: "User ID", header: "User ID",
accessorKey: "user_id", accessorKey: "user_id",
@ -218,9 +233,9 @@ export function AllKeysTable({
const userId = info.getValue() as string; const userId = info.getValue() as string;
return userId ? ( return userId ? (
<Tooltip title={userId}> <Tooltip title={userId}>
<span>{userId.slice(0, 5)}...</span> <span>{userId.slice(0, 7)}...</span>
</Tooltip> </Tooltip>
) : "Not Set"; ) : "-";
}, },
}, },
{ {

View file

@ -21,6 +21,8 @@ import {
} from "./networking"; } from "./networking";
import BulkCreateUsers from "./bulk_create_users_button"; import BulkCreateUsers from "./bulk_create_users_button";
const { Option } = Select; const { Option } = Select;
import { Tooltip } from "antd";
import { InfoCircleOutlined } from '@ant-design/icons';
interface CreateuserProps { interface CreateuserProps {
userID: string; userID: string;
@ -258,7 +260,15 @@ const Createuser: React.FC<CreateuserProps> = ({
<Form.Item label="User Email" name="user_email"> <Form.Item label="User Email" name="user_email">
<TextInput placeholder="" /> <TextInput placeholder="" />
</Form.Item> </Form.Item>
<Form.Item label="User Role" name="user_role"> <Form.Item label={
<span>
Global Proxy Role{' '}
<Tooltip title="This is the role that the user will globally on the proxy. This role is independent of any team/org specific roles.">
<InfoCircleOutlined/>
</Tooltip>
</span>
}
name="user_role">
<Select2> <Select2>
{possibleUIRoles && {possibleUIRoles &&
Object.entries(possibleUIRoles).map( Object.entries(possibleUIRoles).map(
@ -278,7 +288,7 @@ const Createuser: React.FC<CreateuserProps> = ({
)} )}
</Select2> </Select2>
</Form.Item> </Form.Item>
<Form.Item label="Team ID" name="team_id"> <Form.Item label="Team ID" name="team_id" help="If selected, user will be added as a 'user' role to the team.">
<Select placeholder="Select Team ID" style={{ width: "100%" }}> <Select placeholder="Select Team ID" style={{ width: "100%" }}>
{teams ? ( {teams ? (
teams.map((team: any) => ( teams.map((team: any) => (

View file

@ -651,6 +651,66 @@ export const teamDeleteCall = async (accessToken: String, teamID: String) => {
} }
}; };
export const userListCall = async (
accessToken: String,
userIDs: string[] | null = null,
page: number | null = null,
page_size: number | null = null,
) => {
/**
* Get all available teams on proxy
*/
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/user/list` : `/user/list`;
console.log("in userListCall");
const queryParams = new URLSearchParams();
if (userIDs && userIDs.length > 0) {
// Convert array to comma-separated string
const userIDsString = userIDs.join(',');
queryParams.append('user_ids', userIDsString);
}
if (page) {
queryParams.append('page', page.toString());
}
if (page_size) {
queryParams.append('page_size', page_size.toString());
}
const queryString = queryParams.toString();
if (queryString) {
url += `?${queryString}`;
}
const response = await fetch(url, {
method: "GET",
headers: {
[globalLitellmHeaderName]: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
});
if (!response.ok) {
const errorData = await response.text();
handleError(errorData);
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log("/user/list API Response:", data);
return data;
// Handle success - you might want to update some state or UI based on the created key
} catch (error) {
console.error("Failed to create key:", error);
throw error;
}
};
export const userInfoCall = async ( export const userInfoCall = async (
accessToken: String, accessToken: String,
userID: String | null, userID: String | null,

View file

@ -2,7 +2,7 @@ import { ColumnDef } from "@tanstack/react-table";
import { Badge, Grid, Icon } from "@tremor/react"; import { Badge, Grid, Icon } from "@tremor/react";
import { Tooltip } from "antd"; import { Tooltip } from "antd";
import { UserInfo } from "./types"; import { UserInfo } from "./types";
import { PencilAltIcon, TrashIcon } from "@heroicons/react/outline"; import { PencilAltIcon, TrashIcon, InformationCircleIcon } from "@heroicons/react/outline";
export const columns = ( export const columns = (
possibleUIRoles: Record<string, Record<string, string>>, possibleUIRoles: Record<string, Record<string, string>>,
@ -14,7 +14,7 @@ export const columns = (
accessorKey: "user_id", accessorKey: "user_id",
cell: ({ row }) => ( cell: ({ row }) => (
<Tooltip title={row.original.user_id}> <Tooltip title={row.original.user_id}>
<span className="text-xs">{row.original.user_id ? `${row.original.user_id.slice(0, 4)}...` : "-"}</span> <span className="text-xs">{row.original.user_id ? `${row.original.user_id.slice(0, 7)}...` : "-"}</span>
</Tooltip> </Tooltip>
), ),
}, },
@ -26,7 +26,7 @@ export const columns = (
), ),
}, },
{ {
header: "Role", header: "Global Proxy Role",
accessorKey: "user_role", accessorKey: "user_role",
cell: ({ row }) => ( cell: ({ row }) => (
<span className="text-xs"> <span className="text-xs">
@ -52,6 +52,22 @@ export const columns = (
</span> </span>
), ),
}, },
{
header: () => (
<div className="flex items-center gap-2">
<span>SSO ID</span>
<Tooltip title="SSO ID is the ID of the user in the SSO provider. If the user is not using SSO, this will be null.">
<InformationCircleIcon className="w-4 h-4" />
</Tooltip>
</div>
),
accessorKey: "sso_user_id",
cell: ({ row }) => (
<span className="text-xs">
{row.original.sso_user_id !== null ? row.original.sso_user_id : "-"}
</span>
),
},
{ {
header: "API Keys", header: "API Keys",
accessorKey: "key_count", accessorKey: "key_count",

View file

@ -7,4 +7,5 @@ export interface UserInfo {
key_count: number; key_count: number;
created_at: string; created_at: string;
updated_at: string; updated_at: string;
} sso_user_id: string | null;
}