Merge pull request #17506 from BerriAI/litellm_ui_customer_usage

[Feature] Customer Usage UI
This commit is contained in:
yuneng-jiang
2025-12-05 16:18:42 -08:00
committed by GitHub
10 changed files with 173 additions and 12 deletions

View File

@@ -0,0 +1,41 @@
import { allEndUsersCall } from "@/components/networking";
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { all_admin_roles } from "@/utils/roles";
const customersKeys = createQueryKeys("customers");
export interface Customer {
user_id: string;
alias?: string | null;
spend: number;
blocked: boolean;
allowed_model_region?: string | null;
default_model?: string | null;
budget_id?: string | null;
litellm_budget_table?: {
budget_id: string;
max_budget?: number | null;
soft_budget?: number | null;
max_parallel_requests?: number | null;
tpm_limit?: number | null;
rpm_limit?: number | null;
model_max_budget?: Record<string, unknown> | null;
budget_duration?: string | null;
budget_reset_at?: string | null;
created_at: string;
created_by: string;
updated_at: string;
updated_by: string;
} | null;
}
export type CustomersResponse = Customer[];
export const useCustomers = (accessToken: string | null, userRole: string | null) => {
return useQuery<CustomersResponse>({
queryKey: customersKeys.list({}),
queryFn: async () => await allEndUsersCall(accessToken!),
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""),
});
};

View File

@@ -1,11 +1,11 @@
import React from "react";
import { Radio } from "antd";
import type { ExportScope } from "./types";
import type { ExportScope, EntityType } from "./types";
interface ExportTypeSelectorProps {
value: ExportScope;
onChange: (value: ExportScope) => void;
entityType: "tag" | "team" | "organization";
entityType: EntityType;
}
const ExportTypeSelector: React.FC<ExportTypeSelectorProps> = ({ value, onChange, entityType }) => {

View File

@@ -7,7 +7,7 @@ import type { EntitySpendData } from "./types";
interface UsageExportHeaderProps {
dateValue: DateRangePickerValue;
entityType: "tag" | "team" | "organization";
entityType: "tag" | "team" | "organization" | "customer";
spendData: EntitySpendData;
// Optional filter props
showFilters?: boolean;

View File

@@ -2,6 +2,7 @@ import type { DateRangePickerValue } from "@tremor/react";
export type ExportFormat = "csv" | "json";
export type ExportScope = "daily" | "daily_with_models";
export type EntityType = "tag" | "team" | "organization" | "customer";
export interface EntitySpendData {
results: any[];
@@ -17,7 +18,7 @@ export interface EntitySpendData {
export interface EntityUsageExportModalProps {
isOpen: boolean;
onClose: () => void;
entityType: "tag" | "team" | "organization";
entityType: EntityType;
spendData: EntitySpendData;
dateRange: DateRangePickerValue;
selectedFilters: string[];

View File

@@ -1,6 +1,6 @@
import { formatNumberWithCommas } from "@/utils/dataUtils";
import Papa from "papaparse";
import type { EntitySpendData, EntityBreakdown, ExportMetadata, ExportScope } from "./types";
import type { EntitySpendData, EntityBreakdown, ExportMetadata, ExportScope, EntityType } from "./types";
import type { DateRangePickerValue } from "@tremor/react";
export const getEntityBreakdown = (spendData: EntitySpendData): EntityBreakdown[] => {
@@ -139,7 +139,7 @@ export const generateExportData = (
};
export const generateMetadata = (
entityType: "tag" | "team" | "organization",
entityType: EntityType,
dateRange: DateRangePickerValue,
selectedFilters: string[],
exportScope: ExportScope,
@@ -166,7 +166,7 @@ export const handleExportCSV = (
spendData: EntitySpendData,
exportScope: ExportScope,
entityLabel: string,
entityType: "tag" | "team" | "organization",
entityType: EntityType,
): void => {
const data = generateExportData(spendData, exportScope, entityLabel);
const csv = Papa.unparse(data);
@@ -186,7 +186,7 @@ export const handleExportJSON = (
spendData: EntitySpendData,
exportScope: ExportScope,
entityLabel: string,
entityType: "tag" | "team" | "organization",
entityType: EntityType,
dateRange: DateRangePickerValue,
selectedFilters: string[],
): void => {

View File

@@ -18,6 +18,7 @@ vi.mock("./networking", () => ({
tagDailyActivityCall: vi.fn(),
teamDailyActivityCall: vi.fn(),
organizationDailyActivityCall: vi.fn(),
customerDailyActivityCall: vi.fn(),
}));
// Mock the child components to simplify testing
@@ -42,6 +43,7 @@ describe("EntityUsage", () => {
const mockTagDailyActivityCall = vi.mocked(networking.tagDailyActivityCall);
const mockTeamDailyActivityCall = vi.mocked(networking.teamDailyActivityCall);
const mockOrganizationDailyActivityCall = vi.mocked(networking.organizationDailyActivityCall);
const mockCustomerDailyActivityCall = vi.mocked(networking.customerDailyActivityCall);
const mockSpendData = {
results: [
@@ -128,9 +130,11 @@ describe("EntityUsage", () => {
mockTagDailyActivityCall.mockClear();
mockTeamDailyActivityCall.mockClear();
mockOrganizationDailyActivityCall.mockClear();
mockCustomerDailyActivityCall.mockClear();
mockTagDailyActivityCall.mockResolvedValue(mockSpendData);
mockTeamDailyActivityCall.mockResolvedValue(mockSpendData);
mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData);
mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData);
});
it("should render with tag entity type and display spend metrics", async () => {
@@ -182,6 +186,21 @@ describe("EntityUsage", () => {
});
});
it("should render with customer entity type and call customer API", async () => {
render(<EntityUsage {...defaultProps} entityType="customer" />);
await waitFor(() => {
expect(mockCustomerDailyActivityCall).toHaveBeenCalled();
});
expect(screen.getByText("Customer Spend Overview")).toBeInTheDocument();
await waitFor(() => {
const spendElements = screen.getAllByText("$100.50");
expect(spendElements.length).toBeGreaterThan(0);
});
});
it("should switch between tabs", async () => {
render(<EntityUsage {...defaultProps} />);

View File

@@ -23,12 +23,18 @@ import {
} from "@tremor/react";
import { ActivityMetrics, processActivityData } from "./activity_metrics";
import { DailyData, BreakdownMetrics, KeyMetricWithMetadata, EntityMetricWithMetadata, TagUsage } from "./usage/types";
import { organizationDailyActivityCall, tagDailyActivityCall, teamDailyActivityCall } from "./networking";
import {
organizationDailyActivityCall,
tagDailyActivityCall,
teamDailyActivityCall,
customerDailyActivityCall,
} from "./networking";
import TopKeyView from "./top_key_view";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { valueFormatterSpend } from "./usage/utils/value_formatters";
import { getProviderLogoAndName } from "./provider_info_helpers";
import { UsageExportHeader } from "./EntityUsageExport";
import type { EntityType } from "./EntityUsageExport/types";
import TopModelView from "./top_model_view";
interface EntityMetrics {
@@ -68,7 +74,7 @@ export interface EntityList {
interface EntityUsageProps {
accessToken: string | null;
entityType: "tag" | "team" | "organization";
entityType: EntityType;
entityId?: string | null;
userID: string | null;
userRole: string | null;
@@ -135,6 +141,15 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
selectedTags.length > 0 ? selectedTags : null,
);
setSpendData(data);
} else if (entityType === "customer") {
const data = await customerDailyActivityCall(
accessToken,
startTime,
endTime,
1,
selectedTags.length > 0 ? selectedTags : null,
);
setSpendData(data);
} else {
throw new Error("Invalid entity type");
}

View File

@@ -1736,6 +1736,25 @@ export const organizationDailyActivityCall = async (
});
};
export const customerDailyActivityCall = async (
accessToken: string,
startTime: Date,
endTime: Date,
page: number = 1,
customerIds: string[] | null = null,
) => {
return fetchDailyActivity({
accessToken,
endpoint: "/customer/daily/activity",
startTime,
endTime,
page,
extraQueryParams: {
end_user_ids: customerIds,
},
});
};
export const getTotalSpendCall = async (accessToken: string) => {
/**
* Get all models on proxy
@@ -2511,7 +2530,7 @@ export const allEndUsersCall = async (accessToken: string) => {
console.log(data);
return data;
} catch (error) {
console.error("Failed to create key:", error);
console.error("Failed to fetch end users:", error);
throw error;
}
};

View File

@@ -3,6 +3,7 @@ import { describe, it, expect, vi, beforeEach, beforeAll } from "vitest";
import NewUsagePage from "./new_usage";
import type { Organization } from "./networking";
import * as networking from "./networking";
import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers";
// Polyfill ResizeObserver for test environment
beforeAll(() => {
@@ -53,9 +54,14 @@ vi.mock("./EntityUsageExport", () => ({
default: () => <div>Entity Usage Export Modal</div>,
}));
vi.mock("@/app/(dashboard)/hooks/customers/useCustomers", () => ({
useCustomers: vi.fn(),
}));
describe("NewUsage", () => {
const mockUserDailyActivityAggregatedCall = vi.mocked(networking.userDailyActivityAggregatedCall);
const mockTagListCall = vi.mocked(networking.tagListCall);
const mockUseCustomers = vi.mocked(useCustomers);
const mockSpendData = {
results: [
@@ -174,6 +180,19 @@ describe("NewUsage", () => {
},
];
const mockCustomers = [
{
user_id: "customer-123",
alias: "Test Customer",
spend: 0,
blocked: false,
allowed_model_region: null,
default_model: null,
budget_id: null,
litellm_budget_table: null,
},
];
const defaultProps = {
accessToken: "test-token",
userRole: "Admin",
@@ -205,6 +224,11 @@ describe("NewUsage", () => {
mockTagListCall.mockClear();
mockUserDailyActivityAggregatedCall.mockResolvedValue(mockSpendData);
mockTagListCall.mockResolvedValue({});
mockUseCustomers.mockReturnValue({
data: [],
isLoading: false,
error: null,
} as any);
});
it("should render and fetch usage data on mount", async () => {
@@ -289,4 +313,26 @@ describe("NewUsage", () => {
expect(entityUsageElements.length).toBeGreaterThan(0);
});
});
it("should show customer usage tab for admins", async () => {
mockUseCustomers.mockReturnValue({
data: mockCustomers,
isLoading: false,
error: null,
} as any);
const { getByText, getAllByText } = render(<NewUsagePage {...defaultProps} />);
await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});
const customerTab = getByText("Customer Usage");
fireEvent.click(customerTab);
await waitFor(() => {
const entityUsageElements = getAllByText("Entity Usage");
expect(entityUsageElements.length).toBeGreaterThan(0);
});
});
});

View File

@@ -27,9 +27,10 @@ import {
Text,
Title,
} from "@tremor/react";
import React, { useCallback, useEffect, useMemo, useState } from "react";
import { Alert } from "antd";
import React, { useCallback, useEffect, useMemo, useState } from "react";
import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers";
import { formatNumberWithCommas } from "@/utils/dataUtils";
import { Button } from "@tremor/react";
import { all_admin_roles } from "../utils/roles";
@@ -86,6 +87,7 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({
});
const [allTags, setAllTags] = useState<EntityList[]>([]);
const { data: customers = [] } = useCustomers(accessToken, userRole);
const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups");
const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false);
const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false);
@@ -430,6 +432,7 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({
<Tab>Your Organization Usage</Tab>
)}
<Tab>Team Usage</Tab>
{all_admin_roles.includes(userRole || "") ? <Tab>Customer Usage</Tab> : <></>}
{all_admin_roles.includes(userRole || "") ? <Tab>Tag Usage</Tab> : <></>}
{all_admin_roles.includes(userRole || "") ? <Tab>User Agent Activity</Tab> : <></>}
</TabList>
@@ -798,6 +801,23 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({
/>
</TabPanel>
{/* Customer Usage Panel */}
<TabPanel>
<EntityUsage
accessToken={accessToken}
entityType="customer"
userID={userID}
userRole={userRole}
entityList={
customers?.map((customer) => ({
label: customer.alias || customer.user_id,
value: customer.user_id,
})) || null
}
premiumUser={premiumUser}
dateValue={dateValue}
/>
</TabPanel>
{/* Tag Usage Panel */}
<TabPanel>
<EntityUsage