Merge pull request #19047 from BerriAI/litellm_usage_column_names

[Fix] UI - Usage: Team ID and Team Name in Export Report

ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.test.tsx

@@ -10,7 +10,7 @@
  */
 
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render } from "@testing-library/react";
+import { renderWithProviders } from "../../../tests/test-utils";
 import userEvent from "@testing-library/user-event";
 import EntityUsageExportModal from "./EntityUsageExportModal";
 
@@ -35,6 +35,16 @@ vi.mock("../molecules/notifications_manager", () => {
   };
 });
 
+// Mock useTeams hook
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  useTeams: vi.fn(() => ({
+    data: [],
+    isLoading: false,
+    error: null,
+    refetch: vi.fn(),
+  })),
+}));
+
 // JSDOM stubs for download flow used by the modal
 // @ts-ignore
 global.URL.createObjectURL = vi.fn(() => "blob:mock");
@@ -74,7 +84,7 @@ describe("EntityUsageExportModal", () => {
     const user = userEvent.setup();
     const { handleExportCSV } = await import("./utils");
 
-    const { getByRole } = render(<EntityUsageExportModal {...baseProps} />);
+    const { getByRole } = renderWithProviders(<EntityUsageExportModal {...baseProps} />);
 
     // Default primary action reflects CSV export
     expect(getByRole("button", { name: /Export CSV/i })).toBeInTheDocument();
@@ -83,7 +93,7 @@ describe("EntityUsageExportModal", () => {
     await user.click(getByRole("button", { name: /Export CSV/i }));
 
     // Verifies export function was invoked with correct parameters
-    expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily", "Tag", "tag");
+    expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily", "Tag", "tag", {});
 
     // Modal closes after export
     expect(baseProps.onClose).toHaveBeenCalled();
@@ -98,7 +108,7 @@ describe("EntityUsageExportModal", () => {
     const user = userEvent.setup();
     const { handleExportCSV } = await import("./utils");
 
-    const { getByText, getByRole } = render(<EntityUsageExportModal {...baseProps} />);
+    const { getByText, getByRole } = renderWithProviders(<EntityUsageExportModal {...baseProps} />);
 
     // Choose the alternate export type - click the label to trigger radio
     const dailyModelLabel = getByText(/Day-by-day by tag and model/i);
@@ -109,7 +119,7 @@ describe("EntityUsageExportModal", () => {
     await user.click(exportBtn);
 
     // Ensure the selected scope flowed through
-    expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily_with_models", "Tag", "tag");
+    expect(handleExportCSV).toHaveBeenCalledWith(baseProps.spendData, "daily_with_models", "Tag", "tag", {});
 
     // Modal closes after export
     expect(baseProps.onClose).toHaveBeenCalled();

ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx

@@ -1,12 +1,13 @@
-import React, { useState } from "react";
-import { Button } from "@tremor/react";
-import { Modal } from "antd";
+import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
+import { createTeamAliasMap } from "@/utils/teamUtils";
+import { Button, Modal, Skeleton } from "antd";
+import React, { useMemo, useState } from "react";
 import NotificationsManager from "../molecules/notifications_manager";
+import ExportFormatSelector from "./ExportFormatSelector";
 import ExportSummary from "./ExportSummary";
 import ExportTypeSelector from "./ExportTypeSelector";
-import ExportFormatSelector from "./ExportFormatSelector";
-import { handleExportCSV, handleExportJSON } from "./utils";
 import type { EntityUsageExportModalProps, ExportFormat, ExportScope } from "./types";
+import { handleExportCSV, handleExportJSON } from "./utils";
 
 const EntityUsageExportModal: React.FC<EntityUsageExportModalProps> = ({
   isOpen,
@@ -20,19 +21,22 @@ const EntityUsageExportModal: React.FC<EntityUsageExportModalProps> = ({
   const [exportFormat, setExportFormat] = useState<ExportFormat>("csv");
   const [exportScope, setExportScope] = useState<ExportScope>("daily");
   const [isExporting, setIsExporting] = useState(false);
+  const { data: teams, isLoading: isLoadingTeams } = useTeams();
 
   const entityLabel = entityType.charAt(0).toUpperCase() + entityType.slice(1);
   const modalTitle = customTitle || `Export ${entityLabel} Usage`;
 
+  // Cache team alias map using useMemo
+  const teamAliasMap = useMemo(() => createTeamAliasMap(teams), [teams]);
   const handleExport = async (format?: ExportFormat) => {
     const formatToUse = format || exportFormat;
     setIsExporting(true);
     try {
       if (formatToUse === "csv") {
-        handleExportCSV(spendData, exportScope, entityLabel, entityType);
+        handleExportCSV(spendData, exportScope, entityLabel, entityType, teamAliasMap);
         NotificationsManager.success(`${entityLabel} usage data exported successfully as CSV`);
       } else {
-        handleExportJSON(spendData, exportScope, entityLabel, entityType, dateRange, selectedFilters);
+        handleExportJSON(spendData, exportScope, entityLabel, entityType, dateRange, selectedFilters, teamAliasMap);
         NotificationsManager.success(`${entityLabel} usage data exported successfully as JSON`);
       }
       onClose();
@@ -51,23 +55,37 @@ const EntityUsageExportModal: React.FC<EntityUsageExportModalProps> = ({
       onCancel={onClose}
       footer={null}
       width={480}
-      destroyOnClose
     >
       <div className="space-y-5 py-2">
-        <ExportSummary dateRange={dateRange} selectedFilters={selectedFilters} />
-
-        <ExportTypeSelector value={exportScope} onChange={setExportScope} entityType={entityType} />
-
-        <ExportFormatSelector value={exportFormat} onChange={setExportFormat} />
-
-        <div className="flex items-center justify-end gap-2 pt-4 border-t">
-          <Button variant="secondary" onClick={onClose} disabled={isExporting} size="sm">
-            Cancel
-          </Button>
-          <Button onClick={() => handleExport()} loading={isExporting} disabled={isExporting} size="sm">
-            {isExporting ? "Exporting..." : `Export ${exportFormat.toUpperCase()}`}
-          </Button>
-        </div>
+        {isLoadingTeams ? (
+          <Skeleton active />
+        ) : (
+          <>
+            <ExportSummary dateRange={dateRange} selectedFilters={selectedFilters} />
+            <ExportTypeSelector value={exportScope} onChange={setExportScope} entityType={entityType} />
+            <ExportFormatSelector value={exportFormat} onChange={setExportFormat} />
+          </>
+        )}
+        {isLoadingTeams ? (
+          <div className="flex items-center justify-end gap-2 pt-4 border-t">
+            <Skeleton.Button active />
+            <Skeleton.Button active />
+          </div>
+        ) : (
+          <div className="flex items-center justify-end gap-2 pt-4 border-t">
+            <Button variant="outlined" onClick={onClose} disabled={isExporting}>
+              Cancel
+            </Button>
+            <Button
+              onClick={() => handleExport()}
+              loading={isExporting || isLoadingTeams}
+              disabled={isExporting || isLoadingTeams}
+              type="primary"
+            >
+              {isExporting ? "Exporting..." : `Export ${exportFormat.toUpperCase()}`}
+            </Button>
+          </div>
+        )}
       </div>
     </Modal>
   );

ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx

@@ -4,6 +4,7 @@ import { Select } from "antd";
 import React, { useState } from "react";
 import EntityUsageExportModal from "./EntityUsageExportModal";
 import type { EntitySpendData, EntityType } from "./types";
+import type { Team } from "@/components/key_team_helpers/key_list";
 
 interface UsageExportHeaderProps {
   dateValue: DateRangePickerValue;
@@ -18,6 +19,7 @@ interface UsageExportHeaderProps {
   filterOptions?: Array<{ label: string; value: string }>;
   customTitle?: string;
   compactLayout?: boolean;
+  teams?: Team[];
 }
 
 const UsageExportHeader: React.FC<UsageExportHeaderProps> = ({
@@ -32,6 +34,7 @@ const UsageExportHeader: React.FC<UsageExportHeaderProps> = ({
   filterOptions = [],
   customTitle,
   compactLayout = false,
+  teams = [],
 }) => {
   const [isExportModalOpen, setIsExportModalOpen] = useState(false);
 
@@ -95,6 +98,7 @@ const UsageExportHeader: React.FC<UsageExportHeaderProps> = ({
         dateRange={dateValue}
         selectedFilters={selectedFilters}
         customTitle={customTitle}
+        teams={teams}
       />
     </>
   );

ui/litellm-dashboard/src/components/EntityUsageExport/types.ts

@@ -1,4 +1,5 @@
 import type { DateRangePickerValue } from "@tremor/react";
+import type { Team } from "@/components/key_team_helpers/key_list";
 
 export type ExportFormat = "csv" | "json";
 export type ExportScope = "daily" | "daily_with_models";
@@ -23,6 +24,7 @@ export interface EntityUsageExportModalProps {
   dateRange: DateRangePickerValue;
   selectedFilters: string[];
   customTitle?: string;
+  teams?: Team[];
 }
 
 export interface ExportMetadata {

ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts

@@ -1,13 +1,43 @@
 import { formatNumberWithCommas } from "@/utils/dataUtils";
-import Papa from "papaparse";
-import type { EntitySpendData, EntityBreakdown, ExportMetadata, ExportScope, EntityType } from "./types";
 import type { DateRangePickerValue } from "@tremor/react";
+import Papa from "papaparse";
+import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types";
 
-export const getEntityBreakdown = (spendData: EntitySpendData): EntityBreakdown[] => {
+// Helper function to extract team_id from api_key_breakdown
+const extractTeamIdFromApiKeyBreakdown = (apiKeyBreakdown: Record<string, any> | undefined): string | null => {
+  if (!apiKeyBreakdown) return null;
+
+  // Look through all API keys to find the first non-null team_id
+  for (const apiKeyData of Object.values(apiKeyBreakdown)) {
+    const teamId = (apiKeyData as any)?.metadata?.team_id;
+    if (teamId) {
+      return teamId;
+    }
+  }
+  return null;
+};
+
+export const getEntityBreakdown = (
+  spendData: EntitySpendData,
+  teamAliasMap: Record<string, string> = {},
+): EntityBreakdown[] => {
   const entitySpend: { [key: string]: EntityBreakdown } = {};
 
   spendData.results.forEach((day) => {
     Object.entries(day.breakdown.entities || {}).forEach(([entity, data]: [string, any]) => {
+      // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty)
+      const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown) || entity;
+      // Extract key_alias from the first API key that has one
+      const apiKeyBreakdown = data.api_key_breakdown || {};
+      let keyAlias: string | null = null;
+      for (const apiKeyData of Object.values(apiKeyBreakdown)) {
+        const alias = (apiKeyData as any)?.metadata?.key_alias;
+        if (alias) {
+          keyAlias = alias;
+          break;
+        }
+      }
+
       if (!entitySpend[entity]) {
         entitySpend[entity] = {
           metrics: {
@@ -22,8 +52,8 @@ export const getEntityBreakdown = (spendData: EntitySpendData): EntityBreakdown[
             cache_creation_input_tokens: 0,
           },
           metadata: {
-            alias: data.metadata?.team_alias || entity,
-            id: entity,
+            alias: keyAlias || teamAliasMap[teamId] || entity,
+            id: teamId,
           },
         };
       }
@@ -42,15 +72,23 @@ export const getEntityBreakdown = (spendData: EntitySpendData): EntityBreakdown[
   return Object.values(entitySpend).sort((a, b) => b.metrics.spend - a.metrics.spend);
 };
 
-export const generateDailyData = (spendData: EntitySpendData, entityLabel: string): any[] => {
+export const generateDailyData = (
+  spendData: EntitySpendData,
+  entityLabel: string,
+  teamAliasMap: Record<string, string> = {},
+): any[] => {
   const dailyBreakdown: any[] = [];
 
   spendData.results.forEach((day) => {
     Object.entries(day.breakdown.entities || {}).forEach(([entity, data]: [string, any]) => {
+      // Extract team_id from api_key_breakdown metadata (not data.metadata which is empty)
+      const teamId = extractTeamIdFromApiKeyBreakdown(data.api_key_breakdown);
+      const teamAlias = teamId ? teamAliasMap[teamId] || null : null;
+
       dailyBreakdown.push({
         Date: day.date,
-        [entityLabel]: data.metadata?.team_alias || entity,
-        [`${entityLabel} ID`]: entity,
+        [entityLabel]: teamAlias || "-",
+        [`${entityLabel} ID`]: teamId || "-",
         "Spend ($)": formatNumberWithCommas(data.metrics.spend, 4),
         Requests: data.metrics.api_requests,
         "Successful Requests": data.metrics.successful_requests,
@@ -65,15 +103,17 @@ export const generateDailyData = (spendData: EntitySpendData, entityLabel: strin
   return dailyBreakdown.sort((a, b) => new Date(a.Date).getTime() - new Date(b.Date).getTime());
 };
 
-export const generateDailyWithModelsData = (spendData: EntitySpendData, entityLabel: string): any[] => {
+export const generateDailyWithModelsData = (
+  spendData: EntitySpendData,
+  entityLabel: string,
+  teamAliasMap: Record<string, string> = {},
+): any[] => {
   const dailyModelBreakdown: any[] = [];
 
   spendData.results.forEach((day) => {
     const dailyEntityModels: { [key: string]: { [key: string]: any } } = {};
 
     Object.entries(day.breakdown.entities || {}).forEach(([entity, entityData]: [string, any]) => {
-      const entityName = entityData.metadata?.team_alias || entity;
-
       if (!dailyEntityModels[entity]) {
         dailyEntityModels[entity] = {};
       }
@@ -102,13 +142,15 @@ export const generateDailyWithModelsData = (spendData: EntitySpendData, entityLa
 
     Object.entries(dailyEntityModels).forEach(([entity, models]) => {
       const entityData = day.breakdown.entities?.[entity];
-      const entityName = entityData?.metadata?.team_alias || entity;
+      // Extract team_id from api_key_breakdown metadata (not entityData.metadata which is empty)
+      const teamId = extractTeamIdFromApiKeyBreakdown(entityData?.api_key_breakdown);
+      const teamAlias = teamId ? teamAliasMap[teamId] || null : null;
 
       Object.entries(models).forEach(([model, metrics]: [string, any]) => {
         dailyModelBreakdown.push({
           Date: day.date,
-          [entityLabel]: entityName,
-          [`${entityLabel} ID`]: entity,
+          [entityLabel]: teamAlias || "-",
+          [`${entityLabel} ID`]: teamId || "-",
           Model: model,
           "Spend ($)": formatNumberWithCommas(metrics.spend, 4),
           Requests: metrics.requests,
@@ -127,14 +169,15 @@ export const generateExportData = (
   spendData: EntitySpendData,
   exportScope: ExportScope,
   entityLabel: string,
+  teamAliasMap: Record<string, string> = {},
 ): any[] => {
   switch (exportScope) {
     case "daily":
-      return generateDailyData(spendData, entityLabel);
+      return generateDailyData(spendData, entityLabel, teamAliasMap);
     case "daily_with_models":
-      return generateDailyWithModelsData(spendData, entityLabel);
+      return generateDailyWithModelsData(spendData, entityLabel, teamAliasMap);
     default:
-      return generateDailyData(spendData, entityLabel);
+      return generateDailyData(spendData, entityLabel, teamAliasMap);
   }
 };
 
@@ -167,8 +210,9 @@ export const handleExportCSV = (
   exportScope: ExportScope,
   entityLabel: string,
   entityType: EntityType,
+  teamAliasMap: Record<string, string> = {},
 ): void => {
-  const data = generateExportData(spendData, exportScope, entityLabel);
+  const data = generateExportData(spendData, exportScope, entityLabel, teamAliasMap);
   const csv = Papa.unparse(data);
   const blob = new Blob([csv], { type: "text/csv;charset=utf-8;" });
   const url = window.URL.createObjectURL(blob);
@@ -189,8 +233,9 @@ export const handleExportJSON = (
   entityType: EntityType,
   dateRange: DateRangePickerValue,
   selectedFilters: string[],
+  teamAliasMap: Record<string, string> = {},
 ): void => {
-  const data = generateExportData(spendData, exportScope, entityLabel);
+  const data = generateExportData(spendData, exportScope, entityLabel, teamAliasMap);
   const metadata = generateMetadata(entityType, dateRange, selectedFilters, exportScope, spendData);
   const exportObject = {
     metadata,

ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx

@@ -36,7 +36,11 @@ vi.mock("./TopModelView", () => ({
   default: () => <div>Top Models</div>,
 }));
 
-vi.mock("./EntityUsageExport", () => ({
+vi.mock("../../../EntityUsageExport/EntityUsageExportModal", () => ({
+  default: () => <div>Entity Usage Export Modal</div>,
+}));
+
+vi.mock("../../../EntityUsageExport", () => ({
   UsageExportHeader: () => <div>Usage Export Header</div>,
 }));
 

ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx

@@ -399,6 +399,7 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
         selectedFilters={selectedTags}
         onFiltersChange={setSelectedTags}
         filterOptions={getAllTags() || undefined}
+        teams={teams || []}
       />
       <TabGroup>
         <NewBadge>