From 87e915e4e22e44c57077133383daf7600864ece0 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 19 Dec 2025 16:43:52 -0800 Subject: [PATCH] Resolve team alias in usage page --- .../EntityUsage/EntityUsage.test.tsx | 8 + .../components/EntityUsage/EntityUsage.tsx | 6 +- .../UsagePage/components/UsagePageView.tsx | 6 +- .../src/components/activity_metrics.test.tsx | 213 +++++++++++++++++- .../src/components/activity_metrics.tsx | 13 +- .../src/utils/teamUtils.test.ts | 37 +++ ui/litellm-dashboard/src/utils/teamUtils.ts | 6 + 7 files changed, 279 insertions(+), 10 deletions(-) create mode 100644 ui/litellm-dashboard/src/utils/teamUtils.test.ts create mode 100644 ui/litellm-dashboard/src/utils/teamUtils.ts diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx index 9a9fa68f762..46714857c1f 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.test.tsx @@ -40,6 +40,14 @@ vi.mock("./EntityUsageExport", () => ({ UsageExportHeader: () =>
Usage Export Header
, })); +// Mock useTeams hook +vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ + default: vi.fn(() => ({ + teams: [], + setTeams: vi.fn(), + })), +})); + describe("EntityUsage", () => { const mockTagDailyActivityCall = vi.mocked(networking.tagDailyActivityCall); const mockTeamDailyActivityCall = vi.mocked(networking.teamDailyActivityCall); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx index 8cb318b06b2..75cf8c9292c 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/EntityUsage.tsx @@ -37,6 +37,7 @@ import { BreakdownMetrics, DailyData, EntityMetricWithMetadata, KeyMetricWithMet import { valueFormatterSpend } from "../../utils/value_formatters"; import TopKeyView from "./TopKeyView"; import TopModelView from "./TopModelView"; +import useTeams from "@/app/(dashboard)/hooks/useTeams"; interface EntityMetrics { metrics: { @@ -104,9 +105,10 @@ const EntityUsage: React.FC = ({ total_tokens: 0, }, }); + const { teams } = useTeams(); - const modelMetrics = processActivityData(spendData, "models"); - const keyMetrics = processActivityData(spendData, "api_keys"); + const modelMetrics = processActivityData(spendData, "models", teams || []); + const keyMetrics = processActivityData(spendData, "api_keys", teams || []); const [selectedTags, setSelectedTags] = useState([]); const fetchSpendData = async () => { diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 7f4f98cbc51..9766983c36a 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -371,9 +371,9 @@ const UsagePage: React.FC = ({ teams, organizations }) => { return () => clearTimeout(timeoutId); }, [fetchUserSpendData]); - const modelMetrics = processActivityData(userSpendData, "models"); - const keyMetrics = processActivityData(userSpendData, "api_keys"); - const mcpServerMetrics = processActivityData(userSpendData, "mcp_servers"); + const modelMetrics = processActivityData(userSpendData, "models", teams); + const keyMetrics = processActivityData(userSpendData, "api_keys", teams); + const mcpServerMetrics = processActivityData(userSpendData, "mcp_servers", teams); return (
diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index 85bcabcfda3..afe6f303bd3 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -1,8 +1,9 @@ import { render, screen } from "@testing-library/react"; import React from "react"; import { beforeAll, describe, expect, it, vi } from "vitest"; -import { ActivityMetrics } from "./activity_metrics"; -import { ModelActivityData } from "./UsagePage/types"; +import { ActivityMetrics, processActivityData, formatKeyLabel } from "./activity_metrics"; +import { ModelActivityData, DailyData, KeyMetricWithMetadata } from "./UsagePage/types"; +import { Team } from "./key_team_helpers/key_list"; beforeAll(() => { if (typeof window !== "undefined" && !window.ResizeObserver) { @@ -101,3 +102,211 @@ describe("ActivityMetrics", () => { expect(screen.queryByText("Prompt Caching Metrics")).not.toBeInTheDocument(); }); }); + +describe("processActivityData", () => { + const mockDailyActivity: { results: DailyData[] } = { + results: [ + { + date: "2025-01-01", + metrics: { + spend: 100.5, + prompt_tokens: 30000, + completion_tokens: 20000, + total_tokens: 50000, + api_requests: 100, + successful_requests: 95, + failed_requests: 5, + cache_read_input_tokens: 1000, + cache_creation_input_tokens: 500, + }, + breakdown: { + models: {}, + model_groups: {}, + mcp_servers: {}, + providers: {}, + api_keys: { + key1: { + metrics: { + spend: 50.25, + prompt_tokens: 15000, + completion_tokens: 10000, + total_tokens: 25000, + api_requests: 50, + successful_requests: 47, + failed_requests: 3, + cache_read_input_tokens: 500, + cache_creation_input_tokens: 250, + }, + metadata: { + key_alias: "test-key-1", + team_id: "team1", + }, + }, + }, + entities: {}, + }, + }, + ], + }; + + const mockTeams: Team[] = [ + { + team_id: "team1", + team_alias: "Test Team 1", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org1", + created_at: "2025-01-01", + keys: [], + members_with_roles: [], + }, + ]; + + it("should process data for models key without teams parameter", () => { + const result = processActivityData(mockDailyActivity, "models"); + + expect(result).toEqual({}); + }); + + it("should process data for api_keys key with teams parameter", () => { + const result = processActivityData(mockDailyActivity, "api_keys", mockTeams); + + expect(result).toHaveProperty("key1"); + expect(result["key1"].label).toBe("test-key-1 (team: Test Team 1)"); + expect(result["key1"].total_requests).toBe(50); + expect(result["key1"].total_spend).toBe(50.25); + }); + + it("should process data for api_keys key without teams parameter", () => { + const result = processActivityData(mockDailyActivity, "api_keys"); + + expect(result).toHaveProperty("key1"); + expect(result["key1"].label).toBe("test-key-1 (team_id: team1)"); + }); +}); + +describe("formatKeyLabel", () => { + const mockTeams: Team[] = [ + { + team_id: "team1", + team_alias: "Test Team 1", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org1", + created_at: "2025-01-01", + keys: [], + members_with_roles: [], + }, + { + team_id: "team2", + team_alias: "Test Team 2", + models: [], + max_budget: null, + budget_duration: null, + tpm_limit: null, + rpm_limit: null, + organization_id: "org2", + created_at: "2025-01-01", + keys: [], + members_with_roles: [], + }, + ]; + + it("should return key_alias when no team_id is present", () => { + const modelData: KeyMetricWithMetadata = { + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: "test-key", + team_id: null, + }, + }; + + const result = formatKeyLabel(modelData, "test-key", mockTeams); + expect(result).toBe("test-key"); + }); + + it("should return key_alias with team alias when team_id matches", () => { + const modelData: KeyMetricWithMetadata = { + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: "test-key", + team_id: "team1", + }, + }; + + const result = formatKeyLabel(modelData, "test-key", mockTeams); + expect(result).toBe("test-key (team: Test Team 1)"); + }); + + it("should return key_alias with team_id when team is not found", () => { + const modelData: KeyMetricWithMetadata = { + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: "test-key", + team_id: "nonexistent-team", + }, + }; + + const result = formatKeyLabel(modelData, "test-key", mockTeams); + expect(result).toBe("test-key (team_id: nonexistent-team)"); + }); + + it("should use key-hash fallback when key_alias is null", () => { + const modelData: KeyMetricWithMetadata = { + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + metadata: { + key_alias: null, + team_id: "team1", + }, + }; + + const result = formatKeyLabel(modelData, "actual-key", mockTeams); + expect(result).toBe("key-hash-actual-key (team: Test Team 1)"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index 7a387c26600..51a7a803fb6 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -5,6 +5,8 @@ import React from "react"; import { CustomLegend, CustomTooltip } from "./common_components/chartUtils"; import { DailyData, KeyMetricWithMetadata, ModelActivityData, TopApiKeyData } from "./UsagePage/types"; import { valueFormatter } from "./UsagePage/utils/value_formatters"; +import { Team } from "./key_team_helpers/key_list"; +import { resolveTeamAliasFromTeamID } from "@/utils/teamUtils"; interface ActivityMetricsProps { modelMetrics: Record; @@ -337,16 +339,21 @@ export const ActivityMetrics: React.FC = ({ modelMetrics, }; // Helper function to format key label -const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string): string => { +export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, teams: Team[]): string => { const keyAlias = modelData.metadata.key_alias || `key-hash-${model}`; const teamId = modelData.metadata.team_id; - return teamId ? `${keyAlias} (team_id: ${teamId})` : keyAlias; + if (teamId) { + const teamAlias = resolveTeamAliasFromTeamID(teamId, teams); + return teamAlias ? `${keyAlias} (team: ${teamAlias})` : `${keyAlias} (team_id: ${teamId})`; + } + return keyAlias; }; // Process data function export const processActivityData = ( dailyActivity: { results: DailyData[] }, key: "models" | "api_keys" | "mcp_servers", + teams: Team[] = [], ): Record => { const modelMetrics: Record = {}; @@ -354,7 +361,7 @@ export const processActivityData = ( Object.entries(day.breakdown[key] || {}).forEach(([model, modelData]) => { if (!modelMetrics[model]) { modelMetrics[model] = { - label: key === "api_keys" ? formatKeyLabel(modelData as KeyMetricWithMetadata, model) : model, + label: key === "api_keys" ? formatKeyLabel(modelData as KeyMetricWithMetadata, model, teams) : model, total_requests: 0, total_successful_requests: 0, total_failed_requests: 0, diff --git a/ui/litellm-dashboard/src/utils/teamUtils.test.ts b/ui/litellm-dashboard/src/utils/teamUtils.test.ts new file mode 100644 index 00000000000..1151d532611 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/teamUtils.test.ts @@ -0,0 +1,37 @@ +import { describe, expect, it } from "vitest"; +import { resolveTeamAliasFromTeamID } from "./teamUtils"; +import type { Team } from "@/components/networking"; + +describe("resolveTeamAliasFromTeamID", () => { + it("should return team alias when team is found", () => { + const teams = [ + { + team_id: "team1", + team_alias: "Team One", + }, + { + team_id: "team2", + team_alias: "Team Two", + }, + ] as unknown as Team[]; + + const result = resolveTeamAliasFromTeamID("team1", teams); + expect(result).toBe("Team One"); + }); + + it("should return null when team is not found", () => { + const teams = [ + { + team_id: "team1", + team_alias: "Team One", + }, + { + team_id: "team2", + team_alias: "Team Two", + }, + ] as unknown as Team[]; + + const result = resolveTeamAliasFromTeamID("team3", teams); + expect(result).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/teamUtils.ts b/ui/litellm-dashboard/src/utils/teamUtils.ts new file mode 100644 index 00000000000..1916e1c98aa --- /dev/null +++ b/ui/litellm-dashboard/src/utils/teamUtils.ts @@ -0,0 +1,6 @@ +import { Team } from "@/components/networking"; + +export const resolveTeamAliasFromTeamID = (teamID: string, teams: Team[]): string | null => { + const team = teams.find((team) => team.team_id === teamID); + return team ? team.team_alias : null; +};