Merge branch 'master' into task_nodes_aarav

2 files changed, 213 insertions, 0 deletions

diff --git a/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts
new file mode 100644
index 000000000..526fcf9c4
--- /dev/null
+++ b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateManager.ts
@@ -0,0 +1,91 @@
+import { action, makeAutoObservable } from 'mobx';
+import { Col } from '../DocCreatorMenu';
+import { FieldSettings, TemplateField } from '../TemplateFieldTypes/TemplateField';
+import { Template } from '../Template';
+import { NumListCast } from '../../../../../../fields/Doc';
+import { DataVizBox } from '../../DataVizBox';
+import { TemplateFieldType } from '../TemplateBackend';
+import { TemplateMenuAIUtils } from './TemplateMenuAIUtils';
+
+export type Conditional = {
+    field: string;
+    operator: '=' | '>' | '<' | 'contains';
+    condition: string;
+    target: string;
+    attribute: string;
+    value: string;
+};
+
+export class TemplateManager {
+    _templates: Template[] = [];
+
+    _conditionalFieldLogic: Record<string, Conditional[]> = {};
+
+    constructor(templateSettings: FieldSettings[]) {
+        makeAutoObservable(this);
+        this._templates = templateSettings.map(settings => new Template(settings));
+    }
+
+    getValidTemplates = (cols: Col[]) => this._templates.filter(template => template.isValidTemplate(cols));
+
+    addTemplate = (newTemplate: Template) => this._templates.push(newTemplate);
+
+    removeTemplate = (template: Template) => {
+        if (this._templates.includes(template)) {
+            this._templates.splice(this._templates.indexOf(template), 1);
+        }
+        template.cleanup();
+    };
+
+    addFieldCondition = (fieldTitle: string, condition: Conditional) => {
+        if (this._conditionalFieldLogic[fieldTitle] === undefined) {
+            this._conditionalFieldLogic[fieldTitle] = [condition];
+        } else {
+            this._conditionalFieldLogic[fieldTitle].push(condition);
+        }
+    };
+
+    removeFieldCondition = (fieldTitle: string, condition: Conditional) => (this._conditionalFieldLogic[fieldTitle] = this._conditionalFieldLogic[fieldTitle]?.filter(cond => cond !== condition));
+
+    addDataField = (title: string) => this._templates.forEach(template => template.addDataField(title));
+
+    removeDataField = (title: string) => this._templates.forEach(template => template.removeDataField(title));
+
+    createDocsFromTemplate = action((dv: DataVizBox, template: Template, cols: Col[], debug: boolean = false) => {
+        const csvFields = Array.from(Object.keys(dv.records[0]));
+
+        const processContent = (content: { [title: string]: string }) => {
+            const templateCopy = template.clone();
+
+            csvFields
+                .filter(title => title)
+                .forEach(title => {
+                    const field = templateCopy.getFieldByTitle(title);
+                    field?.setContent(content[title], field.viewType);
+                });
+
+            const gptFunc = (type: TemplateFieldType) => (type === TemplateFieldType.VISUAL ? TemplateMenuAIUtils.renderGPTImageCall : TemplateMenuAIUtils.renderGPTTextCall);
+
+            const generateGptContent = cols
+                .map(field => ({ field, templateField: field?.AIGenerated && templateCopy.getFieldByTitle(field.title) }))
+                .filter(({ templateField }) => templateField instanceof TemplateField)
+                .map(({ field, templateField }) => gptFunc(field.type)(templateCopy, field, (templateField as TemplateField).getID));
+
+            return Promise.all(generateGptContent).then(() => templateCopy.applyConditionalLogic(this._conditionalFieldLogic));
+        };
+
+        const rowContents = debug
+            ? [{}, {}, {}, {}]
+            : NumListCast(dv.layoutDoc.dataViz_selectedRows).map(row =>
+                  csvFields.reduce(
+                      (values, col) => {
+                          values[col] = dv.records[row][col];
+                          return values;
+                      },
+                      {} as { [title: string]: string }
+                  )
+              );
+
+        return Promise.all(rowContents.map(processContent));
+    });
+}
diff --git a/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts
new file mode 100644
index 000000000..08818dd6c
--- /dev/null
+++ b/src/client/views/nodes/DataVizBox/DocCreatorMenu/Backend/TemplateMenuAIUtils.ts
@@ -0,0 +1,122 @@
+import { ClientUtils } from '../../../../../../ClientUtils';
+import { Networking } from '../../../../../Network';
+import { gptImageCall, gptAPICall, GPTCallType } from '../../../../../apis/gpt/GPT';
+import { Col } from '../DocCreatorMenu';
+import { TemplateFieldSize, TemplateFieldType } from '../TemplateBackend';
+import { ViewType } from '../TemplateFieldTypes/TemplateField';
+import { Template } from '../Template';
+import { Upload } from '../../../../../../server/SharedMediaTypes';
+
+export class TemplateMenuAIUtils {
+    public static generateGPTImage = async (prompt: string): Promise<string | undefined> => {
+        try {
+            const res = await gptImageCall(prompt);
+
+            if (res) {
+                const result = (await Networking.PostToServer('/uploadRemoteImage', { sources: res })) as Upload.FileInformation[];
+                const source = ClientUtils.prepend(result[0].accessPaths.agnostic.client);
+                return source;
+            }
+        } catch (e) {
+            console.log(e);
+        }
+    };
+
+    public static renderGPTImageCall = async (template: Template, col: Col, fieldNumber: number): Promise<boolean> => {
+        const generateAndLoadImage = async (id: number, prompt: string) => {
+            const field = template.getFieldByID(id);
+            const url = await this.generateGPTImage(prompt);
+            field?.setContent(url ?? '', ViewType.IMG);
+            field?.setTitle(col.title);
+        };
+
+        const fieldContent: string = template.compiledContent;
+
+        try {
+            const sysPrompt =
+                `#${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(fieldNumber, prompt);
+        } catch (e) {
+            console.log(e);
+        }
+        return true;
+    };
+
+    public static renderGPTTextCall = async (template: Template, col: Col, fieldNum: number | undefined): Promise<boolean> => {
+        const wordLimit = (size: TemplateFieldSize) => {
+            switch (size) {
+                case TemplateFieldSize.TINY:
+                    return 2;
+                case TemplateFieldSize.SMALL:
+                    return 5;
+                case TemplateFieldSize.MEDIUM:
+                    return 20;
+                case TemplateFieldSize.LARGE:
+                    return 50;
+                case TemplateFieldSize.HUGE:
+                    return 100;
+                default:
+                    return 10;
+            }
+        };
+
+        const textAssignment = `--- title: ${col.title}, prompt: ${col.desc}, word limit: ${wordLimit(col.sizes[0])} words, assigned field: ${fieldNum} ---`;
+
+        const fieldContent: string = template.compiledContent;
+
+        try {
+            const prompt = fieldContent + textAssignment;
+
+            const res = await gptAPICall(`${Math.random() * 100000}: ${prompt}`, GPTCallType.FILL);
+
+            if (res) {
+                const assignments: { [title: string]: { number: string; content: string } } = JSON.parse(res);
+                Object.entries(assignments).forEach(([, /* title */ info]) => {
+                    const field = template.getFieldByID(Number(info.number));
+
+                    field?.setContent(info.content ?? '', ViewType.TEXT);
+                    field?.setTitle(col.title);
+                });
+            }
+        } catch (err) {
+            console.log(err);
+        }
+
+        return true;
+    };
+
+    /**
+     * Populates a preset template framework with content from a datavizbox or any AI-generated content.
+     * @param template the preloaded template framework being filled in
+     * @param assignments a list of template field numbers (from top to bottom) and their assigned columns from the linked dataviz
+     * @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);
+
+        if (GPTTextCalls.length) {
+            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 ([id, col]) => {
+                return TemplateMenuAIUtils.renderGPTImageCall(template, col, Number(id));
+            });
+
+            await Promise.all(promises);
+        }
+
+        return template;
+    };
+}