diff --git a/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx b/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx new file mode 100644 index 00000000000..aff629ed4f4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/policies/ai_suggestion_modal.tsx @@ -0,0 +1,315 @@ +import React, { useState } from "react"; +import { Modal, Spin, Checkbox } from "antd"; +import { Button, TextInput } from "@tremor/react"; +import { suggestPolicyTemplates } from "../networking"; + +interface SuggestedTemplate { + template_id: string; + reason: string; +} + +interface AiSuggestionModalProps { + visible: boolean; + onSelectTemplates: (templates: any[]) => void; + onCancel: () => void; + accessToken: string | null; + allTemplates: any[]; +} + +const MAX_EXAMPLES = 4; + +const AiSuggestionModal: React.FC = ({ + visible, + onSelectTemplates, + onCancel, + accessToken, + allTemplates, +}) => { + const [attackExamples, setAttackExamples] = useState([""]); + const [description, setDescription] = useState(""); + const [isLoading, setIsLoading] = useState(false); + const [suggestions, setSuggestions] = useState(null); + const [explanation, setExplanation] = useState(null); + const [selectedIds, setSelectedIds] = useState>(new Set()); + + const resetState = () => { + setAttackExamples([""]); + setDescription(""); + setIsLoading(false); + setSuggestions(null); + setExplanation(null); + setSelectedIds(new Set()); + }; + + const handleCancel = () => { + resetState(); + onCancel(); + }; + + const handleAddExample = () => { + if (attackExamples.length < MAX_EXAMPLES) { + setAttackExamples([...attackExamples, ""]); + } + }; + + const handleRemoveExample = (index: number) => { + setAttackExamples(attackExamples.filter((_, i) => i !== index)); + }; + + const handleExampleChange = (index: number, value: string) => { + const updated = [...attackExamples]; + updated[index] = value; + setAttackExamples(updated); + }; + + const hasInput = + attackExamples.some((e) => e.trim().length > 0) || + description.trim().length > 0; + + const handleSuggest = async () => { + if (!accessToken || !hasInput) return; + + setIsLoading(true); + try { + const result = await suggestPolicyTemplates( + accessToken, + attackExamples, + description + ); + setSuggestions(result.selected_templates || []); + setExplanation(result.explanation || null); + // Pre-select all suggested templates + setSelectedIds( + new Set( + (result.selected_templates || []).map( + (s: SuggestedTemplate) => s.template_id + ) + ) + ); + } catch { + setSuggestions([]); + setExplanation("Failed to get suggestions. Please try again."); + } finally { + setIsLoading(false); + } + }; + + const handleBack = () => { + setSuggestions(null); + setExplanation(null); + setSelectedIds(new Set()); + }; + + const handleUseSelected = () => { + const selected = allTemplates.filter((t) => selectedIds.has(t.id)); + resetState(); + onSelectTemplates(selected); + }; + + const toggleTemplate = (templateId: string) => { + setSelectedIds((prev) => { + const next = new Set(prev); + if (next.has(templateId)) { + next.delete(templateId); + } else { + next.add(templateId); + } + return next; + }); + }; + + const getTemplateById = (id: string) => + allTemplates.find((t) => t.id === id); + + const showResults = suggestions !== null && !isLoading; + + return ( + +

AI Policy Suggestion

+

+ {showResults + ? "Select which templates to use" + : "Describe what you want to block and we'll suggest the best policy templates"} +

+ + } + open={visible} + onCancel={handleCancel} + width={600} + footer={ + showResults + ? [ + , + , + ] + : [ + , + , + ] + } + > + {!showResults ? ( +
+
+ + {attackExamples.map((example, index) => ( +
+ handleExampleChange(index, e.target.value)} + /> + {attackExamples.length > 1 && ( + + )} +
+ ))} + {attackExamples.length < MAX_EXAMPLES && ( + + )} +
+ +
+ + setDescription(e.target.value)} + /> +
+ +
+

+ This feature uses AI to analyze your requirements and match them + against available policy templates. An onboarded LLM will be + called to identify the best templates. +

+
+ + {isLoading && ( +
+ + + Using AI to find matching policy templates... + +
+ )} +
+ ) : ( +
+ {suggestions && suggestions.length > 0 ? ( + <> +
+ {suggestions.length} template + {suggestions.length !== 1 ? "s" : ""} suggested +
+
+ {suggestions.map((suggestion) => { + const template = getTemplateById(suggestion.template_id); + if (!template) return null; + return ( +
toggleTemplate(suggestion.template_id)} + > +
+ + toggleTemplate(suggestion.template_id) + } + /> +
+
+ {template.title} +
+
+ {template.description} +
+
+ {suggestion.reason} +
+
+
+
+ ); + })} +
+ + ) : ( +
+

No matching templates found for your requirements.

+

+ Try adjusting your examples or description. +

+
+ )} + + {explanation && ( +
+

{explanation}

+
+ )} +
+ )} +
+ ); +}; + +export default AiSuggestionModal;