Merge pull request #17883 from BerriAI/litellm_ui_new_badge

[Feature] New Badge for Agent Usage
This commit is contained in:
yuneng-jiang 2025-12-12 10:02:31 -08:00 committed by GitHub
commit 25a38b0f27
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 102 additions and 26 deletions

View file

@ -99,7 +99,18 @@ vi.mock("antd", async () => {
}
(Alert as any).displayName = "AntdAlert";
return { Select, Alert };
function Badge(props: any) {
const { count, color, children, ...rest } = props;
return React.createElement(
"div",
{ ...rest, "data-testid": "antd-badge", "data-color": color },
count && React.createElement("span", { "data-testid": "antd-badge-count" }, count),
children,
);
}
(Badge as any).displayName = "AntdBadge";
return { Select, Alert, Badge };
});
vi.mock("@ant-design/icons", async () => {

View file

@ -27,7 +27,7 @@ import {
Text,
Title,
} from "@tremor/react";
import { Alert } from "antd";
import { Alert, Badge } from "antd";
import React, { useCallback, useEffect, useMemo, useState } from "react";
import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents";
@ -419,11 +419,13 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
<div className="flex items-end justify-between gap-6 mb-6">
<div className="flex-1">
<div className="flex items-end justify-between gap-6 mb-4 w-full">
<UsageViewSelect
value={usageView}
onChange={(value) => setUsageView(value)}
isAdmin={all_admin_roles.includes(userRole || "")}
/>
<Badge color="blue" count="New">
<UsageViewSelect
value={usageView}
onChange={(value) => setUsageView(value)}
isAdmin={all_admin_roles.includes(userRole || "")}
/>
</Badge>
<AdvancedDatePicker value={dateValue} onValueChange={handleDateChange} />
</div>
{/* Your Usage Panel */}

View file

@ -6,21 +6,63 @@ vi.mock("antd", async () => {
const React = await import("react");
function Select(props: any) {
const { value, onChange, options, ...rest } = props;
const { value, onChange, options, optionRender, labelRender, ...rest } = props;
const selectedOption = options?.find((opt: any) => opt.value === value);
const renderedLabel = labelRender ? labelRender({ value, label: selectedOption?.label }) : selectedOption?.label;
const optionElements = options?.map((opt: any) => {
const rendered = optionRender ? optionRender({ value: opt.value, label: opt.label }) : opt.label;
return React.createElement("option", { key: opt.value, value: opt.value }, opt.label);
});
const optionRenderOutputs = options
?.map((opt: any) => {
if (optionRender) {
const rendered = optionRender({ value: opt.value, label: opt.label });
return React.createElement(
"div",
{
key: `option-render-${opt.value}`,
"data-testid": `option-render-${opt.value}`,
style: { display: "none" },
},
rendered,
);
}
return null;
})
.filter(Boolean);
return React.createElement(
"select",
{
...rest,
value,
onChange: (e: any) => onChange?.(e.target.value),
role: "combobox",
},
options?.map((opt: any) => React.createElement("option", { key: opt.value, value: opt.value }, opt.label)),
React.Fragment,
null,
React.createElement(
"select",
{
...rest,
value,
onChange: (e: any) => onChange?.(e.target.value),
role: "combobox",
},
optionElements,
),
...(optionRenderOutputs || []),
);
}
(Select as any).displayName = "AntdSelect";
return { Select };
function Badge(props: any) {
const { count, color, children, ...rest } = props;
return React.createElement(
"span",
{ ...rest, "data-testid": "antd-badge", "data-color": color },
count && React.createElement("span", { "data-testid": "antd-badge-count" }, count),
children,
);
}
(Badge as any).displayName = "AntdBadge";
return { Select, Badge };
});
vi.mock("@ant-design/icons", async () => {
@ -67,4 +109,13 @@ describe("UsageViewSelect", () => {
expect(mockOnChange).toHaveBeenCalledWith("team");
});
it("should render badge when option has badgeText", () => {
render(<UsageViewSelect value="agent" onChange={mockOnChange} isAdmin={true} />);
const badge = screen.getByTestId("antd-badge");
expect(badge).toBeInTheDocument();
expect(badge).toHaveAttribute("data-color", "blue");
expect(screen.getByTestId("antd-badge-count")).toHaveTextContent("New");
});
});

View file

@ -8,7 +8,7 @@ import {
TagsOutlined,
TeamOutlined,
} from "@ant-design/icons";
import { Select } from "antd";
import { Badge, Select } from "antd";
import React from "react";
export type UsageOption = "global" | "organization" | "team" | "customer" | "tag" | "agent" | "user-agent-activity";
export interface UsageViewSelectProps {
@ -29,6 +29,7 @@ interface OptionConfig {
showForNonAdmin?: string;
descriptionForAdmin?: string;
descriptionForNonAdmin?: string;
badgeText?: string;
}
const OPTIONS: OptionConfig[] = [
{
@ -37,8 +38,8 @@ const OPTIONS: OptionConfig[] = [
showForAdmin: "Global Usage",
showForNonAdmin: "Your Usage",
description: "View usage across all resources",
descriptionForAdmin: "View usage across all resources and users",
descriptionForNonAdmin: "View your personal usage statistics",
descriptionForAdmin: "View usage across all resources",
descriptionForNonAdmin: "View your usage",
icon: <GlobalOutlined style={{ fontSize: "16px" }} />,
},
{
@ -48,7 +49,7 @@ const OPTIONS: OptionConfig[] = [
showForNonAdmin: "Your Organization Usage",
description: "View organization-level usage",
descriptionForAdmin: "View usage across all organizations",
descriptionForNonAdmin: "View your organization's usage statistics",
descriptionForNonAdmin: "View your organization's usage",
icon: <BankOutlined style={{ fontSize: "16px" }} />,
},
{
@ -77,6 +78,7 @@ const OPTIONS: OptionConfig[] = [
description: "View usage by AI agents",
icon: <RobotOutlined style={{ fontSize: "16px" }} />,
adminOnly: true,
badgeText: "New",
},
{
value: "user-agent-activity",
@ -114,6 +116,7 @@ export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
label,
description: desc,
icon: option.icon,
badgeText: option.badgeText,
};
});
};
@ -134,7 +137,7 @@ export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
<Select
value={value}
onChange={onChange}
className="w-48 sm:w-64 md:w-72"
className="w-54 sm:w-64 md:w-72"
size="large"
options={filteredOptions.map((opt) => ({
value: opt.value,
@ -144,12 +147,17 @@ export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
const opt = filteredOptions.find((o) => o.value === option.value);
if (!opt) return option.label;
return (
<div className="flex items-start gap-2 py-1">
<div className="flex items-center gap-2 py-1">
<div className="flex-shrink-0 mt-0.5">{opt.icon}</div>
<div className="flex-1 min-w-0">
<div className="text-sm font-medium text-gray-900">{opt.label}</div>
<div className="text-xs text-gray-600 mt-0.5">{opt.description}</div>
</div>
{opt.badgeText && (
<div className="items-center">
<Badge color="blue" count={opt.badgeText} />
</div>
)}
</div>
);
}}

View file

@ -21,7 +21,7 @@ import {
ToolOutlined,
UserOutlined,
} from "@ant-design/icons";
import { ConfigProvider, Layout, Menu } from "antd";
import { Badge, ConfigProvider, Layout, Menu } from "antd";
import { all_admin_roles, internalUserRoles, isAdminRole, rolesWithWriteAccess } from "../utils/roles";
import UsageIndicator from "./usage_indicator";
const { Sider } = Layout;
@ -39,7 +39,7 @@ interface SidebarProps {
interface MenuItem {
key: string;
page: string;
label: string;
label: string | React.ReactNode;
roles?: string[];
children?: MenuItem[]; // Add children property for submenus
icon?: React.ReactNode;
@ -71,7 +71,11 @@ const Sidebar: React.FC<SidebarProps> = ({ accessToken, setPage, userRole, defau
{
key: "new_usage",
page: "new_usage",
label: "Usage",
label: (
<span className="flex items-center gap-4">
Usage <Badge color="blue" count="New" />
</span>
),
icon: <BarChartOutlined style={{ fontSize: "18px" }} />,
roles: [...all_admin_roles, ...internalUserRoles],
},