mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
feat(ui): edit logging_only_scope in the custom code guardrail modal
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
f8c790a7b0
commit
625eb49c85
2 changed files with 169 additions and 3 deletions
|
|
@ -1,5 +1,5 @@
|
|||
import { describe, it, expect, vi, beforeEach } from "vitest";
|
||||
import { render, screen, waitFor } from "@testing-library/react";
|
||||
import { fireEvent, render, screen, waitFor } from "@testing-library/react";
|
||||
import userEvent from "@testing-library/user-event";
|
||||
import CustomCodeModal from "./CustomCodeModal";
|
||||
import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
|
||||
|
|
@ -130,6 +130,127 @@ describe("CustomCodeModal", () => {
|
|||
expect(screen.queryByText("logging_only")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should update the logging-only scope in edit mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderModal({
|
||||
editData: {
|
||||
guardrail_id: "g-1",
|
||||
guardrail_name: "existing-guardrail",
|
||||
litellm_params: {
|
||||
mode: ["logging_only"],
|
||||
logging_only_scope: "input",
|
||||
custom_code: "def apply_guardrail(): pass",
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
const scopeSelect = await screen.findByRole("combobox", { name: "Logging only scope" });
|
||||
expect(scopeSelect).toHaveTextContent("Input only (request)");
|
||||
|
||||
await user.click(scopeSelect);
|
||||
await user.click(await screen.findByRole("option", { name: "Output only (response)" }));
|
||||
await user.click(screen.getByRole("button", { name: /update guardrail/i }));
|
||||
|
||||
await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1));
|
||||
expect(mockUpdate.mock.calls[0][2]).toMatchObject({
|
||||
litellm_params: { logging_only_scope: "output" },
|
||||
});
|
||||
});
|
||||
|
||||
it("should clear the logging-only scope when Default is selected in edit mode", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderModal({
|
||||
editData: {
|
||||
guardrail_id: "g-1",
|
||||
guardrail_name: "existing-guardrail",
|
||||
litellm_params: {
|
||||
mode: ["logging_only"],
|
||||
logging_only_scope: "input",
|
||||
custom_code: "def apply_guardrail(): pass",
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
await user.click(await screen.findByRole("combobox", { name: "Logging only scope" }));
|
||||
await user.click(await screen.findByRole("option", { name: "Default (request and response)" }));
|
||||
await user.click(screen.getByRole("button", { name: /update guardrail/i }));
|
||||
|
||||
await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1));
|
||||
expect(mockUpdate.mock.calls[0][2]).toHaveProperty("litellm_params.logging_only_scope", null);
|
||||
});
|
||||
|
||||
it("should omit an unchanged logging-only scope from the edit payload", async () => {
|
||||
renderModal({
|
||||
editData: {
|
||||
guardrail_id: "g-1",
|
||||
guardrail_name: "existing-guardrail",
|
||||
litellm_params: {
|
||||
mode: ["logging_only"],
|
||||
logging_only_scope: "input",
|
||||
custom_code: "def apply_guardrail(): pass",
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
fireEvent.change(await screen.findByPlaceholderText("e.g., block-pii-custom"), {
|
||||
target: { value: "renamed-guardrail" },
|
||||
});
|
||||
await userEvent.setup().click(screen.getByRole("button", { name: /update guardrail/i }));
|
||||
|
||||
await waitFor(() => expect(mockUpdate).toHaveBeenCalledTimes(1));
|
||||
expect(mockUpdate.mock.calls[0][2]).not.toHaveProperty("litellm_params.logging_only_scope");
|
||||
});
|
||||
|
||||
it("should not show logging-only scope for other modes in edit mode", async () => {
|
||||
renderModal({
|
||||
editData: {
|
||||
guardrail_id: "g-1",
|
||||
guardrail_name: "existing-guardrail",
|
||||
litellm_params: { mode: "pre_call", custom_code: "def apply_guardrail(): pass" },
|
||||
},
|
||||
});
|
||||
|
||||
expect(await screen.findByText("Edit Custom Guardrail")).toBeInTheDocument();
|
||||
expect(screen.queryByRole("combobox", { name: "Logging only scope" })).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should create a logging-only guardrail with its selected scope", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderModal();
|
||||
|
||||
await user.click(screen.getAllByRole("combobox")[0]);
|
||||
await user.keyboard("logging_only");
|
||||
await user.click(await screen.findByRole("option", { name: "logging_only" }));
|
||||
await user.click(await screen.findByRole("combobox", { name: "Logging only scope" }));
|
||||
await user.click(await screen.findByRole("option", { name: "Input only (request)" }));
|
||||
fireEvent.change(screen.getByPlaceholderText("e.g., block-pii-custom"), {
|
||||
target: { value: "logging-only-guardrail" },
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /save guardrail/i }));
|
||||
|
||||
await waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1));
|
||||
expect(mockCreate.mock.calls[0][1]).toMatchObject({
|
||||
litellm_params: { logging_only_scope: "input" },
|
||||
});
|
||||
});
|
||||
|
||||
it("should omit the default logging-only scope when creating a guardrail", async () => {
|
||||
const user = userEvent.setup();
|
||||
renderModal();
|
||||
|
||||
await user.click(screen.getAllByRole("combobox")[0]);
|
||||
await user.keyboard("logging_only");
|
||||
await user.click(await screen.findByRole("option", { name: "logging_only" }));
|
||||
await screen.findByRole("combobox", { name: "Logging only scope" });
|
||||
fireEvent.change(screen.getByPlaceholderText("e.g., block-pii-custom"), {
|
||||
target: { value: "logging-only-guardrail" },
|
||||
});
|
||||
await user.click(screen.getByRole("button", { name: /save guardrail/i }));
|
||||
|
||||
await waitFor(() => expect(mockCreate).toHaveBeenCalledTimes(1));
|
||||
expect(mockCreate.mock.calls[0][1]).not.toHaveProperty("litellm_params.logging_only_scope");
|
||||
});
|
||||
|
||||
it("should expand the test section and run a test against the backend", async () => {
|
||||
const user = userEvent.setup();
|
||||
mockTest.mockResolvedValue({ success: true, result: { action: "allow" } } as never);
|
||||
|
|
|
|||
|
|
@ -2,7 +2,14 @@ import React, { useState, useRef, useEffect } from "react";
|
|||
import { CheckCircle2, ChevronRight, Code, ExternalLink, PlayCircle, Save, Users, XCircle } from "lucide-react";
|
||||
import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
|
||||
import { toast } from "@/lib/toast";
|
||||
import { type LoggingOnlyScope } from "../guardrail_info_helpers";
|
||||
import {
|
||||
choiceToLoggingOnlyScope,
|
||||
getLoggingOnlyScopeOptions,
|
||||
getLoggingOnlyScopeUpdate,
|
||||
loggingOnlyScopeToChoice,
|
||||
type LoggingOnlyScope,
|
||||
type LoggingOnlyScopeChoice,
|
||||
} from "../guardrail_info_helpers";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
|
||||
import {
|
||||
|
|
@ -165,6 +172,8 @@ const TEMPLATE_ITEMS = Object.entries(CODE_TEMPLATES).map(([key, template]) => (
|
|||
label: template.name,
|
||||
}));
|
||||
|
||||
const LOGGING_ONLY_SCOPE_ITEMS = getLoggingOnlyScopeOptions(true);
|
||||
|
||||
type ModeOption = (typeof MODE_OPTIONS)[number];
|
||||
|
||||
const MODE_OPTION_BY_VALUE: Record<string, ModeOption> = Object.fromEntries(
|
||||
|
|
@ -193,11 +202,13 @@ interface CustomCodeModalProps {
|
|||
editData?: EditGuardrailData | null;
|
||||
}
|
||||
|
||||
/* eslint-disable max-lines -- The existing modal remains the requested location for this scope control */
|
||||
const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onSuccess, accessToken, editData }) => {
|
||||
const anchor = useComboboxAnchor();
|
||||
const isEditMode = !!editData;
|
||||
const [guardrailName, setGuardrailName] = useState("");
|
||||
const [mode, setMode] = useState<string[]>(["pre_call"]);
|
||||
const [loggingOnlyScopeChoice, setLoggingOnlyScopeChoice] = useState<LoggingOnlyScopeChoice>("default");
|
||||
const [defaultOn, setDefaultOn] = useState(false);
|
||||
const [selectedTemplate, setSelectedTemplate] = useState<string>("empty");
|
||||
const [code, setCode] = useState(CODE_TEMPLATES.empty.code);
|
||||
|
|
@ -322,6 +333,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
|
|||
// Edit mode: populate with existing data
|
||||
setGuardrailName(editData.guardrail_name || "");
|
||||
setMode(normalizeMode(editData.litellm_params?.mode));
|
||||
setLoggingOnlyScopeChoice(loggingOnlyScopeToChoice(editData.litellm_params?.logging_only_scope));
|
||||
setDefaultOn(editData.litellm_params?.default_on || false);
|
||||
setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code);
|
||||
setSelectedTemplate(""); // No template selected in edit mode
|
||||
|
|
@ -329,6 +341,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
|
|||
// Create mode: reset to defaults
|
||||
setGuardrailName("");
|
||||
setMode(["pre_call"]);
|
||||
setLoggingOnlyScopeChoice("default");
|
||||
setDefaultOn(false);
|
||||
setSelectedTemplate("empty");
|
||||
setCode(CODE_TEMPLATES.empty.code);
|
||||
|
|
@ -401,11 +414,17 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
|
|||
if (defaultOn !== editData.litellm_params?.default_on) {
|
||||
updateData.litellm_params.default_on = defaultOn;
|
||||
}
|
||||
|
||||
if (mode.includes("logging_only")) {
|
||||
Object.assign(
|
||||
updateData.litellm_params,
|
||||
getLoggingOnlyScopeUpdate(editData.litellm_params, loggingOnlyScopeChoice),
|
||||
);
|
||||
}
|
||||
await updateGuardrailCall(accessToken, editData.guardrail_id, updateData);
|
||||
toast.success("Custom code guardrail updated successfully");
|
||||
} else {
|
||||
// Create new guardrail
|
||||
const loggingOnlyScope = choiceToLoggingOnlyScope(loggingOnlyScopeChoice);
|
||||
const guardrailData = {
|
||||
guardrail_name: guardrailName,
|
||||
litellm_params: {
|
||||
|
|
@ -413,6 +432,9 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
|
|||
mode: mode,
|
||||
default_on: defaultOn,
|
||||
custom_code: code,
|
||||
...(mode.includes("logging_only") && loggingOnlyScope !== null
|
||||
? { logging_only_scope: loggingOnlyScope }
|
||||
: {}),
|
||||
},
|
||||
guardrail_info: {},
|
||||
};
|
||||
|
|
@ -549,6 +571,29 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
|
|||
</ComboboxContent>
|
||||
</Combobox>
|
||||
</div>
|
||||
{mode.includes("logging_only") && (
|
||||
<div className="w-[200px]">
|
||||
<label className="mb-1 block text-xs font-medium text-muted-foreground">Logging only scope</label>
|
||||
<Select
|
||||
items={LOGGING_ONLY_SCOPE_ITEMS}
|
||||
value={loggingOnlyScopeChoice}
|
||||
onValueChange={(value: string | null) =>
|
||||
value && setLoggingOnlyScopeChoice(value as LoggingOnlyScopeChoice)
|
||||
}
|
||||
>
|
||||
<SelectTrigger className="w-full" aria-label="Logging only scope">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{LOGGING_ONLY_SCOPE_ITEMS.map((option) => (
|
||||
<SelectItem key={option.value} value={option.value}>
|
||||
{option.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
)}
|
||||
<div className="w-[180px]">
|
||||
<label className="mb-1 block text-xs font-medium text-muted-foreground">Template</label>
|
||||
<Select
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue