a million and one things

2 files changed, 32 insertions, 22 deletions

diff --git a/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts
index ef7dbc7ab..d11f05766 100644
--- a/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts
+++ b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts
@@ -53,11 +53,11 @@ export class TemplateManager {
         }
     }
 
-    createDocsFromTemplate = action((dv: DataVizBox, template: Template, csvColumns: Col[], debug: boolean = false) => {
+    createDocsFromTemplate = action((dv: DataVizBox, template: Template, cols: Col[], debug: boolean = false) => {
         const fields = Array.from(Object.keys(dv.records[0]));
 
-        const processContent = (content: { [title: string]: string }) => {
-            const templateCopy = template.cloneBase();
+        const processContent = async (content: { [title: string]: string }) => {
+            const templateCopy = template.clone();
 
             fields
                 .filter(title => title)
@@ -67,14 +67,22 @@ export class TemplateManager {
                 });
 
             const gptFunc = (type: TemplateFieldType) => (type === TemplateFieldType.VISUAL ? TemplateMenuAIUtils.renderGPTImageCall : TemplateMenuAIUtils.renderGPTTextCall);
-            const gptPromises = csvColumns
-                .filter(field => field.type !== TemplateFieldType.UNSET && field.AIGenerated)
-                .map(field => {
-                    const templateField = templateCopy.getFieldByTitle(field.title);
-                    if (templateField !== undefined) {
-                        return gptFunc(field.type)(templateCopy, field, templateField.getID);
-                    }
-                });
+            const applyGPTContent = async () => {
+                const promises = cols
+                    .filter(field => field.AIGenerated)
+                    .map(field => {
+                        const templateField = templateCopy.getFieldByTitle(field.title);
+                        if (templateField !== undefined) {
+                            return gptFunc(field.type)(templateCopy, field, templateField.getID);
+                        }
+                        return null; 
+                    })
+                    .filter(p => p !== null);
+            
+                await Promise.all(promises);
+            };
+
+            await applyGPTContent();
 
             return templateCopy.getRenderedDoc();
         };
@@ -90,6 +98,7 @@ export class TemplateManager {
                       {} as { [title: string]: string }
                   )
               );
+
         return Promise.all(rowContents.map(processContent)).then(
             action(renderedDocs => {
                 return renderedDocs;
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);