Skip to content

fix(ui): use entity key for usage export display #25153

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ryan-crabbe-berri merged 2 commits into from
Apr 4, 2026
Merged

fix(ui): use entity key for usage export display #25153

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9dca431
fix(ui): use entity key for export instead of extracting team_id from…
ryan-crabbe-berri Mar 31, 2026
5452692
test(ui): align EntityUsageExport tests with entity-key display logic
ryan-crabbe-berri Apr 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 26 additions & 18 deletions ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,17 @@ vi.mock("papaparse", () => ({
}));

describe("EntityUsageExport utils", () => {
// Entity keys match team_ids because that's how the backend shapes team exports
// (breakdown.entities is keyed by team_id). The fix under test uses the entity key
// directly for display, so the key_alias/team_id in api_key_breakdown metadata is
// no longer consulted — it's retained here only to mirror real payload shape.
const mockSpendData: EntitySpendData = {
results: [
{
date: "2025-01-01",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 10.5,
api_requests: 100,
Expand Down Expand Up @@ -64,7 +68,7 @@ describe("EntityUsageExport utils", () => {
},
},
},
entity2: {
"team-2": {
metrics: {
spend: 20.3,
api_requests: 200,
Expand Down Expand Up @@ -99,7 +103,7 @@ describe("EntityUsageExport utils", () => {
date: "2025-01-02",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 15.2,
api_requests: 150,
Expand Down Expand Up @@ -184,22 +188,24 @@ describe("EntityUsageExport utils", () => {
expect(entity1?.metrics.cache_creation_input_tokens).toBe(75);
});

it("should use key alias when available", () => {
it("should use entity key as alias when no team alias map is provided", () => {
// Non-team exports (tags, orgs, customers, …) pass no teamAliasMap.
// For teams, this is also the fallback when a team is missing from the map.
const result = getEntityBreakdown(mockSpendData);
const entity1 = result.find((e) => e.metadata.id === "team-1");

expect(entity1?.metadata.alias).toBe("alias-1");
expect(entity1?.metadata.alias).toBe("team-1");
});

it("should use team alias map when key alias is not available", () => {
it("should use team alias map to resolve alias from entity key", () => {
const spendDataWithoutAlias: EntitySpendData = {
...mockSpendData,
results: [
{
date: "2025-01-01",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 10.5,
api_requests: 100,
Expand Down Expand Up @@ -299,7 +305,7 @@ describe("EntityUsageExport utils", () => {
date: "2025-01-01",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 10.5,
api_requests: 100,
Expand Down Expand Up @@ -379,15 +385,17 @@ describe("EntityUsageExport utils", () => {
}
});

it("should use dash when team id is not available", () => {
const spendDataWithoutTeamId: EntitySpendData = {
it("should fall back to the entity key when there is no team alias mapping", () => {
// e.g. tag/org/customer exports where teamAliasMap has no entry for the entity,
// or a team that isn't in the alias map — the entity key itself is the label.
const spendDataWithoutAlias: EntitySpendData = {
...mockSpendData,
results: [
{
date: "2025-01-01",
breakdown: {
entities: {
entity1: {
"my-tag": {
metrics: {
spend: 10.5,
api_requests: 100,
Expand All @@ -406,11 +414,11 @@ describe("EntityUsageExport utils", () => {
metadata: mockSpendData.metadata,
};

const result = generateDailyData(spendDataWithoutTeamId, "Team");
const result = generateDailyData(spendDataWithoutAlias, "Tag");
const entry = result[0];

expect(entry["Team ID"]).toBe("-");
expect(entry["Team"]).toBe("-");
expect(entry["Tag ID"]).toBe("my-tag");
expect(entry["Tag"]).toBe("my-tag");
});

it("should format spend values correctly", () => {
Expand Down Expand Up @@ -471,7 +479,7 @@ describe("EntityUsageExport utils", () => {
date: "2025-01-01",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 10.5,
api_requests: 100,
Expand Down Expand Up @@ -514,7 +522,7 @@ describe("EntityUsageExport utils", () => {
},
},
},
entity2: {
"team-2": {
metrics: {
spend: 20.3,
api_requests: 200,
Expand Down Expand Up @@ -549,7 +557,7 @@ describe("EntityUsageExport utils", () => {
date: "2025-01-02",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 15.2,
api_requests: 150,
Expand Down Expand Up @@ -979,7 +987,7 @@ describe("EntityUsageExport utils", () => {
date: "2025-01-01",
breakdown: {
entities: {
entity1: {
"team-1": {
metrics: {
spend: 10.5,
api_requests: 100,
Expand Down
80 changes: 30 additions & 50 deletions ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,16 @@ import type { DateRangePickerValue } from "@tremor/react";
import Papa from "papaparse";
import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types";

// 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;
};
// Resolve display name for an entity. For teams the teamAliasMap provides
// a human-readable alias; for every other entity type the entity key itself
// (tag name, org id, customer id, …) is already the correct label.
const resolveEntityDisplay = (
entity: string,
teamAliasMap: Record<string, string>,
): { id: string; alias: string } => ({
id: entity,
alias: teamAliasMap[entity] || entity,
});

// Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py).
// If the backend adds a field, add it here too.
Expand Down Expand Up @@ -68,18 +65,7 @@ export const getEntityBreakdown = (

spendData.results.forEach((day) => {
Object.entries(resolveEntities(day.breakdown)).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;
}
}
const { id, alias } = resolveEntityDisplay(entity, teamAliasMap);

if (!entitySpend[entity]) {
entitySpend[entity] = {
Expand All @@ -95,8 +81,8 @@ export const getEntityBreakdown = (
cache_creation_input_tokens: 0,
},
metadata: {
alias: keyAlias || teamAliasMap[teamId] || entity,
id: teamId,
alias,
id,
},
};
}
Expand Down Expand Up @@ -124,14 +110,12 @@ export const generateDailyData = (

spendData.results.forEach((day) => {
Object.entries(resolveEntities(day.breakdown)).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;
const { id, alias } = resolveEntityDisplay(entity, teamAliasMap);

dailyBreakdown.push({
Date: day.date,
[entityLabel]: teamAlias || "-",
[`${entityLabel} ID`]: teamId || "-",
[entityLabel]: alias,
[`${entityLabel} ID`]: id,
"Spend ($)": formatNumberWithCommas(data.metrics.spend, 4),
Requests: data.metrics.api_requests,
"Successful Requests": data.metrics.successful_requests,
Expand All @@ -151,12 +135,12 @@ export const generateDailyWithKeysData = (
entityLabel: string,
teamAliasMap: Record<string, string> = {},
): any[] => {
// Aggregate by unique (Date, Team ID, Key ID) combination to prevent duplicates
// Aggregate by unique (Date, Entity ID, Key ID) combination to prevent duplicates
const aggregatedData: {
[key: string]: {
Date: string;
teamId: string;
teamAlias: string | null;
entityId: string;
entityAlias: string;
keyId: string;
keyAlias: string | null;
metrics: {
Expand All @@ -173,23 +157,22 @@ export const generateDailyWithKeysData = (

spendData.results.forEach((day) => {
Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => {
const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap);
const apiKeyBreakdown = data.api_key_breakdown || {};

// Iterate through each API key in the breakdown
Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => {
const keyAlias = keyData?.metadata?.key_alias || null;
const teamId = keyData?.metadata?.team_id || entity;
const teamAlias = teamId ? teamAliasMap[teamId] || null : null;

// Create unique key for aggregation: Date_TeamID_KeyID
const uniqueKey = `${day.date}_${teamId}_${keyId}`;
// Create unique key for aggregation: Date_EntityID_KeyID
const uniqueKey = `${day.date}_${entityId}_${keyId}`;

if (!aggregatedData[uniqueKey]) {
// First time seeing this (Date, Team ID, Key ID) combination
// First time seeing this (Date, Entity ID, Key ID) combination
aggregatedData[uniqueKey] = {
Date: day.date,
teamId,
teamAlias,
entityId,
entityAlias,
keyId,
keyAlias,
metrics: {
Expand Down Expand Up @@ -219,8 +202,8 @@ export const generateDailyWithKeysData = (
// Convert aggregated data to array format
const dailyKeyBreakdown = Object.values(aggregatedData).map((item) => ({
Date: item.Date,
[entityLabel]: item.teamAlias || "-",
[`${entityLabel} ID`]: item.teamId || "-",
[entityLabel]: item.entityAlias,
[`${entityLabel} ID`]: item.entityId,
"Key Alias": item.keyAlias || "-",
"Key ID": item.keyId,
"Spend ($)": formatNumberWithCommas(item.metrics.spend, 4),
Expand Down Expand Up @@ -273,16 +256,13 @@ export const generateDailyWithModelsData = (
});

Object.entries(dailyEntityModels).forEach(([entity, models]) => {
const entityData = resolveEntities(day.breakdown)[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;
const { id, alias } = resolveEntityDisplay(entity, teamAliasMap);

Object.entries(models).forEach(([model, metrics]: [string, any]) => {
dailyModelBreakdown.push({
Date: day.date,
[entityLabel]: teamAlias || "-",
[`${entityLabel} ID`]: teamId || "-",
[entityLabel]: alias,
[`${entityLabel} ID`]: id,
Model: model,
"Spend ($)": formatNumberWithCommas(metrics.spend, 4),
Requests: metrics.requests,
Expand Down
Loading