From 15881e1d70208d946aea4260bd6382a275afb5f4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 5 Feb 2026 15:30:32 -0800 Subject: [PATCH] style(code-editor): allow editing custom code guardrails on ui + add examples of pre/post calls for custom code guardrails --- .../custom_code/CustomCodeModal.tsx | 199 +++++++++++++++--- .../components/guardrails/guardrail_info.tsx | 107 +++++----- 2 files changed, 218 insertions(+), 88 deletions(-) diff --git a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx index a4d33f889b3..5535688a9fd 100644 --- a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx @@ -9,7 +9,7 @@ import { CaretRightOutlined, SaveOutlined, } from "@ant-design/icons"; -import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking"; +import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; const { Panel } = Collapse; @@ -139,11 +139,25 @@ const MODE_OPTIONS = [ { value: "logging_only", label: "logging_only" }, ]; +// Data for editing an existing guardrail +export interface EditGuardrailData { + guardrail_id: string; + guardrail_name: string; + litellm_params: { + mode?: string; + default_on?: boolean; + custom_code?: string; + [key: string]: any; + }; +} + interface CustomCodeModalProps { visible: boolean; onClose: () => void; onSuccess: () => void; accessToken: string | null; + /** If provided, the modal will be in edit mode */ + editData?: EditGuardrailData | null; } const CustomCodeModal: React.FC = ({ @@ -151,7 +165,9 @@ const CustomCodeModal: React.FC = ({ onClose, onSuccess, accessToken, + editData, }) => { + const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); const [mode, setMode] = useState("pre_call"); const [defaultOn, setDefaultOn] = useState(false); @@ -160,7 +176,61 @@ const CustomCodeModal: React.FC = ({ const [isSaving, setIsSaving] = useState(false); const [isTesting, setIsTesting] = useState(false); const [testExpanded, setTestExpanded] = useState(false); - const [testInput, setTestInput] = useState('{"texts": ["Hello, my SSN is 123-45-6789"], "images": [], "tools": [], "tool_calls": [], "structured_messages": [], "model": "gpt-4"}'); + + // Test input examples for pre_call and post_call + const TEST_INPUT_EXAMPLES = { + pre_call: { + name: "Pre-call (Request)", + data: { + texts: ["Hello, my SSN is 123-45-6789"], + images: [], + tools: [ + { + type: "function", + function: { + name: "get_weather", + description: "Get the current weather in a location", + parameters: { + type: "object", + properties: { + location: { type: "string", description: "City name" } + }, + required: ["location"] + } + } + } + ], + tool_calls: [], + structured_messages: [ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "Hello, my SSN is 123-45-6789" } + ], + model: "gpt-4" + } + }, + post_call: { + name: "Post-call (Response)", + data: { + texts: ["The weather in San Francisco is 72°F and sunny."], + images: [], + tools: [], + tool_calls: [ + { + id: "call_abc123", + type: "function", + function: { + name: "get_weather", + arguments: "{\"location\": \"San Francisco\"}" + } + } + ], + structured_messages: [], + model: "gpt-4" + } + } + }; + + const [testInput, setTestInput] = useState(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2)); const [testResult, setTestResult] = useState(null); const [copiedPrimitive, setCopiedPrimitive] = useState(null); const textareaRef = useRef(null); @@ -171,18 +241,28 @@ const CustomCodeModal: React.FC = ({ setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code); }; - // Reset form when modal opens + // Reset form when modal opens or editData changes useEffect(() => { if (visible) { - setGuardrailName(""); - setMode("pre_call"); - setDefaultOn(false); - setSelectedTemplate("empty"); - setCode(CODE_TEMPLATES.empty.code); + if (editData) { + // Edit mode: populate with existing data + setGuardrailName(editData.guardrail_name || ""); + setMode(editData.litellm_params?.mode || "pre_call"); + setDefaultOn(editData.litellm_params?.default_on || false); + setCode(editData.litellm_params?.custom_code || CODE_TEMPLATES.empty.code); + setSelectedTemplate(""); // No template selected in edit mode + } else { + // Create mode: reset to defaults + setGuardrailName(""); + setMode("pre_call"); + setDefaultOn(false); + setSelectedTemplate("empty"); + setCode(CODE_TEMPLATES.empty.code); + } setTestResult(null); setTestExpanded(false); } - }, [visible]); + }, [visible, editData]); // Copy primitive to clipboard const copyPrimitive = async (primitive: string) => { @@ -210,7 +290,7 @@ const CustomCodeModal: React.FC = ({ } }; - // Save guardrail + // Save guardrail (create or update) const handleSave = async () => { if (!guardrailName.trim()) { NotificationsManager.fromBackend("Please enter a guardrail name"); @@ -227,25 +307,49 @@ const CustomCodeModal: React.FC = ({ setIsSaving(true); try { - const guardrailData = { - guardrail_name: guardrailName, - litellm_params: { - guardrail: "custom_code", - mode: mode, - default_on: defaultOn, - custom_code: code, - }, - guardrail_info: {}, - }; + if (isEditMode && editData) { + // Update existing guardrail + const updateData: any = { + litellm_params: { + custom_code: code, + }, + }; - await createGuardrailCall(accessToken, guardrailData); - NotificationsManager.success("Custom code guardrail created successfully"); + // Only include changed fields + if (guardrailName !== editData.guardrail_name) { + updateData.guardrail_name = guardrailName; + } + if (mode !== editData.litellm_params?.mode) { + updateData.litellm_params.mode = mode; + } + if (defaultOn !== editData.litellm_params?.default_on) { + updateData.litellm_params.default_on = defaultOn; + } + + await updateGuardrailCall(accessToken, editData.guardrail_id, updateData); + NotificationsManager.success("Custom code guardrail updated successfully"); + } else { + // Create new guardrail + const guardrailData = { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "custom_code", + mode: mode, + default_on: defaultOn, + custom_code: code, + }, + guardrail_info: {}, + }; + + await createGuardrailCall(accessToken, guardrailData); + NotificationsManager.success("Custom code guardrail created successfully"); + } onSuccess(); onClose(); } catch (error) { - console.error("Failed to create guardrail:", error); + console.error("Failed to save guardrail:", error); NotificationsManager.fromBackend( - "Failed to create guardrail: " + (error instanceof Error ? error.message : String(error)) + `Failed to ${isEditMode ? "update" : "create"} guardrail: ` + (error instanceof Error ? error.message : String(error)) ); } finally { setIsSaving(false); @@ -323,7 +427,9 @@ const CustomCodeModal: React.FC = ({
{/* Header */}
-

Create Custom Guardrail

+

+ {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"} +

Define custom logic using Python-like syntax

@@ -371,12 +477,12 @@ const CustomCodeModal: React.FC = ({ {/* Main Content */}
{/* Code Editor */} -
-
+
+
Python Logic Restricted environment (no imports)
-
+
{/* Line numbers */}
= ({ setTestExpanded(keys.includes("test"))} - className="mt-3 bg-white border border-gray-200 rounded-lg" + className="mt-3 bg-white border border-gray-200 rounded-lg flex-shrink-0" expandIcon={({ isActive }) => } > = ({ >
- +
+ +
+ Load example: + + +
+
+
+
+
texts: Message content (always)
+
images: Base64 images (vision)
+
tools: Tool definitions (pre_call)
+
tool_calls: LLM tool calls (post_call)
+
structured_messages: Full messages (pre_call)
+
model: Model name (always)
+
+