2 files changed, 82 insertions, 4 deletions

diff --git a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
index a4871f7fd..4c059177b 100644
--- a/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
+++ b/src/client/views/nodes/chatbot/tools/CreateAnyDocTool.ts
@@ -7,8 +7,8 @@ import { DocumentOptions, Docs } from '../../../../documents/Documents';
 /**
  * List of supported document types that can be created via text LLM.
  */
-type supportedDocumentTypesType = 'text' | 'html' | 'equation' | 'functionPlot' | 'dataviz' | 'noteTaking' | 'rtf' | 'message';
-const supportedDocumentTypes: supportedDocumentTypesType[] = ['text', 'html', 'equation', 'functionPlot', 'dataviz', 'noteTaking', 'rtf', 'message'];
+type supportedDocumentTypesType = 'text' | 'html' | 'equation' | 'function_plot' | 'dataviz' | 'note_taking' | 'rtf' | 'message' | 'mermaid_diagram';
+const supportedDocumentTypes: supportedDocumentTypesType[] = ['text', 'html', 'equation', 'function_plot', 'dataviz', 'note_taking', 'rtf', 'message', 'mermaid_diagram'];
 
 /**
  * Description of document options and data field for each type.
@@ -26,7 +26,7 @@ const documentTypesInfo = {
         options: ['title', 'backgroundColor', 'fontColor', 'layout'],
         dataDescription: 'The equation content as a string.',
     },
-    functionPlot: {
+    function_plot: {
         options: ['title', 'backgroundColor', 'layout', 'function_definition'],
         dataDescription: 'The function definition(s) for plotting. Provide as a string or array of function definitions.',
     },
@@ -34,7 +34,7 @@ const documentTypesInfo = {
         options: ['title', 'backgroundColor', 'layout', 'chartType'],
         dataDescription: 'A string of comma-separated values representing the CSV data.',
     },
-    noteTaking: {
+    note_taking: {
         options: ['title', 'backgroundColor', 'layout'],
         dataDescription: 'The initial content or structure for note-taking.',
     },
@@ -46,6 +46,10 @@ const documentTypesInfo = {
         options: ['title', 'backgroundColor', 'layout'],
         dataDescription: 'The message content of the document.',
     },
+    mermaid_diagram: {
+        options: ['title', 'backgroundColor', 'layout'],
+        dataDescription: 'The Mermaid diagram content.',
+    },
 };
 
 const createAnyDocumentToolParams = [
diff --git a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
new file mode 100644
index 000000000..cf9e8cfc8
--- /dev/null
+++ b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
@@ -0,0 +1,74 @@
+import { v4 as uuidv4 } from 'uuid';
+import { Networking } from '../../../../Network';
+import { BaseTool } from './BaseTool';
+import { Observation } from '../types/types';
+import { ParametersType, ToolInfo } from '../types/tool_types';
+import { DocumentOptions } from '../../../../documents/Documents';
+
+const imageCreationToolParams = [
+    {
+        name: 'image_prompt',
+        type: 'string',
+        description: 'The prompt for the image to be created. This should be a string that describes the image to be created in extreme detail for an AI image generator.',
+        required: true,
+    },
+] as const;
+
+type ImageCreationToolParamsType = typeof imageCreationToolParams;
+
+const imageCreationToolInfo: ToolInfo<ImageCreationToolParamsType> = {
+    name: 'imageCreationTool',
+    citationRules: 'No citation needed. Cannot cite image generation for a response.',
+    parameterRules: imageCreationToolParams,
+    description: 'Create an image of any style, content, or design, based on a prompt. The prompt should be a detailed description of the image to be created.',
+};
+
+export class ImageCreationTool extends BaseTool<ImageCreationToolParamsType> {
+    private _addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void;
+    constructor(addLinkedDoc: (doc_type: string, data: string | undefined, options: DocumentOptions, id: string) => void) {
+        super(imageCreationToolInfo);
+        this._addLinkedDoc = addLinkedDoc;
+    }
+
+    async execute(args: ParametersType<ImageCreationToolParamsType>): Promise<Observation[]> {
+        const image_prompt = args.image_prompt;
+
+        console.log(`Generating image for prompt: ${image_prompt}`);
+        // Create an array of promises, each one handling a search for a query
+        try {
+            try {
+                const { image_url } = await Networking.PostToServer('/generateImage', {
+                    image_prompt,
+                });
+                if (res) {
+                    const result = await Networking.PostToServer('/uploadRemoteImage', { sources: res });
+                    const source = ClientUtils.prepend(result[0].accessPaths.agnostic.client);
+                    return source;
+                }
+            } catch (e) {
+                console.log(e);
+            }
+
+            const { base64_data, image_path } = await Networking.PostToServer('/generateImage', {
+                image_prompt,
+            });
+            const id = uuidv4();
+
+            this._addLinkedDoc('image', image_path, {}, id);
+            return [
+                {
+                    type: 'image_url',
+                    image_url: { url: `data:image/jpeg;base64,${base64_data}` },
+                },
+            ];
+        } catch (error) {
+            console.log(error);
+            return [
+                {
+                    type: 'text',
+                    text: `An error occurred while generating image.`,
+                },
+            ];
+        }
+    }
+}