diff --git a/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx b/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx index ce1e5bc22e9..67948a075cd 100644 --- a/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx +++ b/ui/litellm-dashboard/src/components/policies/template_parameter_modal.tsx @@ -1,6 +1,7 @@ import React, { useState, useEffect } from "react"; -import { Modal, Spin } from "antd"; +import { Modal, Spin, Radio, Select } from "antd"; import { Button, TextInput } from "@tremor/react"; +import { modelHubCall, enrichPolicyTemplateStream } from "../networking"; interface TemplateParameter { name: string; @@ -13,9 +14,13 @@ interface TemplateParameter { interface TemplateParameterModalProps { visible: boolean; template: any; - onConfirm: (parameters: Record) => void; + onConfirm: ( + parameters: Record, + enrichmentOptions?: { model?: string; competitors?: string[] } + ) => void; onCancel: () => void; isLoading?: boolean; + accessToken: string; } const TemplateParameterModal: React.FC = ({ @@ -24,10 +29,24 @@ const TemplateParameterModal: React.FC = ({ onConfirm, onCancel, isLoading = false, + accessToken, }) => { const [parameterValues, setParameterValues] = useState>({}); + const [competitorMode, setCompetitorMode] = useState<"ai" | "manual">("ai"); + const [selectedModel, setSelectedModel] = useState(undefined); + const [availableModels, setAvailableModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [competitorTags, setCompetitorTags] = useState([]); + const [variationsMap, setVariationsMap] = useState>({}); + const [isGenerating, setIsGenerating] = useState(false); const parameters: TemplateParameter[] = template?.parameters || []; + const hasEnrichment = !!template?.llm_enrichment; + const enrichmentParam = hasEnrichment ? template.llm_enrichment.parameter : null; + + const nonEnrichmentParams = hasEnrichment + ? parameters.filter((p) => p.name !== enrichmentParam) + : parameters; useEffect(() => { if (visible && template) { @@ -36,15 +55,85 @@ const TemplateParameterModal: React.FC = ({ initial[p.name] = ""; }); setParameterValues(initial); + setCompetitorMode("ai"); + setSelectedModel(undefined); + setCompetitorTags([]); + setVariationsMap({}); + setIsGenerating(false); } }, [visible, template]); - const allRequiredFilled = parameters + useEffect(() => { + if (visible && hasEnrichment && competitorMode === "ai" && availableModels.length === 0) { + loadModels(); + } + }, [visible, hasEnrichment, competitorMode]); + + const loadModels = async () => { + if (!accessToken) return; + setIsLoadingModels(true); + try { + const fetchedModels = await modelHubCall(accessToken); + if (fetchedModels?.data?.length > 0) { + const models = fetchedModels.data + .map((item: any) => item.model_group as string) + .sort(); + setAvailableModels(models); + } + } catch (error) { + console.error("Error fetching models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + const handleGenerateNames = async () => { + if (!accessToken || !selectedModel || !template) return; + const brandName = (parameterValues[enrichmentParam || "brand_name"] || "").trim(); + if (!brandName) return; + + setIsGenerating(true); + setCompetitorTags([]); + setVariationsMap({}); + try { + await enrichPolicyTemplateStream( + accessToken, + template.id, + parameterValues, + selectedModel, + (name) => { + setCompetitorTags((prev) => [...prev, name]); + }, + (result) => { + setCompetitorTags(result.competitors); + setVariationsMap(result.competitor_variations || {}); + setIsGenerating(false); + }, + (error) => { + console.error("Streaming error:", error); + setIsGenerating(false); + } + ); + } catch (error) { + console.error("Error generating competitor names:", error); + setIsGenerating(false); + } + }; + + const allNonEnrichmentFilled = nonEnrichmentParams .filter((p) => p.required) .every((p) => (parameterValues[p.name] || "").trim().length > 0); + const brandNameFilled = enrichmentParam + ? (parameterValues[enrichmentParam] || "").trim().length > 0 + : true; + + const canContinue = hasEnrichment + ? allNonEnrichmentFilled && brandNameFilled && competitorTags.length > 0 + : allNonEnrichmentFilled && brandNameFilled; + const handleConfirm = () => { - onConfirm(parameterValues); + onConfirm(parameterValues, { competitors: competitorTags }); }; return ( @@ -53,15 +142,13 @@ const TemplateParameterModal: React.FC = ({

{template?.title}

- {template?.llm_enrichment - ? "Enter your brand name to auto-discover competitors and configure guardrails" - : "Configure template parameters"} + Configure competitor blocking for your brand

} open={visible} onCancel={onCancel} - width={500} + width={550} footer={[ , ]} >
- {parameters.map((param) => ( + {nonEnrichmentParams.map((param) => (
))} - {template?.llm_enrichment && ( -
-

- This template uses AI to automatically discover your competitors and configure - guardrails. An onboarded LLM will be called to identify competitor names. -

-
+ {hasEnrichment && ( + <> +
+ + setCompetitorMode(e.target.value)} + className="w-full" + > +
+ + ✨ Use AI + + + Enter Manually + +
+
+
+ + {/* Brand Name */} +
+ + + setParameterValues((prev) => ({ + ...prev, + [enrichmentParam || "brand_name"]: e.target.value, + })) + } + /> +
+ + {competitorMode === "ai" && ( + <> +
+ + setCompetitorTags(values)} + tokenSeparators={[","]} + open={false} + suffixIcon={null} + /> +

+ Type a name and press Enter to add. Click ✕ to remove. +

+ {Object.keys(variationsMap).length > 0 && ( +

+ ✓ {Object.values(variationsMap).flat().length} alternate spellings & variations auto-generated for guardrail matching +

+ )} +
+ )} - {isLoading && ( -
- - - {template?.llm_enrichment - ? "Using AI to discover competitors..." - : "Processing template..."} - -
- )} + {!hasEnrichment && + parameters.map((param) => ( +
+ + + setParameterValues((prev) => ({ + ...prev, + [param.name]: e.target.value, + })) + } + /> +
+ ))}
);