Skip to content

Commit c66c821

Browse files
krrishdholakiakrrishdholakia
authored
UI (Keys Page) - Support cross filtering, filter by user id, filter by key hash (#10322)
* feat(filter.tsx): initial commit making filter component more generic - same style as user table filters * refactor(all_keys_table.tsx): refactor to simplify update logic * fix: partially revert changes - reduce scope of pr * fix(filter_logic.tsx): fix filter update logic * fix(all_keys_table.tsx): fix filtering + search logic * refactor: cleanup unused params * Revert "fix(all_keys_table.tsx): fix filtering + search logic" This reverts commit 5fbc331. * feat(filter_logic.tsx): allow filter by user id * fix(key_management_endpoints.py): support filtering `/key/list` by key hash Enables lookup by key hash on ui * fix(key_list.tsx): fix update * fix(key_management.py): fix linting error * test: update testing * fix(prometheus.py): fix key hash * style(all_keys_table.tsx): style improvements * test: fix test
1 parent 339351c commit c66c821

File tree

12 files changed

+292
-350
lines changed

12 files changed

+292
-350
lines changed

‎litellm/integrations/prometheus.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1000,9 +1000,9 @@ def set_llm_deployment_success_metrics(
10001000
):
10011001
try:
10021002
verbose_logger.debug("setting remaining tokens requests metric")
1003-
standard_logging_payload: Optional[StandardLoggingPayload] = (
1004-
request_kwargs.get("standard_logging_object")
1005-
)
1003+
standard_logging_payload: Optional[
1004+
StandardLoggingPayload
1005+
] = request_kwargs.get("standard_logging_object")
10061006

10071007
if standard_logging_payload is None:
10081008
return
@@ -1453,6 +1453,7 @@ async def fetch_keys(
14531453
user_id=None,
14541454
team_id=None,
14551455
key_alias=None,
1456+
key_hash=None,
14561457
exclude_team_id=UI_SESSION_TOKEN_TEAM_ID,
14571458
return_full_object=True,
14581459
organization_id=None,
@@ -1771,10 +1772,10 @@ def initialize_budget_metrics_cron_job(scheduler: AsyncIOScheduler):
17711772
from litellm.integrations.custom_logger import CustomLogger
17721773
from litellm.integrations.prometheus import PrometheusLogger
17731774

1774-
prometheus_loggers: List[CustomLogger] = (
1775-
litellm.logging_callback_manager.get_custom_loggers_for_type(
1776-
callback_type=PrometheusLogger
1777-
)
1775+
prometheus_loggers: List[
1776+
CustomLogger
1777+
] = litellm.logging_callback_manager.get_custom_loggers_for_type(
1778+
callback_type=PrometheusLogger
17781779
)
17791780
# we need to get the initialized prometheus logger instance(s) and call logger.initialize_remaining_budget_metrics() on them
17801781
verbose_logger.debug("found %s prometheus loggers", len(prometheus_loggers))

‎litellm/proxy/management_endpoints/key_management_endpoints.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1861,6 +1861,7 @@ async def validate_key_list_check(
18611861
team_id: Optional[str],
18621862
organization_id: Optional[str],
18631863
key_alias: Optional[str],
1864+
key_hash: Optional[str],
18641865
prisma_client: PrismaClient,
18651866
) -> Optional[LiteLLM_UserTable]:
18661867
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value:
@@ -1924,6 +1925,31 @@ async def validate_key_list_check(
19241925
param="organization_id",
19251926
code=status.HTTP_403_FORBIDDEN,
19261927
)
1928+
1929+
if key_hash:
1930+
try:
1931+
key_info = await prisma_client.db.litellm_verificationtoken.find_unique(
1932+
where={"token": key_hash},
1933+
)
1934+
except Exception:
1935+
raise ProxyException(
1936+
message="Key Hash not found.",
1937+
type=ProxyErrorTypes.bad_request_error,
1938+
param="key_hash",
1939+
code=status.HTTP_403_FORBIDDEN,
1940+
)
1941+
can_user_query_key_info = await _can_user_query_key_info(
1942+
user_api_key_dict=user_api_key_dict,
1943+
key=key_hash,
1944+
key_info=key_info,
1945+
)
1946+
if not can_user_query_key_info:
1947+
raise HTTPException(
1948+
status_code=status.HTTP_403_FORBIDDEN,
1949+
detail="You are not allowed to access this key's info. Your role={}".format(
1950+
user_api_key_dict.user_role
1951+
),
1952+
)
19271953
return complete_user_info
19281954

19291955

@@ -1972,6 +1998,7 @@ async def list_keys(
19721998
organization_id: Optional[str] = Query(
19731999
None, description="Filter keys by organization ID"
19742000
),
2001+
key_hash: Optional[str] = Query(None, description="Filter keys by key hash"),
19752002
key_alias: Optional[str] = Query(None, description="Filter keys by key alias"),
19762003
return_full_object: bool = Query(False, description="Return full key object"),
19772004
include_team_keys: bool = Query(
@@ -2004,6 +2031,7 @@ async def list_keys(
20042031
team_id=team_id,
20052032
organization_id=organization_id,
20062033
key_alias=key_alias,
2034+
key_hash=key_hash,
20072035
prisma_client=prisma_client,
20082036
)
20092037

@@ -2029,6 +2057,7 @@ async def list_keys(
20292057
user_id=user_id,
20302058
team_id=team_id,
20312059
key_alias=key_alias,
2060+
key_hash=key_hash,
20322061
return_full_object=return_full_object,
20332062
organization_id=organization_id,
20342063
admin_team_ids=admin_team_ids,
@@ -2065,6 +2094,7 @@ async def _list_key_helper(
20652094
team_id: Optional[str],
20662095
organization_id: Optional[str],
20672096
key_alias: Optional[str],
2097+
key_hash: Optional[str],
20682098
exclude_team_id: Optional[str] = None,
20692099
return_full_object: bool = False,
20702100
admin_team_ids: Optional[
@@ -2111,6 +2141,8 @@ async def _list_key_helper(
21112141
user_condition["team_id"] = {"not": exclude_team_id}
21122142
if organization_id and isinstance(organization_id, str):
21132143
user_condition["organization_id"] = organization_id
2144+
if key_hash and isinstance(key_hash, str):
2145+
user_condition["token"] = key_hash
21142146

21152147
if user_condition:
21162148
or_conditions.append(user_condition)

‎tests/litellm/proxy/management_endpoints/test_key_management_endpoints.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ async def test_list_keys():
3030
"team_id": None,
3131
"organization_id": None,
3232
"key_alias": None,
33+
"key_hash": None,
3334
"exclude_team_id": None,
3435
"return_full_object": True,
3536
"admin_team_ids": ["28bd3181-02c5-48f2-b408-ce790fb3d5ba"],

‎tests/proxy_admin_ui_tests/test_key_management.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -989,6 +989,7 @@ async def test_list_key_helper(prisma_client):
989989
user_id=None,
990990
team_id=None,
991991
key_alias=None,
992+
key_hash=None,
992993
organization_id=None,
993994
)
994995
assert len(result["keys"]) == 2, "Should return exactly 2 keys"
@@ -1004,6 +1005,7 @@ async def test_list_key_helper(prisma_client):
10041005
user_id=test_user_id,
10051006
team_id=None,
10061007
key_alias=None,
1008+
key_hash=None,
10071009
organization_id=None,
10081010
)
10091011
assert len(result["keys"]) == 3, "Should return exactly 3 keys for test user"
@@ -1016,6 +1018,7 @@ async def test_list_key_helper(prisma_client):
10161018
user_id=None,
10171019
team_id=test_team_id,
10181020
key_alias=None,
1021+
key_hash=None,
10191022
organization_id=None,
10201023
)
10211024
assert len(result["keys"]) == 2, "Should return exactly 2 keys for test team"
@@ -1028,6 +1031,7 @@ async def test_list_key_helper(prisma_client):
10281031
user_id=None,
10291032
team_id=None,
10301033
key_alias=test_key_alias,
1034+
key_hash=None,
10311035
organization_id=None,
10321036
)
10331037
assert len(result["keys"]) == 1, "Should return exactly 1 key with test alias"
@@ -1040,6 +1044,7 @@ async def test_list_key_helper(prisma_client):
10401044
user_id=test_user_id,
10411045
team_id=None,
10421046
key_alias=None,
1047+
key_hash=None,
10431048
return_full_object=True,
10441049
organization_id=None,
10451050
)
@@ -1141,6 +1146,7 @@ async def test_list_key_helper_team_filtering(prisma_client):
11411146
user_id=None,
11421147
team_id=None,
11431148
key_alias=None,
1149+
key_hash=None,
11441150
return_full_object=True,
11451151
organization_id=None,
11461152
)

‎ui/litellm-dashboard/src/components/all_keys_table.tsx

Lines changed: 49 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
"use client";
2-
import React, { useEffect, useState } from "react";
2+
import React, { useEffect, useState, useCallback, useRef } from "react";
33
import { ColumnDef, Row } from "@tanstack/react-table";
44
import { DataTable } from "./view_logs/table";
55
import { Select, SelectItem } from "@tremor/react"
@@ -9,13 +9,16 @@ import { Tooltip } from "antd";
99
import { Team, KeyResponse } from "./key_team_helpers/key_list";
1010
import FilterComponent from "./common_components/filter";
1111
import { FilterOption } from "./common_components/filter";
12-
import { Organization, userListCall } from "./networking";
12+
import { keyListCall, Organization, userListCall } from "./networking";
1313
import { createTeamSearchFunction } from "./key_team_helpers/team_search_fn";
1414
import { createOrgSearchFunction } from "./key_team_helpers/organization_search_fn";
1515
import { useFilterLogic } from "./key_team_helpers/filter_logic";
1616
import { Setter } from "@/types";
1717
import { updateExistingKeys } from "@/utils/dataUtils";
18-
18+
import { debounce } from "lodash";
19+
import { defaultPageSize } from "./constants";
20+
import { fetchAllTeams } from "./key_team_helpers/filter_helpers";
21+
import { fetchAllOrganizations } from "./key_team_helpers/filter_helpers";
1922
interface AllKeysTableProps {
2023
keys: KeyResponse[];
2124
setKeys: Setter<KeyResponse[]>;
@@ -90,6 +93,8 @@ const TeamFilter = ({
9093
* AllKeysTable – a new table for keys that mimics the table styling used in view_logs.
9194
* The team selector and filtering have been removed so that all keys are shown.
9295
*/
96+
97+
9398
export function AllKeysTable({
9499
keys,
95100
setKeys,
@@ -111,8 +116,8 @@ export function AllKeysTable({
111116
}: AllKeysTableProps) {
112117
const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null);
113118
const [userList, setUserList] = useState<UserResponse[]>([]);
114-
115119
// Use the filter logic hook
120+
116121
const {
117122
filters,
118123
filteredKeys,
@@ -126,11 +131,10 @@ export function AllKeysTable({
126131
teams,
127132
organizations,
128133
accessToken,
129-
setSelectedTeam,
130-
setCurrentOrg,
131-
setSelectedKeyAlias
132134
});
133135

136+
137+
134138
useEffect(() => {
135139
if (accessToken) {
136140
const user_IDs = keys.map(key => key.user_id).filter(id => id !== null);
@@ -208,7 +212,7 @@ export function AllKeysTable({
208212
accessorKey: "team_id", // Change to access the team_id
209213
cell: ({ row, getValue }) => {
210214
const teamId = getValue() as string;
211-
const team = allTeams?.find(t => t.team_id === teamId);
215+
const team = teams?.find(t => t.team_id === teamId);
212216
return team?.team_alias || "Unknown";
213217
},
214218
},
@@ -379,7 +383,18 @@ export function AllKeysTable({
379383
}
380384
});
381385
}
386+
},
387+
{
388+
name: 'User ID',
389+
label: 'User ID',
390+
isSearchable: false,
391+
},
392+
{
393+
name: 'Key Hash',
394+
label: 'Key Hash',
395+
isSearchable: false,
382396
}
397+
383398
];
384399

385400

@@ -414,34 +429,35 @@ export function AllKeysTable({
414429
/>
415430
) : (
416431
<div className="border-b py-4 flex-1 overflow-hidden">
417-
<div className="flex items-center justify-between w-full mb-2">
432+
<div className="w-full mb-6">
418433
<FilterComponent options={filterOptions} onApplyFilters={handleFilterChange} initialValues={filters} onResetFilters={handleFilterReset}/>
419-
<div className="flex items-center gap-4">
420-
<span className="inline-flex text-sm text-gray-700">
421-
Showing {isLoading ? "..." : `${(pagination.currentPage - 1) * pageSize + 1} - ${Math.min(pagination.currentPage * pageSize, pagination.totalCount)}`} of {isLoading ? "..." : pagination.totalCount} results
434+
</div>
435+
436+
<div className="flex items-center justify-between w-full mb-4">
437+
<span className="inline-flex text-sm text-gray-700">
438+
Showing {isLoading ? "..." : `${(pagination.currentPage - 1) * pageSize + 1} - ${Math.min(pagination.currentPage * pageSize, pagination.totalCount)}`} of {isLoading ? "..." : pagination.totalCount} results
439+
</span>
440+
441+
<div className="inline-flex items-center gap-2">
442+
<span className="text-sm text-gray-700">
443+
Page {isLoading ? "..." : pagination.currentPage} of {isLoading ? "..." : pagination.totalPages}
422444
</span>
423445

424-
<div className="inline-flex items-center gap-2">
425-
<span className="text-sm text-gray-700">
426-
Page {isLoading ? "..." : pagination.currentPage} of {isLoading ? "..." : pagination.totalPages}
427-
</span>
428-
429-
<button
430-
onClick={() => onPageChange(pagination.currentPage - 1)}
431-
disabled={isLoading || pagination.currentPage === 1}
432-
className="px-3 py-1 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-50 disabled:cursor-not-allowed"
433-
>
434-
Previous
435-
</button>
436-
437-
<button
438-
onClick={() => onPageChange(pagination.currentPage + 1)}
439-
disabled={isLoading || pagination.currentPage === pagination.totalPages}
440-
className="px-3 py-1 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-50 disabled:cursor-not-allowed"
441-
>
442-
Next
443-
</button>
444-
</div>
446+
<button
447+
onClick={() => onPageChange(pagination.currentPage - 1)}
448+
disabled={isLoading || pagination.currentPage === 1}
449+
className="px-3 py-1 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-50 disabled:cursor-not-allowed"
450+
>
451+
Previous
452+
</button>
453+
454+
<button
455+
onClick={() => onPageChange(pagination.currentPage + 1)}
456+
disabled={isLoading || pagination.currentPage === pagination.totalPages}
457+
className="px-3 py-1 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-50 disabled:cursor-not-allowed"
458+
>
459+
Next
460+
</button>
445461
</div>
446462
</div>
447463
<div className="h-[75vh] overflow-auto">

0 commit comments

Comments
 (0)