mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
feat(ui): make Headroom compression a single always-on config card
Replace the two-card cost-optimization compression layout (a read-only list plus a separate add form) with one permanent card that reflects and edits the single Headroom guardrail. The card exposes name and Headroom API base as required fields, an apply-to-all-requests toggle mapped to default_on, and a header on/off switch. Turning it on creates the guardrail, editing and saving PATCHes it, and turning it off confirms then deletes it. There is no persisted disabled state for a guardrail, so off removes it.
This commit is contained in:
Tin Chi Lo
parent
64aad5877a
commit
d751c33c52
2 changed files with 306 additions and 71 deletions
|
|
@ -0,0 +1,146 @@
|
|||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import { beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
import {
|
||||
createGuardrailCall,
|
||||
deleteGuardrailCall,
|
||||
getGuardrailsList,
|
||||
updateGuardrailCall,
|
||||
} from "@/components/networking";
|
||||
|
||||
vi.mock("@/components/networking", () => ({
|
||||
getGuardrailsList: vi.fn(),
|
||||
createGuardrailCall: vi.fn(),
|
||||
updateGuardrailCall: vi.fn(),
|
||||
deleteGuardrailCall: vi.fn(),
|
||||
}));
|
||||
|
||||
vi.mock("@/components/molecules/notifications_manager", () => ({
|
||||
__esModule: true,
|
||||
default: { success: vi.fn(), fromBackend: vi.fn(), error: vi.fn() },
|
||||
}));
|
||||
|
||||
import PromptCompressionTab from "./PromptCompressionTab";
|
||||
|
||||
const headroomGuardrail = {
|
||||
guardrail_id: "gr-1",
|
||||
guardrail_name: "headroom-compression",
|
||||
litellm_params: { guardrail: "headroom", api_base: "https://headroom.example.com", default_on: true },
|
||||
};
|
||||
|
||||
const mockedList = vi.mocked(getGuardrailsList);
|
||||
const mockedCreate = vi.mocked(createGuardrailCall);
|
||||
const mockedUpdate = vi.mocked(updateGuardrailCall);
|
||||
const mockedDelete = vi.mocked(deleteGuardrailCall);
|
||||
|
||||
const masterSwitch = () => screen.getByLabelText("Toggle Headroom compression");
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
describe("PromptCompressionTab - no guardrail configured", () => {
|
||||
beforeEach(() => {
|
||||
mockedList.mockResolvedValue({ guardrails: [] });
|
||||
});
|
||||
|
||||
it("renders a single permanent card with no separate add-guardrail section", async () => {
|
||||
render(<PromptCompressionTab accessToken="sk-test" />);
|
||||
|
||||
await waitFor(() => expect(screen.getByText("Disabled")).toBeInTheDocument());
|
||||
expect(screen.getByText("Headroom prompt compression")).toBeInTheDocument();
|
||||
expect(screen.queryByText("Add Headroom compression guardrail")).not.toBeInTheDocument();
|
||||
expect(screen.queryByRole("button", { name: "Add guardrail" })).not.toBeInTheDocument();
|
||||
expect(masterSwitch()).not.toBeChecked();
|
||||
});
|
||||
|
||||
it("does not create a guardrail when enabling with empty required fields", async () => {
|
||||
render(<PromptCompressionTab accessToken="sk-test" />);
|
||||
await waitFor(() => expect(masterSwitch()).toBeEnabled());
|
||||
|
||||
fireEvent.click(masterSwitch());
|
||||
|
||||
await waitFor(() => expect(screen.getByText("Name is required")).toBeInTheDocument());
|
||||
expect(mockedCreate).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it("creates a headroom guardrail with the entered name, api base and apply-to-all", async () => {
|
||||
mockedCreate.mockResolvedValue({});
|
||||
render(<PromptCompressionTab accessToken="sk-test" />);
|
||||
await waitFor(() => expect(masterSwitch()).toBeEnabled());
|
||||
|
||||
fireEvent.change(screen.getByPlaceholderText("headroom-compression"), { target: { value: "my-headroom" } });
|
||||
fireEvent.change(screen.getByPlaceholderText("https://your-headroom-endpoint"), {
|
||||
target: { value: "https://compress.internal" },
|
||||
});
|
||||
fireEvent.click(masterSwitch());
|
||||
|
||||
await waitFor(() => expect(mockedCreate).toHaveBeenCalledTimes(1));
|
||||
expect(mockedCreate).toHaveBeenCalledWith("sk-test", {
|
||||
guardrail_name: "my-headroom",
|
||||
litellm_params: {
|
||||
guardrail: "headroom",
|
||||
mode: "pre_call",
|
||||
api_base: "https://compress.internal",
|
||||
default_on: true,
|
||||
},
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe("PromptCompressionTab - guardrail already configured", () => {
|
||||
beforeEach(() => {
|
||||
mockedList.mockResolvedValue({ guardrails: [headroomGuardrail] });
|
||||
});
|
||||
|
||||
it("shows the enabled card populated from the existing guardrail", async () => {
|
||||
render(<PromptCompressionTab accessToken="sk-test" />);
|
||||
|
||||
await waitFor(() => expect(screen.getByText("Enabled")).toBeInTheDocument());
|
||||
expect(masterSwitch()).toBeChecked();
|
||||
expect(screen.getByPlaceholderText("headroom-compression")).toHaveValue("headroom-compression");
|
||||
expect(screen.getByPlaceholderText("https://your-headroom-endpoint")).toHaveValue("https://headroom.example.com");
|
||||
});
|
||||
|
||||
it("keeps Save changes disabled until a field is edited, then patches the guardrail", async () => {
|
||||
mockedUpdate.mockResolvedValue({});
|
||||
render(<PromptCompressionTab accessToken="sk-test" />);
|
||||
await waitFor(() => expect(screen.getByText("Enabled")).toBeInTheDocument());
|
||||
|
||||
const save = screen.getByRole("button", { name: "Save changes" });
|
||||
expect(save).toBeDisabled();
|
||||
|
||||
fireEvent.change(screen.getByPlaceholderText("https://your-headroom-endpoint"), {
|
||||
target: { value: "https://new-endpoint.internal" },
|
||||
});
|
||||
await waitFor(() => expect(save).toBeEnabled());
|
||||
|
||||
fireEvent.click(save);
|
||||
await waitFor(() => expect(mockedUpdate).toHaveBeenCalledTimes(1));
|
||||
expect(mockedUpdate).toHaveBeenCalledWith("sk-test", "gr-1", {
|
||||
guardrail_name: "headroom-compression",
|
||||
litellm_params: {
|
||||
guardrail: "headroom",
|
||||
mode: "pre_call",
|
||||
api_base: "https://new-endpoint.internal",
|
||||
default_on: true,
|
||||
},
|
||||
});
|
||||
});
|
||||
|
||||
it("deletes the guardrail after confirming the turn-off prompt", async () => {
|
||||
mockedDelete.mockResolvedValue({});
|
||||
render(<PromptCompressionTab accessToken="sk-test" />);
|
||||
await waitFor(() => expect(screen.getByText("Enabled")).toBeInTheDocument());
|
||||
|
||||
fireEvent.click(masterSwitch());
|
||||
|
||||
const confirm = await screen.findByRole("button", { name: "Turn off" });
|
||||
fireEvent.click(confirm);
|
||||
|
||||
await waitFor(() => expect(mockedDelete).toHaveBeenCalledWith("sk-test", "gr-1"));
|
||||
await waitFor(() => expect(screen.getByText("Disabled")).toBeInTheDocument());
|
||||
expect(mockedUpdate).not.toHaveBeenCalled();
|
||||
expect(mockedCreate).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
@ -1,10 +1,15 @@
|
|||
"use client";
|
||||
|
||||
import React, { useCallback, useEffect, useState } from "react";
|
||||
import { Button, Form, Input, Switch } from "antd";
|
||||
import React, { useCallback, useEffect, useMemo, useState } from "react";
|
||||
import { Button, Form, Input, Popconfirm, Switch } from "antd";
|
||||
|
||||
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import { createGuardrailCall, getGuardrailsList } from "@/components/networking";
|
||||
import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
|
||||
import {
|
||||
createGuardrailCall,
|
||||
deleteGuardrailCall,
|
||||
getGuardrailsList,
|
||||
updateGuardrailCall,
|
||||
} from "@/components/networking";
|
||||
import NotificationsManager from "@/components/molecules/notifications_manager";
|
||||
import {
|
||||
buildCompressionGuardrailPayload,
|
||||
|
|
@ -20,36 +25,70 @@ interface PromptCompressionTabProps {
|
|||
interface CompressionFormValues {
|
||||
name: string;
|
||||
apiBase: string;
|
||||
defaultOn: boolean;
|
||||
applyToAll: boolean;
|
||||
}
|
||||
|
||||
const EMPTY_FORM: CompressionFormValues = { name: "", apiBase: "", applyToAll: true };
|
||||
|
||||
const toFormValues = (guardrail: GuardrailListItem): CompressionFormValues => ({
|
||||
name: guardrail.guardrail_name ?? "",
|
||||
apiBase: guardrail.litellm_params?.api_base ?? "",
|
||||
applyToAll: guardrail.litellm_params?.default_on ?? false,
|
||||
});
|
||||
|
||||
const PromptCompressionTab: React.FC<PromptCompressionTabProps> = ({ accessToken }) => {
|
||||
const [form] = Form.useForm<CompressionFormValues>();
|
||||
const [guardrails, setGuardrails] = useState<GuardrailListItem[]>([]);
|
||||
const [existing, setExisting] = useState<GuardrailListItem | null>(null);
|
||||
const [isLoading, setIsLoading] = useState<boolean>(true);
|
||||
const [isSaving, setIsSaving] = useState<boolean>(false);
|
||||
const [confirmingDisable, setConfirmingDisable] = useState<boolean>(false);
|
||||
|
||||
const loadGuardrails = useCallback(() => {
|
||||
const watched = Form.useWatch([], form);
|
||||
|
||||
const loadGuardrail = useCallback(() => {
|
||||
if (!accessToken) {
|
||||
return;
|
||||
}
|
||||
getGuardrailsList(accessToken)
|
||||
.then((response) => setGuardrails(compressionGuardrailsOf(response as GuardrailListResponse)))
|
||||
.then((response) => {
|
||||
const [first = null] = compressionGuardrailsOf(response as GuardrailListResponse);
|
||||
setExisting(first);
|
||||
if (first) {
|
||||
form.setFieldsValue(toFormValues(first));
|
||||
}
|
||||
})
|
||||
.catch((error) => {
|
||||
console.error("Failed to load compression guardrails:", error);
|
||||
NotificationsManager.fromBackend("Failed to load compression guardrails");
|
||||
console.error("Failed to load compression guardrail:", error);
|
||||
NotificationsManager.fromBackend("Failed to load Headroom compression guardrail");
|
||||
})
|
||||
.finally(() => setIsLoading(false));
|
||||
}, [accessToken]);
|
||||
}, [accessToken, form]);
|
||||
|
||||
useEffect(() => {
|
||||
loadGuardrails();
|
||||
}, [loadGuardrails]);
|
||||
loadGuardrail();
|
||||
}, [loadGuardrail]);
|
||||
|
||||
const handleAdd = async (values: CompressionFormValues) => {
|
||||
const persisted = existing ? toFormValues(existing) : null;
|
||||
const isDirty = useMemo(() => {
|
||||
if (!persisted || !watched) {
|
||||
return false;
|
||||
}
|
||||
return (
|
||||
(watched.name ?? "") !== persisted.name ||
|
||||
(watched.apiBase ?? "") !== persisted.apiBase ||
|
||||
(watched.applyToAll ?? false) !== persisted.applyToAll
|
||||
);
|
||||
}, [persisted, watched]);
|
||||
|
||||
const enableGuardrail = useCallback(async () => {
|
||||
if (!accessToken) {
|
||||
return;
|
||||
}
|
||||
const values = await form.validateFields().catch(() => null);
|
||||
if (!values) {
|
||||
NotificationsManager.fromBackend("Enter a name and Headroom API base to turn on compression");
|
||||
return;
|
||||
}
|
||||
setIsSaving(true);
|
||||
try {
|
||||
await createGuardrailCall(
|
||||
|
|
@ -57,27 +96,104 @@ const PromptCompressionTab: React.FC<PromptCompressionTabProps> = ({ accessToken
|
|||
buildCompressionGuardrailPayload({
|
||||
name: values.name,
|
||||
apiBase: values.apiBase,
|
||||
defaultOn: values.defaultOn ?? true,
|
||||
defaultOn: values.applyToAll ?? true,
|
||||
}),
|
||||
);
|
||||
NotificationsManager.success("Compression guardrail created");
|
||||
form.resetFields();
|
||||
await loadGuardrails();
|
||||
NotificationsManager.success("Headroom compression turned on");
|
||||
loadGuardrail();
|
||||
} catch (error) {
|
||||
console.error("Failed to create compression guardrail:", error);
|
||||
NotificationsManager.fromBackend("Failed to create compression guardrail");
|
||||
console.error("Failed to turn on Headroom compression:", error);
|
||||
NotificationsManager.fromBackend("Failed to turn on Headroom compression");
|
||||
} finally {
|
||||
setIsSaving(false);
|
||||
}
|
||||
}, [accessToken, form, loadGuardrail]);
|
||||
|
||||
const saveChanges = useCallback(async () => {
|
||||
if (!accessToken || !existing) {
|
||||
return;
|
||||
}
|
||||
const values = await form.validateFields().catch(() => null);
|
||||
if (!values) {
|
||||
return;
|
||||
}
|
||||
setIsSaving(true);
|
||||
try {
|
||||
await updateGuardrailCall(
|
||||
accessToken,
|
||||
existing.guardrail_id,
|
||||
buildCompressionGuardrailPayload({
|
||||
name: values.name,
|
||||
apiBase: values.apiBase,
|
||||
defaultOn: values.applyToAll ?? false,
|
||||
}),
|
||||
);
|
||||
NotificationsManager.success("Headroom compression updated");
|
||||
loadGuardrail();
|
||||
} catch (error) {
|
||||
console.error("Failed to update Headroom compression:", error);
|
||||
NotificationsManager.fromBackend("Failed to update Headroom compression");
|
||||
} finally {
|
||||
setIsSaving(false);
|
||||
}
|
||||
}, [accessToken, existing, form, loadGuardrail]);
|
||||
|
||||
const disableGuardrail = useCallback(async () => {
|
||||
setConfirmingDisable(false);
|
||||
if (!accessToken || !existing) {
|
||||
return;
|
||||
}
|
||||
setIsSaving(true);
|
||||
try {
|
||||
await deleteGuardrailCall(accessToken, existing.guardrail_id);
|
||||
NotificationsManager.success("Headroom compression turned off");
|
||||
setExisting(null);
|
||||
} catch (error) {
|
||||
console.error("Failed to turn off Headroom compression:", error);
|
||||
NotificationsManager.fromBackend("Failed to turn off Headroom compression");
|
||||
} finally {
|
||||
setIsSaving(false);
|
||||
}
|
||||
}, [accessToken, existing]);
|
||||
|
||||
const handleEnabledChange = (next: boolean) => {
|
||||
if (next) {
|
||||
enableGuardrail();
|
||||
} else {
|
||||
setConfirmingDisable(true);
|
||||
}
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="w-full space-y-6">
|
||||
<div className="w-full">
|
||||
<Card>
|
||||
<CardHeader>
|
||||
<CardHeader className="border-b pb-4">
|
||||
<CardTitle>Headroom prompt compression</CardTitle>
|
||||
<CardAction>
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="text-sm text-muted-foreground">{existing ? "Enabled" : "Disabled"}</span>
|
||||
<Popconfirm
|
||||
open={confirmingDisable}
|
||||
title="Turn off Headroom compression?"
|
||||
description="Requests will stop being compressed and this configuration will be removed"
|
||||
okText="Turn off"
|
||||
okButtonProps={{ danger: true, loading: isSaving }}
|
||||
cancelText="Cancel"
|
||||
onConfirm={disableGuardrail}
|
||||
onCancel={() => setConfirmingDisable(false)}
|
||||
>
|
||||
<Switch
|
||||
aria-label="Toggle Headroom compression"
|
||||
checked={!!existing}
|
||||
loading={isSaving}
|
||||
disabled={isLoading}
|
||||
onChange={handleEnabledChange}
|
||||
/>
|
||||
</Popconfirm>
|
||||
</div>
|
||||
</CardAction>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<CardContent className="pt-6">
|
||||
<p className="mb-4 text-sm text-muted-foreground">
|
||||
Headroom is a native LiteLLM guardrail that compresses your prompts before they reach the model, so you pay
|
||||
for fewer input tokens. The tokens it removes are priced and shown on the Usage tab as compression savings.{" "}
|
||||
|
|
@ -90,48 +206,8 @@ const PromptCompressionTab: React.FC<PromptCompressionTabProps> = ({ accessToken
|
|||
Headroom setup docs
|
||||
</a>
|
||||
</p>
|
||||
{isLoading && <p className="text-sm text-muted-foreground">Loading...</p>}
|
||||
{!isLoading && guardrails.length === 0 && (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
No prompt compression guardrails configured yet. Add one below to start saving on input tokens
|
||||
</p>
|
||||
)}
|
||||
{!isLoading && guardrails.length > 0 && (
|
||||
<ul className="divide-y divide-gray-200">
|
||||
{guardrails.map((guardrail) => (
|
||||
<li key={guardrail.guardrail_id} className="flex items-center justify-between py-3">
|
||||
<div>
|
||||
<p className="text-sm font-medium text-foreground">{guardrail.guardrail_name}</p>
|
||||
<p className="text-xs text-muted-foreground">{guardrail.litellm_params?.api_base ?? ""}</p>
|
||||
</div>
|
||||
<span
|
||||
className={`rounded-full px-2 py-0.5 text-xs font-medium ${
|
||||
guardrail.litellm_params?.default_on
|
||||
? "bg-emerald-100 text-emerald-800"
|
||||
: "bg-gray-100 text-gray-600"
|
||||
}`}
|
||||
>
|
||||
{guardrail.litellm_params?.default_on ? "Always on" : "Opt-in"}
|
||||
</span>
|
||||
</li>
|
||||
))}
|
||||
</ul>
|
||||
)}
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
||||
<Card>
|
||||
<CardHeader>
|
||||
<CardTitle>Add Headroom compression guardrail</CardTitle>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<Form
|
||||
form={form}
|
||||
layout="vertical"
|
||||
requiredMark={false}
|
||||
onFinish={handleAdd}
|
||||
initialValues={{ defaultOn: true }}
|
||||
>
|
||||
<Form form={form} layout="vertical" requiredMark={false} initialValues={EMPTY_FORM} disabled={isSaving}>
|
||||
<Form.Item name="name" label="Name" rules={[{ required: true, message: "Name is required" }]}>
|
||||
<Input placeholder="headroom-compression" />
|
||||
</Form.Item>
|
||||
|
|
@ -140,13 +216,19 @@ const PromptCompressionTab: React.FC<PromptCompressionTabProps> = ({ accessToken
|
|||
label="Headroom API base"
|
||||
tooltip="Base URL of your Headroom compression service (LiteLLM calls its /v1/compress endpoint)"
|
||||
extra="The URL where your Headroom compression service is hosted"
|
||||
rules={[{ required: true, message: "API base is required" }]}
|
||||
rules={[{ required: true, message: "Headroom API base is required" }]}
|
||||
>
|
||||
<Input placeholder="https://your-headroom-endpoint" />
|
||||
</Form.Item>
|
||||
<Form.Item name="defaultOn" label="Apply to all requests" valuePropName="checked">
|
||||
<Form.Item
|
||||
name="applyToAll"
|
||||
label="Apply to all requests"
|
||||
valuePropName="checked"
|
||||
extra="On: every request is compressed. Off: compression is available for keys or teams to opt in"
|
||||
>
|
||||
<Switch />
|
||||
</Form.Item>
|
||||
|
||||
<div className="mb-4 rounded-lg border border-yellow-200 bg-yellow-50 p-3">
|
||||
<p className="text-sm text-yellow-800">
|
||||
Applying compression to all requests is available to all users. Enabling it selectively per key or team
|
||||
|
|
@ -161,11 +243,18 @@ const PromptCompressionTab: React.FC<PromptCompressionTabProps> = ({ accessToken
|
|||
</a>
|
||||
</p>
|
||||
</div>
|
||||
<div className="flex justify-end">
|
||||
<Button type="primary" htmlType="submit" loading={isSaving}>
|
||||
Add guardrail
|
||||
</Button>
|
||||
</div>
|
||||
|
||||
{existing ? (
|
||||
<div className="flex justify-end">
|
||||
<Button type="primary" onClick={saveChanges} loading={isSaving} disabled={!isDirty}>
|
||||
Save changes
|
||||
</Button>
|
||||
</div>
|
||||
) : (
|
||||
<p className="text-sm text-muted-foreground">
|
||||
Enter a name and Headroom API base, then switch this on to start saving on input tokens
|
||||
</p>
|
||||
)}
|
||||
</Form>
|
||||
</CardContent>
|
||||
</Card>
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue