import React, { useState, useRef, useEffect } from "react"; import { Modal, Select, Switch, Collapse, Input, Spin } from "antd"; import { Button, TextInput } from "@tremor/react"; import { CodeOutlined, PlayCircleOutlined, CheckCircleOutlined, CloseCircleOutlined, CaretRightOutlined, SaveOutlined, } from "@ant-design/icons"; import { createGuardrailCall, testCustomCodeGuardrail } from "../../networking"; import NotificationsManager from "../../molecules/notifications_manager"; const { Panel } = Collapse; const { TextArea } = Input; // Code templates const CODE_TEMPLATES = { empty: { name: "Empty Template", 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" return allow()`, }, blockSSN: { name: "Block SSN", code: `def apply_guardrail(inputs, request_data, input_type): for text in inputs["texts"]: if regex_match(text, r"\\d{3}-\\d{2}-\\d{4}"): return block("SSN detected") return allow()`, }, redactEmail: { name: "Redact Emails", code: `def apply_guardrail(inputs, request_data, input_type): pattern = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\\.[a-zA-Z]{2,}" modified = [] for text in inputs["texts"]: modified.append(regex_replace(text, pattern, "[EMAIL REDACTED]")) return modify(texts=modified)`, }, blockSQL: { name: "Block SQL Injection", code: `def apply_guardrail(inputs, request_data, input_type): if input_type != "request": return allow() for text in inputs["texts"]: if contains_code_language(text, ["sql"]): return block("SQL code not allowed") return allow()`, }, validateJSON: { name: "Validate JSON", code: `def apply_guardrail(inputs, request_data, input_type): if input_type != "response": return allow() schema = {"type": "object", "required": ["name", "value"]} for text in inputs["texts"]: obj = json_parse(text) if obj is None: return block("Invalid JSON response") if not json_schema_valid(obj, schema): return block("Response missing required fields") return allow()`, }, }; // Available primitives organized by category const PRIMITIVES = { "Return Values": [ { name: "allow()", desc: "Let request/response through" }, { name: "block(reason)", desc: "Reject with message" }, { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, ], "Regex Functions": [ { name: "regex_match(text, pattern)", desc: "Returns True if pattern found" }, { name: "regex_replace(text, pattern, replacement)", desc: "Replace all matches" }, { name: "regex_find_all(text, pattern)", desc: "Return list of matches" }, ], "JSON Functions": [ { name: "json_parse(text)", desc: "Parse JSON string, returns None on error" }, { name: "json_stringify(obj)", desc: "Convert to JSON string" }, { name: "json_schema_valid(obj, schema)", desc: "Validate against JSON schema" }, ], "URL Functions": [ { name: "extract_urls(text)", desc: "Extract all URLs from text" }, { name: "is_valid_url(url)", desc: "Check if URL is valid" }, { name: "all_urls_valid(text)", desc: "Check all URLs in text are valid" }, ], "Code Detection": [ { name: "detect_code(text)", desc: "Returns True if code detected" }, { name: "detect_code_languages(text)", desc: "Returns list of detected languages" }, { name: 'contains_code_language(text, ["sql"])', desc: "Check for specific languages" }, ], "Text Utilities": [ { name: "contains(text, substring)", desc: "Check if substring exists" }, { name: "contains_any(text, [substr1, substr2])", desc: "Check if any substring exists" }, { name: "word_count(text)", desc: "Count words" }, { name: "char_count(text)", desc: "Count characters" }, { name: "lower(text) / upper(text) / trim(text)", desc: "String transforms" }, ], }; const MODE_OPTIONS = [ { value: "pre_call", label: "pre_call (Request)" }, { value: "post_call", label: "post_call (Response)" }, { value: "during_call", label: "during_call (Parallel)" }, { value: "logging_only", label: "logging_only" }, ]; interface CustomCodeModalProps { visible: boolean; onClose: () => void; onSuccess: () => void; accessToken: string | null; } const CustomCodeModal: React.FC = ({ visible, onClose, onSuccess, accessToken, }) => { const [guardrailName, setGuardrailName] = useState(""); 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); 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); // Handle template change const handleTemplateChange = (templateKey: string) => { setSelectedTemplate(templateKey); setCode(CODE_TEMPLATES[templateKey as keyof typeof CODE_TEMPLATES].code); }; // Reset form when modal opens useEffect(() => { if (visible) { setGuardrailName(""); setMode("pre_call"); setDefaultOn(false); setSelectedTemplate("empty"); setCode(CODE_TEMPLATES.empty.code); setTestResult(null); setTestExpanded(false); } }, [visible]); // Copy primitive to clipboard const copyPrimitive = async (primitive: string) => { try { await navigator.clipboard.writeText(primitive); setCopiedPrimitive(primitive); setTimeout(() => setCopiedPrimitive(null), 2000); } catch (err) { console.error("Failed to copy:", err); } }; // Handle tab key in textarea const handleKeyDown = (e: React.KeyboardEvent) => { if (e.key === "Tab") { e.preventDefault(); const textarea = e.currentTarget; const start = textarea.selectionStart; const end = textarea.selectionEnd; const newValue = code.substring(0, start) + " " + code.substring(end); setCode(newValue); setTimeout(() => { textarea.selectionStart = textarea.selectionEnd = start + 4; }, 0); } }; // Save guardrail const handleSave = async () => { if (!guardrailName.trim()) { NotificationsManager.fromBackend("Please enter a guardrail name"); return; } if (!code.trim()) { NotificationsManager.fromBackend("Please enter custom code"); return; } if (!accessToken) { NotificationsManager.fromBackend("No access token available"); return; } setIsSaving(true); try { 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); NotificationsManager.fromBackend( "Failed to create guardrail: " + (error instanceof Error ? error.message : String(error)) ); } finally { setIsSaving(false); } }; // Test guardrail using backend endpoint const handleTest = async () => { if (!accessToken) { setTestResult({ error: "No access token available" }); return; } setIsTesting(true); setTestResult(null); try { // Parse test input JSON let parsedInput; try { parsedInput = JSON.parse(testInput); } catch (e) { setTestResult({ error: "Invalid test input JSON" }); setIsTesting(false); return; } // Ensure texts array exists if (!parsedInput.texts) { parsedInput.texts = []; } const response = await testCustomCodeGuardrail(accessToken, { custom_code: code, test_input: parsedInput, input_type: mode as "request" | "response", request_data: { model: "test-model", metadata: {}, }, }); if (response.success && response.result) { setTestResult(response.result); } else if (response.error) { setTestResult({ error: response.error, error_type: response.error_type, }); } else { setTestResult({ error: "Unknown error occurred" }); } } catch (error) { console.error("Failed to test custom code:", error); setTestResult({ error: error instanceof Error ? error.message : "Failed to test custom code", }); } finally { setIsTesting(false); } }; const lineCount = code.split("\n").length; return (
{/* Header */}

Create Custom Guardrail

Define custom logic using Python-like syntax

{/* Top Controls */}
{Object.entries(CODE_TEMPLATES).map(([key, template]) => ( {template.name} ))}
Default On
{/* Main Content */}
{/* Code Editor */}
Python Logic Restricted environment (no imports)
{/* Line numbers */}
{Array.from({ length: Math.max(lineCount, 20) }, (_, i) => (
{i + 1}
))}
{/* Code textarea */}