diff --git a/ui/litellm-dashboard/src/components/policies/add_policy_form.tsx b/ui/litellm-dashboard/src/components/policies/add_policy_form.tsx index 0f14a07e352..4ccc151bd13 100644 --- a/ui/litellm-dashboard/src/components/policies/add_policy_form.tsx +++ b/ui/litellm-dashboard/src/components/policies/add_policy_form.tsx @@ -1,9 +1,8 @@ import React, { useState, useEffect } from "react"; import { Form, Select, Modal, Divider, Typography, Tag, Alert, Radio } from "antd"; import { Button, TextInput, Textarea } from "@tremor/react"; -import { Policy, PolicyCreateRequest, PolicyUpdateRequest, GuardrailPipeline } from "./types"; +import { Policy, PolicyCreateRequest, PolicyUpdateRequest } from "./types"; import { Guardrail } from "../guardrails/types"; -import PipelineFlowBuilder, { createDefaultStep } from "./pipeline_flow_builder"; import { getResolvedGuardrails, modelAvailableCall } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; @@ -15,6 +14,7 @@ interface AddPolicyFormProps { visible: boolean; onClose: () => void; onSuccess: () => void; + onOpenFlowBuilder: () => void; accessToken: string | null; editingPolicy?: Policy | null; existingPolicies: Policy[]; @@ -23,10 +23,117 @@ interface AddPolicyFormProps { updatePolicy: (accessToken: string, policyId: string, policyData: any) => Promise; } +// ───────────────────────────────────────────────────────────────────────────── +// Mode Picker (Step 1) - shown first when creating a new policy +// ───────────────────────────────────────────────────────────────────────────── + +interface ModePicker { + selected: "simple" | "flow_builder"; + onSelect: (mode: "simple" | "flow_builder") => void; +} + +const ModePicker: React.FC = ({ selected, onSelect }) => ( +
+ {/* Simple Mode Card */} +
onSelect("simple")} + style={{ + flex: 1, + padding: "24px 20px", + border: `2px solid ${selected === "simple" ? "#4f46e5" : "#e5e7eb"}`, + borderRadius: 12, + cursor: "pointer", + backgroundColor: selected === "simple" ? "#eef2ff" : "#fff", + transition: "all 0.15s ease", + }} + > +
+ + + + +
+ + Simple Mode + + + Pick guardrails from a list. All run in parallel. + +
+ + {/* Flow Builder Card */} +
onSelect("flow_builder")} + style={{ + flex: 1, + padding: "24px 20px", + border: `2px solid ${selected === "flow_builder" ? "#4f46e5" : "#e5e7eb"}`, + borderRadius: 12, + cursor: "pointer", + backgroundColor: selected === "flow_builder" ? "#eef2ff" : "#fff", + transition: "all 0.15s ease", + position: "relative", + }} + > + + NEW + +
+ + + +
+ + Flow Builder + + + Define steps, conditions, and error responses. + +
+
+); + +// ───────────────────────────────────────────────────────────────────────────── +// Main Component +// ───────────────────────────────────────────────────────────────────────────── + const AddPolicyForm: React.FC = ({ visible, onClose, onSuccess, + onOpenFlowBuilder, accessToken, editingPolicy, existingPolicies, @@ -40,21 +147,16 @@ const AddPolicyForm: React.FC = ({ const [isLoadingResolved, setIsLoadingResolved] = useState(false); const [modelConditionType, setModelConditionType] = useState<"model" | "regex">("model"); const [availableModels, setAvailableModels] = useState([]); - const [policyMode, setPolicyMode] = useState<"simple" | "flow_builder">("simple"); - const [pipeline, setPipeline] = useState({ - mode: "pre_call", - steps: [createDefaultStep()], - }); + const [step, setStep] = useState<"pick_mode" | "simple_form">("pick_mode"); + const [selectedMode, setSelectedMode] = useState<"simple" | "flow_builder">("simple"); const { userId, userRole } = useAuthorized(); // Only consider it "editing" if editingPolicy has a policy_id (real existing policy) - // If editingPolicy is set but has no policy_id, it's just pre-filled data for a new policy (e.g., from a template) const isEditing = !!editingPolicy?.policy_id; useEffect(() => { if (visible && editingPolicy) { const modelCondition = editingPolicy.condition?.model; - // Detect if it's a regex pattern (contains *, ., [, ], etc.) const isRegex = modelCondition && /[.*+?^${}()|[\]\\]/.test(modelCondition); setModelConditionType(isRegex ? "regex" : "model"); @@ -66,23 +168,25 @@ const AddPolicyForm: React.FC = ({ guardrails_remove: editingPolicy.guardrails_remove || [], model_condition: modelCondition, }); - // Load resolved guardrails for editing + if (editingPolicy.policy_id && accessToken) { loadResolvedGuardrails(editingPolicy.policy_id); } - // Detect pipeline mode + + // If editing a pipeline policy, go directly to flow builder if (editingPolicy.pipeline) { - setPolicyMode("flow_builder"); - setPipeline(editingPolicy.pipeline); - } else { - setPolicyMode("simple"); + onClose(); + onOpenFlowBuilder(); + return; } + // If editing a simple policy, skip mode picker + setStep("simple_form"); } else if (visible) { form.resetFields(); setResolvedGuardrails([]); setModelConditionType("model"); - setPolicyMode("simple"); - setPipeline({ mode: "pre_call", steps: [createDefaultStep()] }); + setSelectedMode("simple"); + setStep("pick_mode"); } // eslint-disable-next-line react-hooks/exhaustive-deps }, [visible, editingPolicy, form]); @@ -96,7 +200,6 @@ const AddPolicyForm: React.FC = ({ const loadAvailableModels = async () => { if (!accessToken) return; - try { const response = await modelAvailableCall(accessToken, userId, userRole); if (response?.data) { @@ -110,7 +213,6 @@ const AddPolicyForm: React.FC = ({ const loadResolvedGuardrails = async (policyId: string) => { if (!accessToken) return; - setIsLoadingResolved(true); try { const data = await getResolvedGuardrails(accessToken, policyId); @@ -130,20 +232,15 @@ const AddPolicyForm: React.FC = ({ let resolved = new Set(); - // If inheriting, find parent policy and get its guardrails if (inheritFrom) { const parentPolicy = existingPolicies.find(p => p.policy_name === inheritFrom); if (parentPolicy) { - // Recursively resolve parent's guardrails const parentResolved = resolveParentGuardrails(parentPolicy); parentResolved.forEach(g => resolved.add(g)); } } - // Add guardrails guardrailsAdd.forEach((g: string) => resolved.add(g)); - - // Remove guardrails guardrailsRemove.forEach((g: string) => resolved.delete(g)); return Array.from(resolved).sort(); @@ -152,32 +249,23 @@ const AddPolicyForm: React.FC = ({ const resolveParentGuardrails = (policy: Policy): string[] => { let resolved = new Set(); - // If parent inherits, resolve recursively if (policy.inherit) { const grandparent = existingPolicies.find(p => p.policy_name === policy.inherit); if (grandparent) { - const grandparentResolved = resolveParentGuardrails(grandparent); - grandparentResolved.forEach(g => resolved.add(g)); + resolveParentGuardrails(grandparent).forEach(g => resolved.add(g)); } } - - // Add parent's guardrails if (policy.guardrails_add) { policy.guardrails_add.forEach(g => resolved.add(g)); } - - // Remove parent's removed guardrails if (policy.guardrails_remove) { policy.guardrails_remove.forEach(g => resolved.delete(g)); } - return Array.from(resolved); }; - // Recompute resolved guardrails when form values change const handleFormChange = () => { - const resolved = computeResolvedGuardrails(); - setResolvedGuardrails(resolved); + setResolvedGuardrails(computeResolvedGuardrails()); }; const resetForm = () => { @@ -186,9 +274,20 @@ const AddPolicyForm: React.FC = ({ const handleClose = () => { resetForm(); + setStep("pick_mode"); + setSelectedMode("simple"); onClose(); }; + const handleModeConfirm = () => { + if (selectedMode === "flow_builder") { + onClose(); + onOpenFlowBuilder(); + } else { + setStep("simple_form"); + } + }; + const handleSubmit = async () => { try { setIsSubmitting(true); @@ -199,36 +298,16 @@ const AddPolicyForm: React.FC = ({ throw new Error("No access token available"); } - let data: PolicyCreateRequest | PolicyUpdateRequest; - - if (policyMode === "flow_builder") { - // Flow builder mode: auto-derive guardrails_add from pipeline steps - const guardrailsFromPipeline = pipeline.steps - .map((s) => s.guardrail) - .filter(Boolean); - data = { - policy_name: values.policy_name, - description: values.description || undefined, - inherit: values.inherit || undefined, - guardrails_add: guardrailsFromPipeline, - guardrails_remove: [], - condition: values.model_condition - ? { model: values.model_condition } - : undefined, - pipeline: pipeline, - }; - } else { - data = { - policy_name: values.policy_name, - description: values.description || undefined, - inherit: values.inherit || undefined, - guardrails_add: values.guardrails_add || [], - guardrails_remove: values.guardrails_remove || [], - condition: values.model_condition - ? { model: values.model_condition } - : undefined, - }; - } + const data: PolicyCreateRequest | PolicyUpdateRequest = { + policy_name: values.policy_name, + description: values.description || undefined, + inherit: values.inherit || undefined, + guardrails_add: values.guardrails_add || [], + guardrails_remove: values.guardrails_remove || [], + condition: values.model_condition + ? { model: values.model_condition } + : undefined, + }; if (isEditing && editingPolicy) { await updatePolicy(accessToken, editingPolicy.policy_id, data as PolicyUpdateRequest); @@ -263,13 +342,57 @@ const AddPolicyForm: React.FC = ({ value: p.policy_name, })); + // ── Mode Picker Step ────────────────────────────────────────────────────── + if (step === "pick_mode") { + return ( + + + + {selectedMode === "flow_builder" && ( + + )} + +
+ + +
+
+ ); + } + + // ── Simple Form Step ────────────────────────────────────────────────────── return (
= ({ /> - - Mode - - -
-
setPolicyMode("simple")} - style={{ - flex: 1, - padding: "12px 16px", - border: `2px solid ${policyMode === "simple" ? "#1890ff" : "#d9d9d9"}`, - borderRadius: 8, - cursor: "pointer", - backgroundColor: policyMode === "simple" ? "#e6f7ff" : "#fff", - }} - > - Simple Mode -
- - Add/remove guardrails directly - -
-
setPolicyMode("flow_builder")} - style={{ - flex: 1, - padding: "12px 16px", - border: `2px solid ${policyMode === "flow_builder" ? "#1890ff" : "#d9d9d9"}`, - borderRadius: 8, - cursor: "pointer", - backgroundColor: policyMode === "flow_builder" ? "#e6f7ff" : "#fff", - }} - > - Flow Builder - NEW -
- - Visual pipeline with conditional routing - -
-
- Inheritance @@ -364,75 +445,59 @@ const AddPolicyForm: React.FC = ({ /> - {policyMode === "simple" ? ( - <> - - Guardrails - + + Guardrails + - - + - - + - {resolvedGuardrails.length > 0 && ( - - - These are the final guardrails that will be applied (including inheritance): - -
- {resolvedGuardrails.map((g) => ( - - {g} - - ))} -
- - } - type="info" - showIcon - style={{ marginBottom: 16 }} - /> - )} - - ) : ( - <> - - Pipeline Flow - - - - + {resolvedGuardrails.length > 0 && ( + + + These are the final guardrails that will be applied (including inheritance): + +
+ {resolvedGuardrails.map((g) => ( + + {g} + + ))} +
+ + } + type="info" + showIcon + style={{ marginBottom: 16 }} + /> )}