diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx index b1768d5b81c..78085052655 100644 --- a/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx +++ b/ui/litellm-dashboard/src/components/policies/pipeline_flow_builder.tsx @@ -2,199 +2,30 @@ import React, { useState } from "react"; import { Select, Typography, Spin } from "antd"; import MessageManager from "@/components/molecules/message_manager"; import { Button, TextInput } from "@tremor/react"; -import { ArrowLeftIcon, PlusIcon } from "@heroicons/react/outline"; +import { ArrowLeftIcon } from "@heroicons/react/outline"; import { DotsVerticalIcon } from "@heroicons/react/solid"; -import { GuardrailPipeline, PipelineStep, PipelineTestResult, PolicyCreateRequest, PolicyUpdateRequest, Policy } from "./types"; +import { GuardrailPipeline, PipelineStep, PolicyCreateRequest, PolicyUpdateRequest, Policy } from "./types"; import { Guardrail } from "../guardrails/types"; -import { testPipelineCall, listPolicyVersions, createPolicyVersion, updatePolicyVersionStatus } from "../networking"; +import { listPolicyVersions, createPolicyVersion, updatePolicyVersionStatus } from "../networking"; import NotificationsManager from "../molecules/notifications_manager"; import { - getComplianceDatasetPrompts, - getFrameworks, -} from "../../data/compliancePrompts"; -import type { CompliancePrompt } from "../../data/compliancePrompts"; - -const TEST_SOURCE_QUICK = "quick_chat"; -const TEST_SOURCE_ALL = "__all__"; - -function getPromptsForTestSource(source: string): CompliancePrompt[] { - if (source === TEST_SOURCE_QUICK) return []; - if (source === TEST_SOURCE_ALL) return getComplianceDatasetPrompts(); - const fw = getFrameworks().find((f) => f.name === source); - return fw ? fw.categories.flatMap((c) => c.prompts) : []; -} + ACTION_OPTIONS, + ACTION_LABELS, + createDefaultStep, + insertStep, + removeStep, + updateStepAtIndex, + derivePipelineFromPolicy, + GuardrailIcon, + PlayIcon, + PassIcon, + FailIcon, + Connector, +} from "./pipeline_utils"; +import { PipelineTestPanel } from "./pipeline_test_drawer"; const { Text } = Typography; -const ACTION_OPTIONS = [ - { label: "Next Step", value: "next" }, - { label: "Allow", value: "allow" }, - { label: "Block", value: "block" }, - { label: "Custom Response", value: "modify_response" }, -]; - -const ACTION_LABELS: Record = { - allow: "Allow", - block: "Block", - next: "Next Step", - modify_response: "Custom Response", -}; - -function createDefaultStep(): PipelineStep { - return { - guardrail: "", - on_pass: "next", - on_fail: "block", - pass_data: false, - modify_response_message: null, - }; -} - -function insertStep(steps: PipelineStep[], atIndex: number): PipelineStep[] { - const newSteps = [...steps]; - newSteps.splice(atIndex, 0, createDefaultStep()); - return newSteps; -} - -function removeStep(steps: PipelineStep[], index: number): PipelineStep[] { - if (steps.length <= 1) return steps; - const newSteps = [...steps]; - newSteps.splice(index, 1); - return newSteps; -} - -function updateStepAtIndex( - steps: PipelineStep[], - index: number, - updated: Partial -): PipelineStep[] { - return steps.map((s, i) => (i === index ? { ...s, ...updated } : s)); -} - -/** - * Derives a pipeline from a policy. When the policy has a pipeline, use it. - * When it only has guardrails_add (legacy/simple form), convert those guardrails - * into pipeline steps in order. - */ -function derivePipelineFromPolicy(policy: Policy | null | undefined): GuardrailPipeline { - if (!policy) { - return { mode: "pre_call", steps: [createDefaultStep()] }; - } - if (policy.pipeline?.steps?.length) { - return policy.pipeline; - } - const guardrails = policy.guardrails_add || []; - if (guardrails.length > 0) { - return { - mode: policy.pipeline?.mode ?? "pre_call", - steps: guardrails.map((g) => ({ - guardrail: g, - on_pass: "next" as const, - on_fail: "block" as const, - pass_data: false, - modify_response_message: null, - })), - }; - } - return { mode: "pre_call", steps: [createDefaultStep()] }; -} - -// ───────────────────────────────────────────────────────────────────────────── -// Icons (matching the reference image) -// ───────────────────────────────────────────────────────────────────────────── - -const GuardrailIcon: React.FC = () => ( -
- - - - -
-); - -const PlayIcon: React.FC = () => ( -
- - - -
-); - -const PassIcon: React.FC = () => ( - - - - -); - -const FailIcon: React.FC = () => ( - - - -); - -// ───────────────────────────────────────────────────────────────────────────── -// Connector -// ───────────────────────────────────────────────────────────────────────────── - -interface ConnectorProps { - onInsert: () => void; -} - -const Connector: React.FC = ({ onInsert }) => ( -
-
- -
-
-); - // ───────────────────────────────────────────────────────────────────────────── // Step Card (editable) // ───────────────────────────────────────────────────────────────────────────── @@ -581,432 +412,6 @@ export const PipelineInfoDisplay: React.FC = ({ pipeli
); -// ───────────────────────────────────────────────────────────────────────────── -// Pipeline Test Panel (right drawer) -// ───────────────────────────────────────────────────────────────────────────── - -interface PipelineTestPanelProps { - pipeline: GuardrailPipeline; - accessToken: string | null; - onClose: () => void; -} - -const OUTCOME_STYLES: Record = { - pass: { bg: "#f0fdf4", color: "#16a34a", label: "PASS" }, - fail: { bg: "#fef2f2", color: "#dc2626", label: "FAIL" }, - error: { bg: "#fffbeb", color: "#d97706", label: "ERROR" }, -}; - -const TERMINAL_STYLES: Record = { - allow: { bg: "#f0fdf4", color: "#16a34a" }, - block: { bg: "#fef2f2", color: "#dc2626" }, - modify_response: { bg: "#eff6ff", color: "#2563eb" }, -}; - -interface ComplianceRunEntry { - prompt: CompliancePrompt; - result: PipelineTestResult | null; - error?: string; - matched: boolean; -} - -function complianceMatchExpected(expected: "pass" | "fail", terminalAction: string): boolean { - if (expected === "pass") { - return terminalAction === "allow" || terminalAction === "modify_response"; - } - return terminalAction === "block"; -} - -const testSourceOptions = [ - { value: TEST_SOURCE_QUICK, label: "Quick chat (custom message)" }, - ...getFrameworks().map((f) => ({ value: f.name, label: f.name })), - { value: TEST_SOURCE_ALL, label: "All compliance datasets" }, -]; - -const PipelineTestPanel: React.FC = ({ - pipeline, - accessToken, - onClose, -}) => { - const [testSource, setTestSource] = useState(TEST_SOURCE_QUICK); - const [testMessage, setTestMessage] = useState("Hello, can you help me?"); - const [isRunning, setIsRunning] = useState(false); - const [result, setResult] = useState(null); - const [error, setError] = useState(null); - const [complianceResults, setComplianceResults] = useState([]); - - const isQuickChat = testSource === TEST_SOURCE_QUICK; - const promptsForSource = getPromptsForTestSource(testSource); - const isDataset = promptsForSource.length > 0; - - const handleRunTest = async () => { - if (!accessToken) return; - - const emptySteps = pipeline.steps.filter((s) => !s.guardrail); - if (emptySteps.length > 0) { - setError("All steps must have a guardrail selected"); - return; - } - - setError(null); - setIsRunning(true); - setResult(null); - setComplianceResults([]); - - if (isQuickChat) { - try { - const data = await testPipelineCall( - accessToken, - pipeline, - [{ role: "user", content: testMessage }] - ); - setResult(data); - } catch (e) { - setError(e instanceof Error ? e.message : String(e)); - } finally { - setIsRunning(false); - } - return; - } - - const entries: ComplianceRunEntry[] = []; - for (const prompt of promptsForSource) { - try { - const data = await testPipelineCall(accessToken, pipeline, [ - { role: "user", content: prompt.prompt }, - ]); - const matched = complianceMatchExpected(prompt.expectedResult, data.terminal_action); - entries.push({ prompt, result: data, matched }); - } catch (e) { - const errMsg = e instanceof Error ? e.message : String(e); - entries.push({ - prompt, - result: null, - error: errMsg, - matched: false, - }); - } - } - setComplianceResults(entries); - setIsRunning(false); - }; - - return ( -
- {/* Panel header */} -
- Test Pipeline - -
- - {/* Input section */} -
- -