1 files changed, 12 insertions, 11 deletions

diff --git a/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts
index 446fe3442..9bc2bfce2 100644
--- a/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts
+++ b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts
@@ -27,12 +27,13 @@ export class TemplateMenuAIUtils {
         }
     };
 
-    public static renderGPTImageCall = async (template: Template, col: Col, fieldNumber: number | undefined): Promise<boolean> => {
-        const generateAndLoadImage = async (fieldNum: string, column: Col, prompt: string) => {
+    public static renderGPTImageCall = async (template: Template, col: Col, fieldNumber: number): Promise<boolean> => {
+        const generateAndLoadImage = async (id: number, prompt: string) => {
             const url = await this.generateGPTImage(prompt);
-            const field: TemplateField = template.getFieldByID(Number(fieldNum));
+            var field: TemplateField = template.getFieldByID(id);
 
             field.setContent(url ?? '', ViewType.IMG);
+            field = template.getFieldByID(id);
             field.setTitle(col.title);
         };
 
@@ -40,14 +41,14 @@ export class TemplateMenuAIUtils {
 
         try {
             const sysPrompt =
-                'Your job is to create a prompt for an AI image generator to help it generate an image based on existing content in a template and a user prompt. Your prompt should focus heavily on visual elements to help the image generator; avoid unecessary info that might distract it. ONLY INCLUDE THE PROMPT, NO OTHER TEXT OR EXPLANATION. The existing content is as follows: ' +
+                `#${Math.random() * 100}: Your job is to create a prompt for an AI image generator to help it generate an image based on existing content in a template and a user prompt. Your prompt should focus heavily on visual elements to help the image generator; avoid unecessary info that might distract it. ONLY INCLUDE THE PROMPT, NO OTHER TEXT OR EXPLANATION. The existing content is as follows: ` +
                 fieldContent +
                 ' **** The user prompt is: ' +
                 col.desc;
 
             const prompt = await gptAPICall(sysPrompt, GPTCallType.COMPLETEPROMPT);
 
-            await generateAndLoadImage(String(fieldNumber), col, prompt);
+            await generateAndLoadImage(fieldNumber, prompt);
         } catch (e) {
             console.log(e);
         }
@@ -104,20 +105,20 @@ export class TemplateMenuAIUtils {
      * @returns a doc containing the fully rendered template
      */
     public static applyGPTContentToTemplate = async (template: Template, assignments: { [field: string]: Col }): Promise<Template | undefined> => {
-        const GPTTextCalls = Object.entries(assignments).filter(([, col]) => col.type === TemplateFieldType.TEXT && !col.AIGenerated);
-        const GPTIMGCalls = Object.entries(assignments).filter(([, col]) => col.type === TemplateFieldType.VISUAL && !col.AIGenerated);
+        const GPTTextCalls = Object.entries(assignments).filter(([, col]) => col.type === TemplateFieldType.TEXT && col.AIGenerated);
+        const GPTIMGCalls = Object.entries(assignments).filter(([, col]) => col.type === TemplateFieldType.VISUAL && col.AIGenerated);
 
         if (GPTTextCalls.length) {
-            const promises = GPTTextCalls.map(([str, col]) => {
-                return TemplateMenuAIUtils.renderGPTTextCall(template, col, Number(str));
+            const promises = GPTTextCalls.map(([id, col]) => {
+                return TemplateMenuAIUtils.renderGPTTextCall(template, col, Number(id));
             });
 
             await Promise.all(promises);
         }
 
         if (GPTIMGCalls.length) {
-            const promises = GPTIMGCalls.map(async ([fieldNum, col]) => {
-                return TemplateMenuAIUtils.renderGPTImageCall(template, col, Number(fieldNum));
+            const promises = GPTIMGCalls.map(async ([id, col]) => {
+                return TemplateMenuAIUtils.renderGPTImageCall(template, col, Number(id));
             });
 
             await Promise.all(promises);