From e432c211ae5bc5482b3c101468a473991b523af8 Mon Sep 17 00:00:00 2001
From: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
Date: Thu, 19 Feb 2026 10:59:46 -0800
Subject: [PATCH] add Use AI button and template loading callback to
 PolicyTemplates

---
 .../src/components/policies/policy_templates.tsx  | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx
index fdb9247ba9a..6951b368ef3 100644
--- a/ui/litellm-dashboard/src/components/policies/policy_templates.tsx
+++ b/ui/litellm-dashboard/src/components/policies/policy_templates.tsx
@@ -116,6 +116,8 @@ const PolicyTemplateCard: React.FC<PolicyTemplateCardProps> = ({
 
 interface PolicyTemplatesProps {
   onUseTemplate: (templateData: any) => void;
+  onOpenAiSuggestion: () => void;
+  onTemplatesLoaded?: (templates: any[]) => void;
   accessToken: string | null;
 }
 
@@ -128,7 +130,7 @@ const iconMap: Record<string, React.ComponentType<React.SVGProps<SVGSVGElement>>
   CheckCircleIcon: CheckCircleIcon,
 };
 
-const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, accessToken }) => {
+const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, onOpenAiSuggestion, onTemplatesLoaded, accessToken }) => {
   const [templates, setTemplates] = useState<any[]>([]);
   const [isLoading, setIsLoading] = useState(false);
   const [selectedTags, setSelectedTags] = useState<Set<string>>(new Set());
@@ -179,6 +181,7 @@ const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, access
       try {
         const data = await getPolicyTemplates(accessToken);
         setTemplates(data);
+        onTemplatesLoaded?.(data);
       } catch (error) {
         console.error("Error fetching policy templates:", error);
         message.error("Failed to fetch policy templates");
@@ -210,6 +213,16 @@ const PolicyTemplates: React.FC<PolicyTemplatesProps> = ({ onUseTemplate, access
             guardrails for your organization.
           </p>
         </div>
+        <Button
+          type="default"
+          onClick={onOpenAiSuggestion}
+          className="flex items-center gap-1.5"
+        >
+          <svg className="w-4 h-4" viewBox="0 0 16 16" fill="currentColor">
+            <path d="M8 1l1.5 3.5L13 6l-3.5 1.5L8 11 6.5 7.5 3 6l3.5-1.5L8 1zm4 7l.75 1.75L14.5 10.5l-1.75.75L12 13l-.75-1.75L9.5 10.5l1.75-.75L12 8zM4 9l.75 1.75L6.5 11.5l-1.75.75L4 14l-.75-1.75L1.5 11.5l1.75-.75L4 9z" />
+          </svg>
+          Use AI to find templates
+        </Button>
       </div>
 
       <div className="flex gap-6">