Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
15 changes: 15 additions & 0 deletions ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
import { getAgentsList } from "@/components/networking";
import { AgentsResponse } from "@/components/agents/types";
import { useQuery } from "@tanstack/react-query";
import { createQueryKeys } from "../common/queryKeysFactory";
import { all_admin_roles } from "@/utils/roles";

const agentsKeys = createQueryKeys("agents");

export const useAgents = (accessToken: string | null, userRole: string | null) => {
return useQuery<AgentsResponse>({
queryKey: agentsKeys.list({}),
queryFn: async () => await getAgentsList(accessToken!),
enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""),
});
};
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
import React, { useState } from "react";
import type { DateRangePickerValue } from "@tremor/react";
import { Button, Text } from "@tremor/react";
import { Select } from "antd";
import React, { useState } from "react";
import EntityUsageExportModal from "./EntityUsageExportModal";
import type { DateRangePickerValue } from "@tremor/react";
import type { EntitySpendData } from "./types";
import type { EntitySpendData, EntityType } from "./types";

interface UsageExportHeaderProps {
dateValue: DateRangePickerValue;
entityType: "tag" | "team" | "organization" | "customer";
entityType: EntityType;
spendData: EntitySpendData;
// Optional filter props
showFilters?: boolean;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +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 type EntityType = "tag" | "team" | "organization" | "customer" | "agent";

export interface EntitySpendData {
results: any[];
Expand Down
19 changes: 19 additions & 0 deletions ui/litellm-dashboard/src/components/entity_usage.test.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ vi.mock("./networking", () => ({
teamDailyActivityCall: vi.fn(),
organizationDailyActivityCall: vi.fn(),
customerDailyActivityCall: vi.fn(),
agentDailyActivityCall: vi.fn(),
}));

// Mock the child components to simplify testing
Expand All @@ -44,6 +45,7 @@ describe("EntityUsage", () => {
const mockTeamDailyActivityCall = vi.mocked(networking.teamDailyActivityCall);
const mockOrganizationDailyActivityCall = vi.mocked(networking.organizationDailyActivityCall);
const mockCustomerDailyActivityCall = vi.mocked(networking.customerDailyActivityCall);
const mockAgentDailyActivityCall = vi.mocked(networking.agentDailyActivityCall);

const mockSpendData = {
results: [
Expand Down Expand Up @@ -131,10 +133,12 @@ describe("EntityUsage", () => {
mockTeamDailyActivityCall.mockClear();
mockOrganizationDailyActivityCall.mockClear();
mockCustomerDailyActivityCall.mockClear();
mockAgentDailyActivityCall.mockClear();
mockTagDailyActivityCall.mockResolvedValue(mockSpendData);
mockTeamDailyActivityCall.mockResolvedValue(mockSpendData);
mockOrganizationDailyActivityCall.mockResolvedValue(mockSpendData);
mockCustomerDailyActivityCall.mockResolvedValue(mockSpendData);
mockAgentDailyActivityCall.mockResolvedValue(mockSpendData);
});

it("should render with tag entity type and display spend metrics", async () => {
Expand Down Expand Up @@ -201,6 +205,21 @@ describe("EntityUsage", () => {
});
});

it("should render with agent entity type and call agent API", async () => {
render(<EntityUsage {...defaultProps} entityType="agent" />);

await waitFor(() => {
expect(mockAgentDailyActivityCall).toHaveBeenCalled();
});

expect(screen.getByText("Agent 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} />);

Expand Down
10 changes: 10 additions & 0 deletions ui/litellm-dashboard/src/components/entity_usage.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ import {
tagDailyActivityCall,
teamDailyActivityCall,
customerDailyActivityCall,
agentDailyActivityCall,
} from "./networking";
import TopKeyView from "./top_key_view";
import { formatNumberWithCommas } from "@/utils/dataUtils";
Expand Down Expand Up @@ -150,6 +151,15 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
selectedTags.length > 0 ? selectedTags : null,
);
setSpendData(data);
} else if (entityType === "agent") {
const data = await agentDailyActivityCall(
accessToken,
startTime,
endTime,
1,
selectedTags.length > 0 ? selectedTags : null,
);
setSpendData(data);
} else {
throw new Error("Invalid entity type");
}
Expand Down
19 changes: 19 additions & 0 deletions ui/litellm-dashboard/src/components/networking.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -1754,6 +1754,25 @@ export const customerDailyActivityCall = async (
});
};

export const agentDailyActivityCall = async (
accessToken: string,
startTime: Date,
endTime: Date,
page: number = 1,
agentIds: string[] | null = null,
) => {
return fetchDailyActivity({
accessToken,
endpoint: "/agent/daily/activity",
startTime,
endTime,
page,
extraQueryParams: {
agent_ids: agentIds,
},
});
};

export const getTotalSpendCall = async (accessToken: string) => {
/**
* Get all models on proxy
Expand Down
74 changes: 62 additions & 12 deletions ui/litellm-dashboard/src/components/new_usage.test.tsx
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import { render, screen, fireEvent, waitFor } from "@testing-library/react";
import { render, screen, fireEvent, waitFor, act } from "@testing-library/react";
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";
import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents";

// Polyfill ResizeObserver for test environment
beforeAll(() => {
Expand Down Expand Up @@ -58,10 +59,15 @@ vi.mock("@/app/(dashboard)/hooks/customers/useCustomers", () => ({
useCustomers: vi.fn(),
}));

vi.mock("@/app/(dashboard)/hooks/agents/useAgents", () => ({
useAgents: vi.fn(),
}));

describe("NewUsage", () => {
const mockUserDailyActivityAggregatedCall = vi.mocked(networking.userDailyActivityAggregatedCall);
const mockTagListCall = vi.mocked(networking.tagListCall);
const mockUseCustomers = vi.mocked(useCustomers);
const mockUseAgents = vi.mocked(useAgents);

const mockSpendData = {
results: [
Expand Down Expand Up @@ -193,6 +199,13 @@ describe("NewUsage", () => {
},
];

const mockAgents = [
{
agent_id: "agent-123",
agent_name: "Test Agent",
},
];

const defaultProps = {
accessToken: "test-token",
userRole: "Admin",
Expand Down Expand Up @@ -229,6 +242,11 @@ describe("NewUsage", () => {
isLoading: false,
error: null,
} as any);
mockUseAgents.mockReturnValue({
data: { agents: [] },
isLoading: false,
error: null,
} as any);
});

it("should render and fetch usage data on mount", async () => {
Expand Down Expand Up @@ -278,7 +296,9 @@ describe("NewUsage", () => {

// Switch to Team Usage tab
const teamUsageTab = screen.getByText("Team Usage");
fireEvent.click(teamUsageTab);
act(() => {
fireEvent.click(teamUsageTab);
});

// Should render EntityUsage component
await waitFor(() => {
Expand All @@ -288,7 +308,9 @@ describe("NewUsage", () => {

// Switch to Tag Usage tab (admin only)
const tagUsageTab = screen.getByText("Tag Usage");
fireEvent.click(tagUsageTab);
act(() => {
fireEvent.click(tagUsageTab);
});

// Should still render EntityUsage component for tags
await waitFor(() => {
Expand All @@ -298,18 +320,20 @@ describe("NewUsage", () => {
});

it("should show organization usage banner and tab for admins", async () => {
const { getByText, getAllByText } = render(<NewUsagePage {...defaultProps} organizations={mockOrganizations} />);
render(<NewUsagePage {...defaultProps} organizations={mockOrganizations} />);

await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});

const organizationTab = getByText("Organization Usage");
fireEvent.click(organizationTab);
const organizationTab = screen.getByText("Organization Usage");
act(() => {
fireEvent.click(organizationTab);
});

await waitFor(() => {
expect(getByText("Organization usage is a new feature.")).toBeInTheDocument();
const entityUsageElements = getAllByText("Entity Usage");
expect(screen.getByText("Organization usage is a new feature.")).toBeInTheDocument();
const entityUsageElements = screen.getAllByText("Entity Usage");
expect(entityUsageElements.length).toBeGreaterThan(0);
});
});
Expand All @@ -321,17 +345,43 @@ describe("NewUsage", () => {
error: null,
} as any);

const { getByText, getAllByText } = render(<NewUsagePage {...defaultProps} />);
render(<NewUsagePage {...defaultProps} />);

await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});

const customerTab = screen.getByText("Customer Usage");
act(() => {
fireEvent.click(customerTab);
});

await waitFor(() => {
const entityUsageElements = screen.getAllByText("Entity Usage");
expect(entityUsageElements.length).toBeGreaterThan(0);
});
});

it("should show agent usage tab for admins", async () => {
mockUseAgents.mockReturnValue({
data: { agents: mockAgents },
isLoading: false,
error: null,
} as any);

render(<NewUsagePage {...defaultProps} />);

await waitFor(() => {
expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
});

const customerTab = getByText("Customer Usage");
fireEvent.click(customerTab);
const agentTab = screen.getByText("Agent Usage");
act(() => {
fireEvent.click(agentTab);
});

await waitFor(() => {
const entityUsageElements = getAllByText("Entity Usage");
const entityUsageElements = screen.getAllByText("Entity Usage");
expect(entityUsageElements.length).toBeGreaterThan(0);
});
});
Expand Down
16 changes: 16 additions & 0 deletions ui/litellm-dashboard/src/components/new_usage.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ import { DailyData, KeyMetricWithMetadata, MetricWithMetadata } from "./usage/ty
import { valueFormatterSpend } from "./usage/utils/value_formatters";
import UserAgentActivity from "./user_agent_activity";
import ViewUserSpend from "./view_user_spend";
import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents";

interface NewUsagePageProps {
accessToken: string | null;
Expand Down Expand Up @@ -88,6 +89,7 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({

const [allTags, setAllTags] = useState<EntityList[]>([]);
const { data: customers = [] } = useCustomers(accessToken, userRole);
const { data: agentsResponse } = useAgents(accessToken, userRole);
const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups");
const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false);
const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false);
Expand Down Expand Up @@ -435,6 +437,7 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({
<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>Agent Usage</Tab> : <></>}
{all_admin_roles.includes(userRole || "") ? <Tab>User Agent Activity</Tab> : <></>}
</TabList>
<AdvancedDatePicker value={dateValue} onValueChange={handleDateChange} />
Expand Down Expand Up @@ -842,6 +845,19 @@ const NewUsagePage: React.FC<NewUsagePageProps> = ({
dateValue={dateValue}
/>
</TabPanel>
<TabPanel>
<EntityUsage
accessToken={accessToken}
entityType="agent"
userID={userID}
userRole={userRole}
entityList={
agentsResponse?.agents?.map((agent) => ({ label: agent.agent_name, value: agent.agent_id })) || null
}
premiumUser={premiumUser}
dateValue={dateValue}
/>
</TabPanel>
{/* User Agent Activity Panel */}
<TabPanel>
<UserAgentActivity accessToken={accessToken} userRole={userRole} dateValue={dateValue} />
Expand Down
Loading