Merge pull request #18503 from BerriAI/litellm_ui_key_list

[Refactor] Key Table to Use Optional "expand" query

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

@@ -93,6 +93,10 @@ const mockKey: KeyResponse = {
   user_tpm_limit: 1000,
   user_rpm_limit: 100,
   user_email: "user@example.com",
+  user: {
+    user_email: "user@example.com",
+    user_id: "user-1",
+  },
 };
 
 const mockTeam: Team = {
@@ -190,3 +194,35 @@ it("should display key information correctly", async () => {
     expect(screen.getByText("5.5000")).toBeInTheDocument();
   });
 });
+
+it("should display user email 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("user@example.com")).toBeInTheDocument();
+  });
+});

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

@@ -1,23 +1,39 @@
 "use client";
-import React, { useEffect, useState } from "react";
-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";
-import { Tooltip } from "antd";
-import { Team, KeyResponse } from "./key_team_helpers/key_list";
-import FilterComponent from "./molecules/filter";
-import { FilterOption } from "./molecules/filter";
-import { Organization, userListCall } from "./networking";
-import { useFilterLogic } from "./key_team_helpers/filter_logic";
 import { Setter } from "@/types";
-import { updateExistingKeys } from "@/utils/dataUtils";
-import { flexRender, getCoreRowModel, getSortedRowModel, SortingState, useReactTable } from "@tanstack/react-table";
-import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell, Icon } from "@tremor/react";
-import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline";
-import { Badge, Text } from "@tremor/react";
+import { formatNumberWithCommas, updateExistingKeys } from "@/utils/dataUtils";
+import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline";
+import {
+  ColumnDef,
+  ColumnResizeDirection,
+  ColumnResizeMode,
+  flexRender,
+  getCoreRowModel,
+  getSortedRowModel,
+  SortingState,
+  useReactTable,
+} from "@tanstack/react-table";
+import {
+  Badge,
+  Button,
+  Icon,
+  Select,
+  SelectItem,
+  Table,
+  TableBody,
+  TableCell,
+  TableHead,
+  TableHeaderCell,
+  TableRow,
+  Text,
+} from "@tremor/react";
+import { Tooltip } from "antd";
+import React, { useEffect, useState } from "react";
 import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key";
-import { formatNumberWithCommas } from "@/utils/dataUtils";
+import { useFilterLogic } from "./key_team_helpers/filter_logic";
+import { KeyResponse, Team } from "./key_team_helpers/key_list";
+import FilterComponent, { FilterOption } from "./molecules/filter";
+import { Organization } from "./networking";
+import KeyInfoView from "./templates/key_info_view";
 
 interface AllKeysTableProps {
   keys: KeyResponse[];
@@ -124,7 +140,6 @@ export function AllKeysTable({
   setAccessToken,
 }: 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>(() => {
@@ -155,17 +170,6 @@ export function AllKeysTable({
       accessToken,
     });
 
-  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]);
-
   // Add a useEffect to call refresh when a key is created
   useEffect(() => {
     if (refresh) {
@@ -269,18 +273,19 @@ export function AllKeysTable({
     },
     {
       id: "user_email",
-      accessorKey: "user_id",
+      accessorKey: "user",
       header: "User Email",
       size: 160,
       cell: (info) => {
-        const userId = info.getValue() as string;
-        const user = userList.find((u) => u.user_id === userId);
-        return user?.user_email ? (
-          <Tooltip title={user?.user_email}>
-            <span>{user?.user_email.slice(0, 20)}...</span>
+        const user = info.getValue() as any;
+        const value = user?.user_email;
+        const width = info.cell.column.getSize();
+        return (
+          <Tooltip title={value}>
+            <span className={`font-mono text-xs truncate block`} style={{ maxWidth: width, overflow: "hidden" }}>
+              {value ?? "-"}
+            </span>
           </Tooltip>
-        ) : (
-          "-"
         );
       },
     },

ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx

@@ -1,6 +1,6 @@
-import { useState, useEffect } from "react";
-import { keyListCall, Member, Organization } from "../networking";
 import { Setter } from "@/types";
+import { useEffect, useState } from "react";
+import { keyListCall, Member, Organization } from "../networking";
 
 export interface Team {
   team_id: string;
@@ -91,6 +91,10 @@ export interface KeyResponse {
   last_rotation_at?: string;
   key_rotation_at?: string;
   next_rotation_at?: string;
+  user?: {
+    user_id: string;
+    user_email: string;
+  };
 }
 
 interface KeyListResponse {
@@ -106,6 +110,7 @@ interface UseKeyListProps {
   selectedKeyAlias: string | null;
   accessToken: string;
   createClicked: boolean;
+  expand?: string[];
 }
 
 interface PaginationData {
@@ -129,6 +134,7 @@ const useKeyList = ({
   selectedKeyAlias,
   accessToken,
   createClicked,
+  expand = [],
 }: UseKeyListProps): UseKeyListReturn => {
   const [keyData, setKeyData] = useState<KeyListResponse>({
     keys: [],
@@ -151,7 +157,19 @@ const useKeyList = ({
       const page = typeof params.page === "number" ? params.page : 1;
       const pageSize = typeof params.pageSize === "number" ? params.pageSize : 100;
 
-      const data = await keyListCall(accessToken, null, null, null, null, null, page, pageSize);
+      const data = await keyListCall(
+        accessToken,
+        null,
+        null,
+        null,
+        null,
+        null,
+        page,
+        pageSize,
+        null,
+        null,
+        expand.join(","),
+      );
       console.log("data", data);
       setKeyData(data);
       setError(null);

ui/litellm-dashboard/src/components/networking.tsx

@@ -3250,6 +3250,7 @@ export const keyListCall = async (
   pageSize: number,
   sortBy: string | null = null,
   sortOrder: string | null = null,
+  expand: string | null = null,
 ) => {
   /**
    * Get all available teams on proxy
@@ -3294,6 +3295,11 @@ export const keyListCall = async (
     if (sortOrder) {
       queryParams.append("sort_order", sortOrder);
     }
+
+    if (expand) {
+      queryParams.append("expand", expand);
+    }
+
     queryParams.append("return_full_object", "true");
     queryParams.append("include_team_keys", "true");
     queryParams.append("include_created_by_keys", "true");

ui/litellm-dashboard/src/components/templates/view_key_table.tsx

@@ -1,17 +1,14 @@
 "use client";
-import React, { useEffect, useState } from "react";
-import { keyDeleteCall, Organization } from "../networking";
-import { add } from "date-fns";
-import { regenerateKeyCall } from "../networking";
-import { Grid, Col, Button, Text, Title, TextInput } from "@tremor/react";
-import { fetchAvailableModelsForTeamOrKey } from "../key_team_helpers/fetch_available_models_team_key";
-import { Modal, Form, InputNumber } from "antd";
-import { CopyToClipboard } from "react-copy-to-clipboard";
-import useKeyList from "../key_team_helpers/key_list";
-import { KeyResponse } from "../key_team_helpers/key_list";
-import { AllKeysTable } from "../all_keys_table";
-import { Team } from "../key_team_helpers/key_list";
 import { Setter } from "@/types";
+import { Button, Col, Grid, Text, TextInput, Title } from "@tremor/react";
+import { Form, InputNumber, Modal } from "antd";
+import { add } from "date-fns";
+import React, { useEffect, useState } from "react";
+import { CopyToClipboard } from "react-copy-to-clipboard";
+import { AllKeysTable } from "../all_keys_table";
+import { fetchAvailableModelsForTeamOrKey } from "../key_team_helpers/fetch_available_models_team_key";
+import useKeyList, { KeyResponse, Team } from "../key_team_helpers/key_list";
+import { keyDeleteCall, Organization, regenerateKeyCall } from "../networking";
 
 import NotificationManager from "../molecules/notifications_manager";
 
@@ -125,6 +122,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     selectedKeyAlias,
     accessToken: accessToken || "",
     createClicked,
+    expand: ["user"],
   });
 
   const handlePageChange = (newPage: number) => {