style(code-editor): allow editing custom code guardrails on ui + add examples of pre/post calls for custom code guardrails

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<CustomCodeModalProps> = ({
@@ -151,7 +165,9 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
   onClose,
   onSuccess,
   accessToken,
+  editData,
 }) => {
+  const isEditMode = !!editData;
   const [guardrailName, setGuardrailName] = useState("");
   const [mode, setMode] = useState<string>("pre_call");
   const [defaultOn, setDefaultOn] = useState(false);
@@ -160,7 +176,61 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
   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<any>(null);
   const [copiedPrimitive, setCopiedPrimitive] = useState<string | null>(null);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
@@ -171,18 +241,28 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
     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<CustomCodeModalProps> = ({
     }
   };
 
-  // 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<CustomCodeModalProps> = ({
 
     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<CustomCodeModalProps> = ({
       <div className="flex flex-col h-[80vh]">
         {/* Header */}
         <div className="pb-4 border-b border-gray-200">
-          <h2 className="text-xl font-semibold text-gray-900">Create Custom Guardrail</h2>
+          <h2 className="text-xl font-semibold text-gray-900">
+            {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"}
+          </h2>
           <p className="text-sm text-gray-500 mt-1">Define custom logic using Python-like syntax</p>
         </div>
 
@@ -371,12 +477,12 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
         {/* Main Content */}
         <div className="flex flex-1 overflow-hidden mt-4 gap-6">
           {/* Code Editor */}
-          <div className="flex-[2] flex flex-col min-w-0">
-            <div className="flex items-center justify-between mb-2">
+          <div className="flex-[2] flex flex-col min-w-0 overflow-y-auto">
+            <div className="flex items-center justify-between mb-2 flex-shrink-0">
               <span className="text-xs font-semibold text-gray-500 uppercase tracking-wide">Python Logic</span>
               <span className="text-xs text-gray-400">Restricted environment (no imports)</span>
             </div>
-            <div className="flex-1 relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e]" style={{ minHeight: "350px" }}>
+            <div className="relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e] flex-shrink-0" style={{ minHeight: "300px", maxHeight: "400px" }}>
               {/* Line numbers */}
               <div 
                 className="absolute left-0 top-0 bottom-0 w-12 bg-[#1e1e1e] border-r border-gray-700 text-right pr-3 pt-3 select-none overflow-hidden"
@@ -402,7 +508,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
             <Collapse
               activeKey={testExpanded ? ["test"] : []}
               onChange={(keys) => 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 }) => <CaretRightOutlined rotate={isActive ? 90 : 0} />}
             >
               <Panel
@@ -416,11 +522,40 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
               >
                 <div className="space-y-3">
                   <div>
-                    <label className="block text-xs font-medium text-gray-600 mb-1">Test Input (JSON)</label>
+                    <div className="flex items-center justify-between mb-2">
+                      <label className="block text-xs font-medium text-gray-600">Test Input (JSON)</label>
+                      <div className="flex items-center gap-2">
+                        <span className="text-xs text-gray-500">Load example:</span>
+                        <button
+                          type="button"
+                          onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2))}
+                          className="px-2 py-1 text-xs rounded border border-orange-200 bg-orange-50 text-orange-700 hover:bg-orange-100 transition-colors"
+                        >
+                          Pre-call
+                        </button>
+                        <button
+                          type="button"
+                          onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.post_call.data, null, 2))}
+                          className="px-2 py-1 text-xs rounded border border-green-200 bg-green-50 text-green-700 hover:bg-green-100 transition-colors"
+                        >
+                          Post-call
+                        </button>
+                      </div>
+                    </div>
+                    <div className="mb-2 p-2 bg-gray-50 rounded text-xs text-gray-600 border border-gray-200">
+                      <div className="grid grid-cols-2 gap-x-4 gap-y-1">
+                        <div><strong>texts</strong>: Message content (always)</div>
+                        <div><strong>images</strong>: Base64 images (vision)</div>
+                        <div><strong>tools</strong>: Tool definitions <span className="text-orange-600">(pre_call)</span></div>
+                        <div><strong>tool_calls</strong>: LLM tool calls <span className="text-green-600">(post_call)</span></div>
+                        <div><strong>structured_messages</strong>: Full messages <span className="text-orange-600">(pre_call)</span></div>
+                        <div><strong>model</strong>: Model name (always)</div>
+                      </div>
+                    </div>
                     <TextArea
                       value={testInput}
                       onChange={(e) => setTestInput(e.target.value)}
-                      rows={4}
+                      rows={8}
                       className="font-mono text-xs"
                       placeholder='{"texts": ["test message"], ...}'
                     />
@@ -535,7 +670,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({
               disabled={isSaving || !guardrailName.trim()}
               icon={SaveOutlined}
             >
-              Save Guardrail
+              {isEditMode ? "Update Guardrail" : "Save Guardrail"}
             </Button>
           </div>
         </div>

ui/litellm-dashboard/src/components/guardrails/guardrail_info.tsx

@@ -1,37 +1,39 @@
+import React, { useState, useEffect, useCallback } from "react";
 import {
-  getGuardrailInfo,
-  getGuardrailProviderSpecificParams,
-  getGuardrailUISettings,
-  updateGuardrailCall,
-} from "@/components/networking";
-import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
-import { CodeOutlined, EyeInvisibleOutlined, InfoCircleOutlined, StopOutlined } from "@ant-design/icons";
-import { ArrowLeftIcon } from "@heroicons/react/outline";
-import {
-  Badge,
   Card,
+  Title,
+  Text,
   Grid,
+  Badge,
+  Button as TremorButton,
   Tab,
   TabGroup,
   TabList,
   TabPanel,
   TabPanels,
-  Text,
-  Title,
+  TextInput,
 } from "@tremor/react";
-import { Button, Divider, Form, Input, Select, Tooltip } from "antd";
-import { CheckIcon, CopyIcon } from "lucide-react";
-import React, { useCallback, useEffect, useState } from "react";
-import NotificationsManager from "../molecules/notifications_manager";
-import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager";
-import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal";
+import { Button, Form, Input, Select, Divider, Tooltip } from "antd";
+import { InfoCircleOutlined, EyeInvisibleOutlined, StopOutlined, CodeOutlined } from "@ant-design/icons";
+import {
+  getGuardrailInfo,
+  updateGuardrailCall,
+  getGuardrailUISettings,
+  getGuardrailProviderSpecificParams,
+} from "@/components/networking";
 import { getGuardrailLogoAndName, guardrail_provider_map } from "./guardrail_info_helpers";
-import GuardrailOptionalParams from "./guardrail_optional_params";
-import GuardrailProviderFields from "./guardrail_provider_fields";
 import PiiConfiguration from "./pii_configuration";
+import GuardrailProviderFields from "./guardrail_provider_fields";
+import GuardrailOptionalParams from "./guardrail_optional_params";
+import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager";
 import ToolPermissionRulesEditor, {
   ToolPermissionConfig,
 } from "./tool_permission/ToolPermissionRulesEditor";
+import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal";
+import { ArrowLeftIcon } from "@heroicons/react/outline";
+import { copyToClipboard as utilCopyToClipboard } from "@/utils/dataUtils";
+import { CheckIcon, CopyIcon } from "lucide-react";
+import NotificationsManager from "../molecules/notifications_manager";
 
 export interface GuardrailInfoProps {
   guardrailId: string;
@@ -81,12 +83,6 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
       }>;
       pattern_categories: string[];
       supported_actions: string[];
-      content_categories?: Array<{
-        name: string;
-        display_name: string;
-        description: string;
-        default_action: string;
-      }>;
     };
   } | null>(null);
   const [copiedStates, setCopiedStates] = useState<Record<string, boolean>>({});
@@ -102,19 +98,14 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
   const [customCodeModalVisible, setCustomCodeModalVisible] = useState(false);
 
   // Content Filter data ref (managed by ContentFilterManager)
-  const contentFilterDataRef = React.useRef<{
-    patterns: any[];
-    blockedWords: any[];
-    categories: any[];
-  }>({
+  const contentFilterDataRef = React.useRef<{ patterns: any[]; blockedWords: any[] }>({
     patterns: [],
     blockedWords: [],
-    categories: [],
   });
 
   // Memoize onDataChange callback to prevent unnecessary re-renders
-  const handleContentFilterDataChange = useCallback((patterns: any[], blockedWords: any[], categories: any[]) => {
-    contentFilterDataRef.current = { patterns, blockedWords, categories: categories || [] };
+  const handleContentFilterDataChange = useCallback((patterns: any[], blockedWords: any[]) => {
+    contentFilterDataRef.current = { patterns, blockedWords };
   }, []);
 
   const fetchGuardrailInfo = async () => {
@@ -283,15 +274,16 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
 
       // Only add Content Filter patterns if there are changes
       if (guardrailData.litellm_params?.guardrail === "litellm_content_filter" && hasUnsavedContentFilterChanges) {
+        const originalPatterns = guardrailData.litellm_params?.patterns || [];
+        const originalBlockedWords = guardrailData.litellm_params?.blocked_words || [];
+
         const formattedData = formatContentFilterDataForAPI(
           contentFilterDataRef.current.patterns || [],
           contentFilterDataRef.current.blockedWords || [],
-          contentFilterDataRef.current.categories || [],
         );
 
         updateData.litellm_params.patterns = formattedData.patterns;
         updateData.litellm_params.blocked_words = formattedData.blocked_words;
-        updateData.litellm_params.categories = formattedData.categories;
       }
 
       if (guardrailData.litellm_params?.guardrail === "tool_permission") {
@@ -362,7 +354,7 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
 
         console.log("allowedParams: ", allowedParams);
         allowedParams.forEach((paramName) => {
-          if (paramName === "patterns" || paramName === "blocked_words" || paramName === "categories") {
+          if (paramName === "patterns" || paramName === "blocked_words") {
             return;
           }
           // Check for both direct parameter name and nested optional_params object
@@ -446,9 +438,9 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
   return (
     <div className="p-4">
       <div>
-        <Button type="text" icon={<ArrowLeftIcon className="w-4 h-4" />} onClick={onClose} className="mb-4">
+        <TremorButton icon={ArrowLeftIcon} variant="light" onClick={onClose} className="mb-4">
           Back to Guardrails
-        </Button>
+        </TremorButton>
         <Title>{guardrailData.guardrail_name || "Unnamed Guardrail"}</Title>
         <div className="flex items-center cursor-pointer">
           <Text className="text-gray-500 font-mono">{guardrailData.guardrail_id}</Text>
@@ -458,10 +450,11 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
             size="small"
             icon={copiedStates["guardrail-id"] ? <CheckIcon size={12} /> : <CopyIcon size={12} />}
             onClick={() => copyToClipboard(guardrailData.guardrail_id, "guardrail-id")}
-            className={`left-2 z-10 transition-all duration-200 ${copiedStates["guardrail-id"]
-              ? "text-green-600 bg-green-50 border-green-200"
-              : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
-              }`}
+            className={`left-2 z-10 transition-all duration-200 ${
+              copiedStates["guardrail-id"]
+                ? "text-green-600 bg-green-50 border-green-200"
+                : "text-gray-500 hover:text-gray-700 hover:bg-gray-100"
+            }`}
           />
         </div>
       </div>
@@ -540,8 +533,9 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                           <Text className="flex-1 font-medium text-gray-900">{key}</Text>
                           <Text className="flex-1">
                             <span
-                              className={`inline-flex items-center gap-1.5 ${value === "MASK" ? "text-blue-600" : "text-red-600"
-                                }`}
+                              className={`inline-flex items-center gap-1.5 ${
+                                value === "MASK" ? "text-blue-600" : "text-red-600"
+                              }`}
                             >
                               {value === "MASK" ? <EyeInvisibleOutlined /> : <StopOutlined />}
                               {String(value)}
@@ -569,13 +563,14 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                     <Text className="font-medium text-lg">Custom Code</Text>
                   </div>
                   {isAdmin && !isConfigGuardrail && (
-                    <Button
-                      size="small"
-                      icon={<CodeOutlined />}
+                    <TremorButton 
+                      size="xs"
+                      variant="secondary"
+                      icon={CodeOutlined}
                       onClick={() => setCustomCodeModalVisible(true)}
                     >
                       Edit Code
-                    </Button>
+                    </TremorButton>
                   )}
                 </div>
                 <div className="relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e]">
@@ -608,14 +603,14 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                   )}
                   {!isEditing && !isConfigGuardrail && (
                     guardrailData.litellm_params?.guardrail === "custom_code" ? (
-                      <Button
-                        icon={<CodeOutlined />}
+                      <TremorButton 
+                        icon={CodeOutlined}
                         onClick={() => setCustomCodeModalVisible(true)}
                       >
                         Edit Code
-                      </Button>
+                      </TremorButton>
                     ) : (
-                      <Button onClick={() => setIsEditing(true)}>Edit Settings</Button>
+                      <TremorButton onClick={() => setIsEditing(true)}>Edit Settings</TremorButton>
                     )
                   )}
                 </div>
@@ -642,7 +637,7 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                       name="guardrail_name"
                       rules={[{ required: true, message: "Please input a guardrail name" }]}
                     >
-                      <Input placeholder="Enter guardrail name" />
+                      <TextInput />
                     </Form.Item>
 
                     <Form.Item label="Default On" name="default_on">
@@ -680,7 +675,7 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                       onUnsavedChanges={setHasUnsavedContentFilterChanges}
                     />
 
-                    {(guardrailData.litellm_params?.guardrail === "tool_permission" || guardrailProviderSpecificParams) && <Divider orientation="left">Provider Settings</Divider>}
+                    <Divider orientation="left">Provider Settings</Divider>
 
                     {guardrailData.litellm_params?.guardrail === "tool_permission" ? (
                       <ToolPermissionRulesEditor
@@ -740,7 +735,7 @@ const GuardrailInfoView: React.FC<GuardrailInfoProps> = ({ guardrailId, onClose,
                       >
                         Cancel
                       </Button>
-                      <Button type="primary" htmlType="submit">Save Changes</Button>
+                      <TremorButton>Save Changes</TremorButton>
                     </div>
                   </Form>
                 ) : (