feat: improve template selection reliability and prompt clarity

AshishKumar4 · AshishKumar4 · commit 8ed1a0840b70 · 2025-11-04T14:29:28.000-05:00
- Added retry mechanism with configurable retry count to handle transient failures
- Enhanced prompt formatting with clearer template list presentation and validation rules
- Improved error handling to retry on failures before falling back to null selection
diff --git a/worker/agents/planning/templateSelector.ts b/worker/agents/planning/templateSelector.ts
@@ -20,23 +20,25 @@ interface SelectTemplateArgs {
 /**
  * Uses AI to select the most suitable template for a given query.
  */
-export async function selectTemplate({ env, query, availableTemplates, inferenceContext, images }: SelectTemplateArgs): Promise<TemplateSelection> {
+export async function selectTemplate({ env, query, availableTemplates, inferenceContext, images }: SelectTemplateArgs, retryCount: number = 3): Promise<TemplateSelection> {
     if (availableTemplates.length === 0) {
         logger.info("No templates available for selection.");
         return { selectedTemplateName: null, reasoning: "No templates were available to choose from.", useCase: null, complexity: null, styleSelection: null, projectName: '' };
     }
 
     try {
-        logger.info("Asking AI to select a template", { 
+        logger.info(`Asking AI to select a template for the ${retryCount} time`, { 
             query, 
             queryLength: query.length,
             imagesCount: images?.length || 0,
             availableTemplates: availableTemplates.map(t => t.name),
             templateCount: availableTemplates.length 
         });
 
+        const validTemplateNames = availableTemplates.map(t => t.name);
+
         const templateDescriptions = availableTemplates.map((t, index) =>
-            `- Template #${index + 1} \n Name - ${t.name} \n Language: ${t.language}, Frameworks: ${t.frameworks?.join(', ') || 'None'}\n ${t.description.selection}`
+            `### Template #${index + 1} \n Name - ${t.name} \n Language: ${t.language}, Frameworks: ${t.frameworks?.join(', ') || 'None'}\n Description: \`\`\`${t.description.selection}\`\`\``
         ).join('\n\n');
 
         const systemPrompt = `You are an Expert Software Architect at Cloudflare specializing in template selection for rapid development. Your task is to select the most suitable starting template based on user requirements.
@@ -81,13 +83,16 @@ Reasoning: "Social template provides user interactions, content sharing, and com
 ## RULES:
 - ALWAYS select a template (never return null)
 - Ignore misleading template names - analyze actual features
+- **ONLY** Choose from the list of available templates
 - Focus on functionality over naming conventions
 - Provide clear, specific reasoning for selection`
 
         const userPrompt = `**User Request:** "${query}"
 
-**Available Templates:**
-${templateDescriptions}
+## **Available Templates:**
+**ONLY** These template names are available for selection: ${validTemplateNames.join(', ')}
+
+Template detail: ${templateDescriptions}
 
 **Task:** Select the most suitable template and provide:
 1. Template name (exact match from list)
@@ -122,7 +127,6 @@ ENTROPY SEED: ${generateSecureToken(64)} - for unique results`;
             maxTokens: 2000,
         });
 
-
         logger.info(`AI template selection result: ${selection.selectedTemplateName || 'None'}, Reasoning: ${selection.reasoning}`);
         return selection;
 
@@ -131,6 +135,10 @@ ENTROPY SEED: ${generateSecureToken(64)} - for unique results`;
         if (error instanceof RateLimitExceededError || error instanceof SecurityError) {
             throw error;
         }
+
+        if (retryCount > 0) {
+            return selectTemplate({ env, query, availableTemplates, inferenceContext, images }, retryCount - 1);
+        }
         // Fallback to no template selection in case of error
         return { selectedTemplateName: null, reasoning: "An error occurred during the template selection process.", useCase: null, complexity: null, styleSelection: null, projectName: '' };
     }
