diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.test.tsx new file mode 100644 index 00000000000..0521f0327fa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.test.tsx @@ -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(); + + 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(); + 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(); + 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(); + + 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(); + 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(); + 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(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx index 071ad1d6bd5..26f87298015 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/PromptCompressionTab.tsx @@ -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 = ({ accessToken }) => { const [form] = Form.useForm(); - const [guardrails, setGuardrails] = useState([]); + const [existing, setExisting] = useState(null); const [isLoading, setIsLoading] = useState(true); const [isSaving, setIsSaving] = useState(false); + const [confirmingDisable, setConfirmingDisable] = useState(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 = ({ 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 ( -
+
- + Headroom prompt compression + +
+ {existing ? "Enabled" : "Disabled"} + setConfirmingDisable(false)} + > + + +
+
- +

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 = ({ accessToken Headroom setup docs

- {isLoading &&

Loading...

} - {!isLoading && guardrails.length === 0 && ( -

- No prompt compression guardrails configured yet. Add one below to start saving on input tokens -

- )} - {!isLoading && guardrails.length > 0 && ( -
    - {guardrails.map((guardrail) => ( -
  • -
    -

    {guardrail.guardrail_name}

    -

    {guardrail.litellm_params?.api_base ?? ""}

    -
    - - {guardrail.litellm_params?.default_on ? "Always on" : "Opt-in"} - -
  • - ))} -
- )} -
-
- - - Add Headroom compression guardrail - - -
+ @@ -140,13 +216,19 @@ const PromptCompressionTab: React.FC = ({ 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" }]} > - + +

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 = ({ accessToken

-
- -
+ + {existing ? ( +
+ +
+ ) : ( +

+ Enter a name and Headroom API base, then switch this on to start saving on input tokens +

+ )}