From 775490bc90af8cd60507749b60bfa7ed61e440c2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 19 Feb 2026 14:07:09 -0800 Subject: [PATCH] feat(ui): add test playground to AI policy suggestion modal --- .../policies/ai_suggestion_modal.tsx | 622 ++++++++++++++---- 1 file changed, 485 insertions(+), 137 deletions(-) diff --git a/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx index 623c291f079..60b49a7755e 100644 --- a/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx +++ b/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx @@ -1,13 +1,24 @@ import React, { useState, useEffect } from "react"; -import { Modal, Spin, Checkbox, Select } from "antd"; -import { Button } from "@tremor/react"; -import { suggestPolicyTemplates, modelHubCall } from "../networking"; +import { Modal, Spin, Checkbox, Select, Input, Typography, Tooltip } from "antd"; +import { Button, Card } from "@tremor/react"; +import { CheckCircleOutlined, CloseCircleOutlined, InfoCircleOutlined, DownOutlined, RightOutlined } from "@ant-design/icons"; +import { suggestPolicyTemplates, modelHubCall, testPolicyTemplate, enrichPolicyTemplate } from "../networking"; + +const { TextArea } = Input; +const { Text } = Typography; interface SuggestedTemplate { template_id: string; reason: string; } +interface GuardrailTestResult { + guardrail_name: string; + action: string; + output_text: string; + details: string; +} + interface AiSuggestionModalProps { visible: boolean; onSelectTemplates: (templates: any[]) => void; @@ -34,6 +45,17 @@ const AiSuggestionModal: React.FC = ({ const [selectedModel, setSelectedModel] = useState(undefined); const [availableModels, setAvailableModels] = useState([]); const [isLoadingModels, setIsLoadingModels] = useState(false); + // Test panel state + const [showTestPanel, setShowTestPanel] = useState(false); + const [testInputText, setTestInputText] = useState(""); + const [isTestLoading, setIsTestLoading] = useState(false); + const [testResults, setTestResults] = useState(null); + const [testOverallAction, setTestOverallAction] = useState(null); + const [collapsedResults, setCollapsedResults] = useState>(new Set()); + // Enrichment state for competitor templates + const [enrichedDefs, setEnrichedDefs] = useState>({}); + const [isEnriching, setIsEnriching] = useState(false); + const [enrichBrandName, setEnrichBrandName] = useState(""); useEffect(() => { if (visible && availableModels.length === 0) { @@ -67,6 +89,15 @@ const AiSuggestionModal: React.FC = ({ setExplanation(null); setSelectedIds(new Set()); setSelectedModel(undefined); + setShowTestPanel(false); + setTestInputText(""); + setIsTestLoading(false); + setTestResults(null); + setTestOverallAction(null); + setCollapsedResults(new Set()); + setEnrichedDefs({}); + setIsEnriching(false); + setEnrichBrandName(""); }; const handleCancel = () => { @@ -126,6 +157,11 @@ const AiSuggestionModal: React.FC = ({ setSuggestions(null); setExplanation(null); setSelectedIds(new Set()); + setShowTestPanel(false); + setTestInputText(""); + setTestResults(null); + setTestOverallAction(null); + setCollapsedResults(new Set()); }; const handleUseSelected = () => { @@ -149,14 +185,433 @@ const AiSuggestionModal: React.FC = ({ const getTemplateById = (id: string) => allTemplates.find((t) => t.id === id); + const toggleResultCollapse = (name: string) => { + setCollapsedResults((prev) => { + const next = new Set(prev); + if (next.has(name)) next.delete(name); + else next.add(name); + return next; + }); + }; + + const getSelectedTemplatesNeedingEnrichment = (): any[] => { + return Array.from(selectedIds) + .map((id) => getTemplateById(id)) + .filter((t) => t?.llm_enrichment); + }; + + const needsEnrichment = getSelectedTemplatesNeedingEnrichment().length > 0; + + const getAllSelectedGuardrailDefs = (): any[] => { + const defs: any[] = []; + for (const id of selectedIds) { + // Use enriched defs if available, otherwise use template's original + if (enrichedDefs[id]) { + defs.push(...enrichedDefs[id]); + } else { + const t = getTemplateById(id); + if (t?.guardrailDefinitions) { + defs.push(...t.guardrailDefinitions); + } + } + } + return defs; + }; + + const handleEnrichCompetitors = async () => { + if (!accessToken || !selectedModel) return; + const templatesToEnrich = getSelectedTemplatesNeedingEnrichment(); + if (templatesToEnrich.length === 0) return; + + setIsEnriching(true); + try { + for (const template of templatesToEnrich) { + const paramName = template.llm_enrichment.parameter; + const result = await enrichPolicyTemplate( + accessToken, + template.id, + { [paramName]: enrichBrandName }, + selectedModel, + ); + setEnrichedDefs((prev) => ({ + ...prev, + [template.id]: result.guardrailDefinitions, + })); + } + } catch (e) { + console.error("Failed to enrich templates:", e); + } finally { + setIsEnriching(false); + } + }; + + const handleRunTest = async () => { + if (!accessToken || !testInputText.trim()) return; + const allDefs = getAllSelectedGuardrailDefs(); + if (allDefs.length === 0) return; + + setIsTestLoading(true); + setTestResults(null); + setTestOverallAction(null); + setCollapsedResults(new Set()); + + try { + const result = await testPolicyTemplate(accessToken, allDefs, testInputText); + setTestResults(result.results || []); + setTestOverallAction(result.overall_action || "passed"); + } catch { + setTestResults([]); + setTestOverallAction("error"); + } finally { + setIsTestLoading(false); + } + }; + + const handleTestKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey && !e.ctrlKey && !e.metaKey) { + e.preventDefault(); + handleRunTest(); + } + }; + const showResults = suggestions !== null && !isLoading; + // Helper to render the suggestions list (reused in both layouts) + const renderSuggestionsList = () => { + if (!suggestions || suggestions.length === 0) { + return ( +
+ + + +

No matching templates found

+

Try adjusting your examples or description.

+
+ ); + } + + return ( +
+ {suggestions.map((suggestion) => { + const template = getTemplateById(suggestion.template_id); + if (!template) return null; + const isSelected = selectedIds.has(suggestion.template_id); + return ( +
+
toggleTemplate(suggestion.template_id)} + > +
+ toggleTemplate(suggestion.template_id)} + className="mt-0.5" + /> +
+
+ + {template.title} + + {template.complexity && ( + + {template.complexity} + + )} +
+

+ {template.description} +

+
+ {template.guardrails && template.guardrails.slice(0, 4).map((g: string) => ( + + {g} + + ))} + {template.guardrails && template.guardrails.length > 4 && ( + + +{template.guardrails.length - 4} more + + )} +
+
+ +

+ {suggestion.reason} +

+
+
+
+
+
+ ); + })} + + {/* Explanation */} + {explanation && ( +
+
+ + + Why these templates + +
+

{explanation}

+
+ )} +
+ ); + }; + + // Helper to render the test panel + const renderTestPanel = () => ( +
+ {/* Test header */} +
+
+

Test Guardrails

+ +
+
+ {Array.from(selectedIds).map((id) => { + const t = getTemplateById(id); + return t ? ( + + {t.title} + + ) : null; + })} +
+

+ {getAllSelectedGuardrailDefs().length} guardrails across {selectedIds.size} template{selectedIds.size !== 1 ? "s" : ""} +

+
+ + {/* Enrichment section for competitor templates */} + {needsEnrichment && Object.keys(enrichedDefs).length > 0 && ( +
+
+ + + Competitor names loaded for {enrichBrandName} + +
+
+ )} + {needsEnrichment && Object.keys(enrichedDefs).length === 0 && ( +
+
+ + + + + Competitor template requires your brand name to discover competitors + +
+
+ setEnrichBrandName(e.target.value)} + onPressEnter={() => enrichBrandName.trim() && handleEnrichCompetitors()} + className="flex-1" + /> + +
+
+ )} + + {/* Input */} +
+
+
+
+ + + + +
+ Characters: {testInputText.length} +
+