tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-13 18:23:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #23415 from BerriAI/litellm_style-created-by-not-uuid

Browse Source 
style: make virtual keys tables' created by not a UUID
This commit is contained in:
ryan-crabbe
2026-03-11 19:53:39 -07:00
committed by GitHub
parent 483a605f6e dfda7c10fc
commit 6229e73dfc
6 changed files with 232 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_types.py
+1
View File
@@ -2458,6 +2458,7 @@ class UserAPIKeyAuth(
user_max_budget: Optional[float] = None
request_route: Optional[str] = None
user: Optional[Any] = None # Expanded user object when expand=user is used
created_by_user: Optional[Any] = None # Expanded created_by user when expand=user is used
end_user_object_permission: Optional[LiteLLM_ObjectPermissionTable] = None
model_config = ConfigDict(arbitrary_types_allowed=True)
litellm/proxy/management_endpoints/key_management_endpoints.py
+17 -7
View File
@@ -4597,9 +4597,11 @@ async def _list_key_helper(
user_map = {}
if expand and "user" in expand:
user_ids = [key.user_id for key in keys if key.user_id]
if user_ids:
created_by_ids = [key.created_by for key in keys if key.created_by]
all_ids = list(set(user_ids + created_by_ids)) # Remove duplicates
if all_ids:
users = await prisma_client.db.litellm_usertable.find_many(
where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates
where={"user_id": {"in": all_ids}}
)
user_map = {user.user_id: user for user in users}
@@ -4617,11 +4619,19 @@ async def _list_key_helper(
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
# Include user information if expand includes "user"
if expand and "user" in expand and key.user_id and key.user_id in user_map:
try:
key_dict["user"] = user_map[key.user_id].model_dump()
except Exception:
key_dict["user"] = user_map[key.user_id].dict()
if expand and "user" in expand:
if key.user_id and key.user_id in user_map:
try:
key_dict["user"] = user_map[key.user_id].model_dump()
except Exception:
key_dict["user"] = user_map[key.user_id].dict()
if key.created_by and key.created_by in user_map:
created_by_user = user_map[key.created_by]
key_dict["created_by_user"] = {
"user_id": created_by_user.user_id,
"user_email": created_by_user.user_email,
"user_alias": created_by_user.user_alias,
}
if return_full_object is True or (expand and "user" in expand):
if use_deleted_table:
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+94
View File
@@ -4003,6 +4003,7 @@ async def test_list_keys_with_expand_user():
mock_key1 = MagicMock()
mock_key1.token = "token1"
mock_key1.user_id = "user123"
mock_key1.created_by = None
# Set up model_dump() to raise AttributeError so it falls back to dict()
mock_key1.model_dump = MagicMock(side_effect=AttributeError("model_dump not available"))
mock_key1.dict = MagicMock(return_value=key1_dict)
@@ -4016,6 +4017,7 @@ async def test_list_keys_with_expand_user():
mock_key2 = MagicMock()
mock_key2.token = "token2"
mock_key2.user_id = "user456"
mock_key2.created_by = None
# Set up model_dump() to raise AttributeError so it falls back to dict()
mock_key2.model_dump = MagicMock(side_effect=AttributeError("model_dump not available"))
mock_key2.dict = MagicMock(return_value=key2_dict)
@@ -4120,6 +4122,98 @@ async def test_list_keys_with_expand_user():
}
@pytest.mark.asyncio
async def test_list_keys_with_expand_user_includes_created_by_user():
"""
Test that expand=user also resolves created_by to a user object.
"""
mock_prisma_client = AsyncMock()
# Key created by user789 but owned by user123
key1_dict = {
"token": "token1",
"user_id": "user123",
"created_by": "user789",
"key_alias": "key1",
"models": ["gpt-4"],
}
mock_key1 = MagicMock()
mock_key1.token = "token1"
mock_key1.user_id = "user123"
mock_key1.created_by = "user789"
mock_key1.model_dump = MagicMock(return_value=key1_dict)
mock_find_many_keys = AsyncMock(return_value=[mock_key1])
mock_count_keys = AsyncMock(return_value=1)
# Create mock users for both user_id and created_by
mock_user_owner = MagicMock()
mock_user_owner.user_id = "user123"
mock_user_owner.user_email = "owner@example.com"
mock_user_owner.user_alias = "Owner"
mock_user_owner.model_dump = MagicMock(return_value={
"user_id": "user123",
"user_email": "owner@example.com",
"user_alias": "Owner",
})
mock_user_creator = MagicMock()
mock_user_creator.user_id = "user789"
mock_user_creator.user_email = "creator@example.com"
mock_user_creator.user_alias = "Creator"
mock_find_many_users = AsyncMock(return_value=[mock_user_owner, mock_user_creator])
mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys
mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys
mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users
async def mock_attach_object_permission(d, _):
return d
with patch(
"litellm.proxy.management_endpoints.key_management_endpoints.attach_object_permission_to_dict",
side_effect=mock_attach_object_permission,
):
args = {
"prisma_client": mock_prisma_client,
"page": 1,
"size": 50,
"user_id": None,
"team_id": None,
"organization_id": None,
"key_alias": None,
"key_hash": None,
"exclude_team_id": None,
"return_full_object": False,
"admin_team_ids": None,
"include_created_by_keys": False,
"expand": ["user"],
}
result = await _list_key_helper(**args)
# Verify that the user lookup included both user_id and created_by
call_args = mock_find_many_users.call_args
user_ids_in_query = set(call_args.kwargs["where"]["user_id"]["in"])
assert user_ids_in_query == {"user123", "user789"}
# Verify created_by_user is attached
key_result = result["keys"][0]
assert key_result.created_by_user == {
"user_id": "user789",
"user_email": "creator@example.com",
"user_alias": "Creator",
}
# Verify user (owner) is also still attached
assert key_result.user == {
"user_id": "user123",
"user_email": "owner@example.com",
"user_alias": "Owner",
}
@pytest.mark.asyncio
async def test_list_keys_with_status_deleted():
"""
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
+88
View File
@@ -555,6 +555,94 @@ it("should display 'Default Proxy Admin' for created_by when value is 'default_u
});
it("should display created_by_user email in 'Created By' column when available", async () => {
const keyWithCreatedByUser = {
...mockKey,
created_by: "some-uuid-1234",
created_by_user: {
user_id: "some-uuid-1234",
user_email: "creator@example.com",
user_alias: null,
},
};
mockUseFilterLogic.mockReturnValue({
filters: {
"Team ID": "",
"Organization ID": "",
"Key Alias": "",
"User ID": "",
"Sort By": "created_at",
"Sort Order": "desc",
},
filteredKeys: [keyWithCreatedByUser],
allTeams: [mockTeam],
allOrganizations: [mockOrganization],
handleFilterChange: vi.fn(),
handleFilterReset: vi.fn(),
});
const mockProps = {
teams: [mockTeam],
organizations: [mockOrganization],
onSortChange: vi.fn(),
currentSort: {
sortBy: "created_at",
sortOrder: "desc" as const,
},
};
renderWithProviders(<VirtualKeysTable {...mockProps} />);
await waitFor(() => {
expect(screen.getByText("creator@example.com")).toBeInTheDocument();
});
});
it("should display created_by_user alias over email when both available", async () => {
const keyWithCreatedByUser = {
...mockKey,
created_by: "some-uuid-1234",
created_by_user: {
user_id: "some-uuid-1234",
user_email: "creator@example.com",
user_alias: "The Creator",
},
};
mockUseFilterLogic.mockReturnValue({
filters: {
"Team ID": "",
"Organization ID": "",
"Key Alias": "",
"User ID": "",
"Sort By": "created_at",
"Sort Order": "desc",
},
filteredKeys: [keyWithCreatedByUser],
allTeams: [mockTeam],
allOrganizations: [mockOrganization],
handleFilterChange: vi.fn(),
handleFilterReset: vi.fn(),
});
const mockProps = {
teams: [mockTeam],
organizations: [mockOrganization],
onSortChange: vi.fn(),
currentSort: {
sortBy: "created_at",
sortOrder: "desc" as const,
},
};
renderWithProviders(<VirtualKeysTable {...mockProps} />);
await waitFor(() => {
expect(screen.getByText("The Creator")).toBeInTheDocument();
});
});
it("should render table without crashing when models is null", async () => {
const keyWithNullModels = {
...mockKey,
ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
+27 -12
View File
@@ -311,25 +311,40 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
cell: (info) => {
const userId = info.getValue() as string | null;
if (!userId) return "-";
const key = info.row.original;
const createdByUser = key.created_by_user;
const userAlias = createdByUser?.user_alias ?? null;
const userEmail = createdByUser?.user_email ?? null;
const isDefaultAdmin = userId === "default_user_id";
const displayValue = userAlias || userEmail || userId;
const width = 160;
const popoverContent = (
<div className="flex flex-col gap-2 text-xs min-w-[200px] max-w-[300px]">
<div className="flex flex-col min-w-0">
<span className="text-gray-400">User ID</span>
<Typography.Text
className="font-mono text-xs"
ellipsis={{ tooltip: userId }}
copyable
>
{userId}
</Typography.Text>
</div>
{[
{ label: "User Alias", value: userAlias },
{ label: "User Email", value: userEmail },
{ label: "User ID", value: userId },
].map(({ label, value }) => (
<div key={label} className="flex flex-col min-w-0">
<span className="text-gray-400">{label}</span>
{value ? (
<Typography.Text
className="font-mono text-xs"
ellipsis={{ tooltip: value }}
copyable
>
{value}
</Typography.Text>
) : (
<span className="font-mono">-</span>
)}
</div>
))}
</div>
);
if (isDefaultAdmin) {
if (isDefaultAdmin && !userAlias && !userEmail) {
return (
<Popover content={popoverContent} trigger="hover" placement="bottomLeft">
<span className="cursor-default">
@@ -345,7 +360,7 @@ export function VirtualKeysTable({ teams, organizations, onSortChange, currentSo
className="font-mono text-xs truncate block cursor-default"
style={{ maxWidth: width, overflow: "hidden" }}
>
{userId}
{displayValue}
</span>
</Popover>
);
ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx
+5
View File
@@ -101,6 +101,11 @@ export interface KeyResponse {
user_email: string;
user_alias: string | null;
};
created_by_user?: {
user_id: string;
user_email: string;
user_alias: string | null;
};
}
interface KeyListResponse {