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 866d24df50b..6531a153c5a 100644 --- a/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/custom_code/CustomCodeModal.tsx @@ -1,5 +1,5 @@ import React, { useState, useRef, useEffect } from "react"; -import { Modal, Select, Switch, Collapse, Input, Divider } from "antd"; +import { Modal, Select, Switch, Collapse, Input, Spin } from "antd"; import { Button, TextInput } from "@tremor/react"; import { CodeOutlined, @@ -8,10 +8,8 @@ import { CloseCircleOutlined, CaretRightOutlined, SaveOutlined, - UsergroupAddOutlined, - ExportOutlined, } from "@ant-design/icons"; -import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "../../networking"; +import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; const { Panel } = Collapse; @@ -21,7 +19,7 @@ const { TextArea } = Input; const CODE_TEMPLATES = { empty: { name: "Empty Template", - code: `async def apply_guardrail(inputs, request_data, input_type): + code: `def apply_guardrail(inputs, request_data, input_type): # inputs: {texts, images, tools, tool_calls, structured_messages, model} # request_data: {model, user_id, team_id, end_user_id, metadata} # input_type: "request" or "response" @@ -70,30 +68,8 @@ const CODE_TEMPLATES = { return block("Response missing required fields") return allow()`, }, - externalAPI: { - name: "External API Check (async)", - code: `async def apply_guardrail(inputs, request_data, input_type): - # Call an external moderation API (async for non-blocking) - for text in inputs["texts"]: - response = await http_post( - "https://api.example.com/moderate", - body={"text": text, "user_id": request_data["user_id"]}, - headers={"Authorization": "Bearer YOUR_API_KEY"}, - timeout=10 - ) - - if not response["success"]: - # API call failed, allow by default or block - return allow() - - if response["body"].get("flagged"): - return block(response["body"].get("reason", "Content flagged")) - - return allow()`, - }, }; - // Available primitives organized by category const PRIMITIVES = { "Return Values": [ @@ -101,11 +77,6 @@ const PRIMITIVES = { { name: "block(reason)", desc: "Reject with message" }, { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, ], - "HTTP Requests (async)": [ - { name: "await http_request(url, method, headers, body)", desc: "Make async HTTP request" }, - { name: "await http_get(url, headers)", desc: "Async GET request" }, - { name: "await http_post(url, body, headers)", desc: "Async POST request" }, - ], "Regex Functions": [ { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" }, { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" }, @@ -140,30 +111,13 @@ const MODE_OPTIONS = [ { value: "post_call", label: "post_call (Response)" }, { value: "during_call", label: "during_call (Parallel)" }, { value: "logging_only", label: "logging_only" }, - { value: "pre_mcp_call", label: "pre_mcp_call (Before MCP Tool Call)" }, - { value: "post_mcp_call", label: "post_mcp_call (After MCP Tool Call)" }, - { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" }, ]; -// Data for editing an existing guardrail -export interface EditGuardrailData { - guardrail_id: string; - guardrail_name: string; - litellm_params: { - mode?: string | 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 = ({ @@ -171,72 +125,16 @@ const CustomCodeModal: React.FC = ({ onClose, onSuccess, accessToken, - editData, }) => { - const isEditMode = !!editData; const [guardrailName, setGuardrailName] = useState(""); - const [mode, setMode] = useState(["pre_call"]); + const [mode, setMode] = useState("pre_call"); const [defaultOn, setDefaultOn] = useState(false); const [selectedTemplate, setSelectedTemplate] = useState("empty"); const [code, setCode] = useState(CODE_TEMPLATES.empty.code); const [isSaving, setIsSaving] = useState(false); const [isTesting, setIsTesting] = useState(false); const [testExpanded, setTestExpanded] = useState(false); - - // 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 [testInput, setTestInput] = useState('{"texts": ["Hello, my SSN is 123-45-6789"], "images": [], "tools": [], "tool_calls": [], "structured_messages": [], "model": "gpt-4"}'); const [testResult, setTestResult] = useState(null); const [copiedPrimitive, setCopiedPrimitive] = useState(null); const textareaRef = useRef(null); @@ -244,40 +142,21 @@ const CustomCodeModal: React.FC = ({ // Handle template change const handleTemplateChange = (templateKey: string) => { setSelectedTemplate(templateKey); - - // Check if it's a standard template setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code); }; - // Normalize mode from API (string or string[]) to string[] - const normalizeMode = (m: string | string[] | undefined): string[] => { - if (m === undefined || m === null) return ["pre_call"]; - if (Array.isArray(m)) return m.length ? m : ["pre_call"]; - return [m]; - }; - - // Reset form when modal opens or editData changes + // Reset form when modal opens useEffect(() => { if (visible) { - if (editData) { - // Edit mode: populate with existing data - setGuardrailName(editData.guardrail_name || ""); - setMode(normalizeMode(editData.litellm_params?.mode)); - 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); - } + setGuardrailName(""); + setMode("pre_call"); + setDefaultOn(false); + setSelectedTemplate("empty"); + setCode(CODE_TEMPLATES.empty.code); setTestResult(null); setTestExpanded(false); } - }, [visible, editData]); + }, [visible]); // Copy primitive to clipboard const copyPrimitive = async (primitive: string) => { @@ -305,7 +184,7 @@ const CustomCodeModal: React.FC = ({ } }; - // Save guardrail (create or update) + // Save guardrail const handleSave = async () => { if (!guardrailName.trim()) { NotificationsManager.fromBackend("Please enter a guardrail name"); @@ -322,53 +201,25 @@ const CustomCodeModal: React.FC = ({ setIsSaving(true); try { - if (isEditMode && editData) { - // Update existing guardrail - const updateData: any = { - litellm_params: { - custom_code: code, - }, - }; + const guardrailData = { + guardrail_name: guardrailName, + litellm_params: { + guardrail: "custom_code", + mode: mode, + default_on: defaultOn, + custom_code: code, + }, + guardrail_info: {}, + }; - // Only include changed fields - if (guardrailName !== editData.guardrail_name) { - updateData.guardrail_name = guardrailName; - } - const existingMode = normalizeMode(editData.litellm_params?.mode); - const modeChanged = - mode.length !== existingMode.length || - mode.some((m, i) => m !== existingMode[i]); - if (modeChanged) { - 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"); - } + await createGuardrailCall(accessToken, guardrailData); + NotificationsManager.success("Custom code guardrail created successfully"); onSuccess(); onClose(); } catch (error) { - console.error("Failed to save guardrail:", error); + console.error("Failed to create guardrail:", error); NotificationsManager.fromBackend( - `Failed to ${isEditMode ? "update" : "create"} guardrail: ` + (error instanceof Error ? error.message : String(error)) + "Failed to create guardrail: " + (error instanceof Error ? error.message : String(error)) ); } finally { setIsSaving(false); @@ -401,20 +252,10 @@ const CustomCodeModal: React.FC = ({ parsedInput.texts = []; } - // Use first request-like or response-like mode for test input_type - const requestModes = ["pre_call", "pre_mcp_call"]; - const responseModes = ["post_call", "post_mcp_call"]; - const testInputType: "request" | "response" = - mode.some((m) => requestModes.includes(m)) - ? "request" - : mode.some((m) => responseModes.includes(m)) - ? "response" - : "request"; - const response = await testCustomCodeGuardrail(accessToken, { custom_code: code, test_input: parsedInput, - input_type: testInputType, + input_type: mode as "request" | "response", request_data: { model: "test-model", metadata: {}, @@ -456,9 +297,7 @@ const CustomCodeModal: React.FC = ({
{/* Header */}
-

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

+

Create Custom Guardrail

Define custom logic using Python-like syntax

@@ -472,16 +311,14 @@ const CustomCodeModal: React.FC = ({ placeholder="e.g., block-pii-custom" />
-
- +
+
@@ -541,12 +345,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 flex-shrink-0" + className="mt-3 bg-white border border-gray-200 rounded-lg" 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)
-
-
+