diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.ts new file mode 100644 index 000000000000..f2b7e76777dc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.ts @@ -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({ + queryKey: agentsKeys.list({}), + queryFn: async () => await getAgentsList(accessToken!), + enabled: Boolean(accessToken) && all_admin_roles.includes(userRole || ""), + }); +}; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx index e326183d8807..b390c19df163 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx @@ -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; diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts index ded2731c945e..81b0307c7124 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/types.ts @@ -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[]; diff --git a/ui/litellm-dashboard/src/components/entity_usage.test.tsx b/ui/litellm-dashboard/src/components/entity_usage.test.tsx index 2b6234c039eb..d0d2337e185d 100644 --- a/ui/litellm-dashboard/src/components/entity_usage.test.tsx +++ b/ui/litellm-dashboard/src/components/entity_usage.test.tsx @@ -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 @@ -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: [ @@ -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 () => { @@ -201,6 +205,21 @@ describe("EntityUsage", () => { }); }); + it("should render with agent entity type and call agent API", async () => { + render(); + + 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(); diff --git a/ui/litellm-dashboard/src/components/entity_usage.tsx b/ui/litellm-dashboard/src/components/entity_usage.tsx index ca30ded9494c..cced9e8ff987 100644 --- a/ui/litellm-dashboard/src/components/entity_usage.tsx +++ b/ui/litellm-dashboard/src/components/entity_usage.tsx @@ -28,6 +28,7 @@ import { tagDailyActivityCall, teamDailyActivityCall, customerDailyActivityCall, + agentDailyActivityCall, } from "./networking"; import TopKeyView from "./top_key_view"; import { formatNumberWithCommas } from "@/utils/dataUtils"; @@ -150,6 +151,15 @@ const EntityUsage: React.FC = ({ 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"); } diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index cd1289f0a889..5f94fc16f013 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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 diff --git a/ui/litellm-dashboard/src/components/new_usage.test.tsx b/ui/litellm-dashboard/src/components/new_usage.test.tsx index aec07765e7e0..340201452fbe 100644 --- a/ui/litellm-dashboard/src/components/new_usage.test.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.test.tsx @@ -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(() => { @@ -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: [ @@ -193,6 +199,13 @@ describe("NewUsage", () => { }, ]; + const mockAgents = [ + { + agent_id: "agent-123", + agent_name: "Test Agent", + }, + ]; + const defaultProps = { accessToken: "test-token", userRole: "Admin", @@ -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 () => { @@ -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(() => { @@ -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(() => { @@ -298,18 +320,20 @@ describe("NewUsage", () => { }); it("should show organization usage banner and tab for admins", async () => { - const { getByText, getAllByText } = render(); + render(); 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); }); }); @@ -321,17 +345,43 @@ describe("NewUsage", () => { error: null, } as any); - const { getByText, getAllByText } = render(); + render(); + + 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(); 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); }); }); diff --git a/ui/litellm-dashboard/src/components/new_usage.tsx b/ui/litellm-dashboard/src/components/new_usage.tsx index 9b6c7e48c554..fd89ddb869b6 100644 --- a/ui/litellm-dashboard/src/components/new_usage.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.tsx @@ -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; @@ -88,6 +89,7 @@ const NewUsagePage: React.FC = ({ const [allTags, setAllTags] = useState([]); 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); @@ -435,6 +437,7 @@ const NewUsagePage: React.FC = ({ Team Usage {all_admin_roles.includes(userRole || "") ? Customer Usage : <>} {all_admin_roles.includes(userRole || "") ? Tag Usage : <>} + {all_admin_roles.includes(userRole || "") ? Agent Usage : <>} {all_admin_roles.includes(userRole || "") ? User Agent Activity : <>} @@ -842,6 +845,19 @@ const NewUsagePage: React.FC = ({ dateValue={dateValue} /> + + ({ label: agent.agent_name, value: agent.agent_id })) || null + } + premiumUser={premiumUser} + dateValue={dateValue} + /> + {/* User Agent Activity Panel */}