Merge pull request #18424 from BerriAI/litellm_ui_keys_columns

[Feature] Allow Column Resizing on Keys Table
This commit is contained in:
yuneng-jiang
2025-12-24 17:52:54 -08:00
committed by GitHub
2 changed files with 263 additions and 4 deletions
@@ -0,0 +1,192 @@
import { screen, waitFor } from "@testing-library/react";
import { vi, it, expect } from "vitest";
import { renderWithProviders } from "../../tests/test-utils";
import { AllKeysTable } from "./all_keys_table";
import { KeyResponse, Team } from "./key_team_helpers/key_list";
import { Organization } from "./networking";
// Mock network calls
vi.mock("./networking", async (importOriginal) => {
const actual = await importOriginal<typeof import("./networking")>();
return {
...actual,
userListCall: vi.fn().mockResolvedValue({
users: [
{
user_id: "user-1",
user_email: "user@example.com",
user_role: "user",
},
],
}),
};
});
// Mock filter helpers
vi.mock("./key_team_helpers/filter_helpers", () => ({
fetchAllKeyAliases: vi.fn().mockResolvedValue(["test-key-alias"]),
fetchAllTeams: vi.fn().mockResolvedValue([
{
team_id: "team-1",
team_alias: "Test Team",
},
]),
fetchAllOrganizations: vi.fn().mockResolvedValue([
{
organization_id: "org-1",
organization_alias: "Test Organization",
},
]),
}));
const mockKey: KeyResponse = {
token: "sk-1234567890abcdef",
token_id: "key-1",
key_name: "test-key",
key_alias: "Test Key Alias",
spend: 5.5,
max_budget: 100,
expires: "2024-12-31T23:59:59Z",
models: ["gpt-3.5-turbo", "gpt-4"],
aliases: {},
config: {},
user_id: "user-1",
team_id: "team-1",
max_parallel_requests: 10,
metadata: {},
tpm_limit: 1000,
rpm_limit: 100,
duration: "30d",
budget_duration: "1m",
budget_reset_at: "2024-12-01T00:00:00Z",
allowed_cache_controls: [],
allowed_routes: [],
permissions: {},
model_spend: { "gpt-3.5-turbo": 2.5, "gpt-4": 3.0 },
model_max_budget: { "gpt-3.5-turbo": 50, "gpt-4": 50 },
soft_budget_cooldown: false,
blocked: false,
litellm_budget_table: {},
organization_id: "org-1",
created_at: "2024-11-01T10:00:00Z",
updated_at: "2024-11-15T10:00:00Z",
team_spend: 5.5,
team_alias: "Test Team",
team_tpm_limit: 5000,
team_rpm_limit: 500,
team_max_budget: 500,
team_models: ["gpt-3.5-turbo", "gpt-4"],
team_blocked: false,
soft_budget: 50,
team_model_aliases: {},
team_member_spend: 0,
team_metadata: {},
end_user_id: "end-user-1",
end_user_tpm_limit: 100,
end_user_rpm_limit: 10,
end_user_max_budget: 10,
last_refreshed_at: Date.now(),
api_key: "sk-1234567890abcdef",
user_role: "user",
rpm_limit_per_model: {},
tpm_limit_per_model: {},
user_tpm_limit: 1000,
user_rpm_limit: 100,
user_email: "user@example.com",
};
const mockTeam: Team = {
team_id: "team-1",
team_alias: "Test Team",
models: ["gpt-3.5-turbo", "gpt-4"],
max_budget: 500,
budget_duration: "1m",
tpm_limit: 5000,
rpm_limit: 500,
organization_id: "org-1",
created_at: "2024-10-01T10:00:00Z",
keys: [],
members_with_roles: [],
};
const mockOrganization: Organization = {
organization_id: "org-1",
organization_alias: "Test Organization",
budget_id: "budget-1",
metadata: {},
models: ["gpt-3.5-turbo", "gpt-4"],
spend: 100,
model_spend: { "gpt-3.5-turbo": 50, "gpt-4": 50 },
created_at: "2024-10-01T10:00:00Z",
created_by: "user-1",
updated_at: "2024-11-01T10:00:00Z",
updated_by: "user-1",
litellm_budget_table: {},
teams: [],
users: [],
members: [],
};
it("should render AllKeysTable component", () => {
const mockProps = {
keys: [mockKey],
setKeys: vi.fn(),
isLoading: false,
pagination: {
currentPage: 1,
totalPages: 1,
totalCount: 1,
},
onPageChange: vi.fn(),
pageSize: 50,
teams: [mockTeam],
selectedTeam: null,
setSelectedTeam: vi.fn(),
selectedKeyAlias: null,
setSelectedKeyAlias: vi.fn(),
accessToken: "test-token",
userID: "user-1",
userRole: "admin",
organizations: [mockOrganization],
setCurrentOrg: vi.fn(),
premiumUser: false,
};
renderWithProviders(<AllKeysTable {...mockProps} />);
expect(screen.getByText("Test Key Alias")).toBeInTheDocument();
});
it("should display key information correctly", async () => {
const mockProps = {
keys: [mockKey],
setKeys: vi.fn(),
isLoading: false,
pagination: {
currentPage: 1,
totalPages: 1,
totalCount: 1,
},
onPageChange: vi.fn(),
pageSize: 50,
teams: [mockTeam],
selectedTeam: null,
setSelectedTeam: vi.fn(),
selectedKeyAlias: null,
setSelectedKeyAlias: vi.fn(),
accessToken: "test-token",
userID: "user-1",
userRole: "admin",
organizations: [mockOrganization],
setCurrentOrg: vi.fn(),
premiumUser: false,
};
renderWithProviders(<AllKeysTable {...mockProps} />);
await waitFor(() => {
expect(screen.getByText("Test Key Alias")).toBeInTheDocument();
expect(screen.getByText("Test Team")).toBeInTheDocument();
expect(screen.getByText("5.5000")).toBeInTheDocument();
});
});
@@ -1,6 +1,6 @@
"use client";
import React, { useEffect, useState } from "react";
import { ColumnDef } from "@tanstack/react-table";
import { ColumnDef, ColumnResizeMode, ColumnResizeDirection } from "@tanstack/react-table";
import { Select, SelectItem } from "@tremor/react";
import { Button } from "@tremor/react";
import KeyInfoView from "./templates/key_info_view";
@@ -125,6 +125,8 @@ export function AllKeysTable({
}: AllKeysTableProps) {
const [selectedKeyId, setSelectedKeyId] = useState<string | null>(null);
const [userList, setUserList] = useState<UserResponse[]>([]);
const [columnResizeMode, setColumnResizeMode] = React.useState<ColumnResizeMode>("onChange");
const [columnResizeDirection, setColumnResizeDirection] = React.useState<ColumnResizeDirection>("ltr");
const [sorting, setSorting] = React.useState<SortingState>(() => {
if (currentSort) {
return [
@@ -184,6 +186,7 @@ export function AllKeysTable({
{
id: "expander",
header: () => null,
size: 40,
cell: ({ row }) =>
row.getCanExpand() ? (
<button onClick={row.getToggleExpandedHandler()} style={{ cursor: "pointer" }}>
@@ -195,6 +198,7 @@ export function AllKeysTable({
id: "token",
accessorKey: "token",
header: "Key ID",
size: 150,
cell: (info) => (
<div className="overflow-hidden">
<Tooltip title={info.getValue() as string}>
@@ -214,10 +218,16 @@ export function AllKeysTable({
id: "key_alias",
accessorKey: "key_alias",
header: "Key Alias",
size: 150,
cell: (info) => {
const value = info.getValue() as string;
const width = info.cell.column.getSize();
return (
<Tooltip title={value}>{value ? (value.length > 20 ? `${value.slice(0, 20)}...` : value) : "-"}</Tooltip>
<Tooltip title={value}>
<span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
{value ?? "-"}
</span>
</Tooltip>
);
},
},
@@ -225,12 +235,14 @@ export function AllKeysTable({
id: "key_name",
accessorKey: "key_name",
header: "Secret Key",
size: 120,
cell: (info) => <span className="font-mono text-xs">{info.getValue() as string}</span>,
},
{
id: "team_alias",
accessorKey: "team_id",
header: "Team Alias",
size: 120,
cell: ({ row, getValue }) => {
const teamId = getValue() as string;
const team = teams?.find((t) => t.team_id === teamId);
@@ -241,6 +253,7 @@ export function AllKeysTable({
id: "team_id",
accessorKey: "team_id",
header: "Team ID",
size: 120,
cell: (info) => (
<Tooltip title={info.getValue() as string}>
{info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"}
@@ -251,12 +264,14 @@ export function AllKeysTable({
id: "organization_id",
accessorKey: "organization_id",
header: "Organization ID",
size: 140,
cell: (info) => (info.getValue() ? info.renderValue() : "-"),
},
{
id: "user_email",
accessorKey: "user_id",
header: "User Email",
size: 160,
cell: (info) => {
const userId = info.getValue() as string;
const user = userList.find((u) => u.user_id === userId);
@@ -273,6 +288,7 @@ export function AllKeysTable({
id: "user_id",
accessorKey: "user_id",
header: "User ID",
size: 120,
cell: (info) => {
const userId = info.getValue() as string | null;
if (userId && userId.length > 15) {
@@ -289,6 +305,7 @@ export function AllKeysTable({
id: "created_at",
accessorKey: "created_at",
header: "Created At",
size: 120,
cell: (info) => {
const value = info.getValue();
return value ? new Date(value as string).toLocaleDateString() : "-";
@@ -298,6 +315,7 @@ export function AllKeysTable({
id: "created_by",
accessorKey: "created_by",
header: "Created By",
size: 120,
cell: (info) => {
const value = info.getValue() as string | null;
if (value && value.length > 15) {
@@ -314,6 +332,7 @@ export function AllKeysTable({
id: "updated_at",
accessorKey: "updated_at",
header: "Updated At",
size: 120,
cell: (info) => {
const value = info.getValue();
return value ? new Date(value as string).toLocaleDateString() : "Never";
@@ -323,6 +342,7 @@ export function AllKeysTable({
id: "expires",
accessorKey: "expires",
header: "Expires",
size: 120,
cell: (info) => {
const value = info.getValue();
return value ? new Date(value as string).toLocaleDateString() : "Never";
@@ -332,12 +352,14 @@ export function AllKeysTable({
id: "spend",
accessorKey: "spend",
header: "Spend (USD)",
size: 100,
cell: (info) => formatNumberWithCommas(info.getValue() as number, 4),
},
{
id: "max_budget",
accessorKey: "max_budget",
header: "Budget (USD)",
size: 110,
cell: (info) => {
const maxBudget = info.getValue() as number | null;
if (maxBudget === null) {
@@ -350,6 +372,7 @@ export function AllKeysTable({
id: "budget_reset_at",
accessorKey: "budget_reset_at",
header: "Budget Reset",
size: 130,
cell: (info) => {
const value = info.getValue();
return value ? new Date(value as string).toLocaleString() : "Never";
@@ -359,6 +382,7 @@ export function AllKeysTable({
id: "models",
accessorKey: "models",
header: "Models",
size: 200,
cell: (info) => {
const models = info.getValue() as string[];
return (
@@ -442,6 +466,7 @@ export function AllKeysTable({
{
id: "rate_limits",
header: "Rate Limits",
size: 140,
cell: ({ row }) => {
const key = row.original;
return (
@@ -527,6 +552,8 @@ export function AllKeysTable({
const table = useReactTable({
data: filteredKeys,
columns: columns.filter((col) => col.id !== "expander"),
columnResizeMode,
columnResizeDirection,
state: {
sorting,
},
@@ -639,18 +666,35 @@ export function AllKeysTable({
<div className="h-[75vh] overflow-auto">
<div className="rounded-lg custom-border relative">
<div className="overflow-x-auto">
<Table className="[&_td]:py-0.5 [&_th]:py-1">
<Table className="[&_td]:py-0.5 [&_th]:py-1" style={{ width: table.getCenterTotalSize() }}>
<TableHead>
{table.getHeaderGroups().map((headerGroup) => (
<TableRow key={headerGroup.id}>
{headerGroup.headers.map((header) => (
<TableHeaderCell
key={header.id}
className={`py-1 h-8 ${
data-header-id={header.id}
className={`py-1 h-8 relative hover:bg-gray-50 ${
header.id === "actions"
? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
: ""
}`}
style={{
width: header.getSize(),
position: "relative",
}}
onMouseEnter={() => {
const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`);
if (resizer) {
(resizer as HTMLElement).style.opacity = "0.5";
}
}}
onMouseLeave={() => {
const resizer = document.querySelector(`[data-header-id="${header.id}"] .resizer`);
if (resizer && !header.column.getIsResizing()) {
(resizer as HTMLElement).style.opacity = "0";
}
}}
onClick={header.column.getToggleSortingHandler()}
>
<div className="flex items-center justify-between gap-2">
@@ -671,6 +715,28 @@ export function AllKeysTable({
)}
</div>
)}
<div
onDoubleClick={() => header.column.resetSize()}
onMouseDown={header.getResizeHandler()}
onTouchStart={header.getResizeHandler()}
className={`resizer ${table.options.columnResizeDirection} ${header.column.getIsResizing() ? "isResizing" : ""}`}
style={{
position: "absolute",
right: 0,
top: 0,
height: "100%",
width: "5px",
background: header.column.getIsResizing() ? "#3b82f6" : "transparent",
cursor: "col-resize",
userSelect: "none",
touchAction: "none",
opacity: header.column.getIsResizing() ? 1 : 0,
transform:
columnResizeMode === "onEnd" && header.column.getIsResizing()
? `translateX(${(table.options.columnResizeDirection === "rtl" ? -1 : 1) * (table.getState().columnSizingInfo.deltaOffset ?? 0)}px)`
: "",
}}
/>
</div>
</TableHeaderCell>
))}
@@ -693,6 +759,7 @@ export function AllKeysTable({
<TableCell
key={cell.id}
style={{
width: cell.column.getSize(),
maxWidth: "8-x",
whiteSpace: "pre-wrap",
overflow: "hidden",