1 files changed, 69 insertions, 0 deletions
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..37907fd4f
--- /dev/null
+++ b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
@@ -0,0 +1,69 @@
+import { RTFCast } from '../../../../../fields/Types';
+import { DocumentOptions } from '../../../../documents/Documents';
+import { Networking } from '../../../../Network';
+import { ParametersType, ToolInfo } from '../types/tool_types';
+import { Observation } from '../types/types';
+import { BaseTool } from './BaseTool';
+import { Upload } from '../../../../../server/SharedMediaTypes';
+import { List } from '../../../../../fields/List';
+
+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 _createImage: (result: Upload.FileInformation & Upload.InspectionResults, options: DocumentOptions) => void;
+    constructor(createImage: (result: Upload.FileInformation & Upload.InspectionResults, options: DocumentOptions) => void) {
+        super(imageCreationToolInfo);
+        this._createImage = createImage;
+    }
+
+    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 {
+            const { result, url } = (await Networking.PostToServer('/generateImage', {
+                image_prompt,
+            })) as { result: Upload.FileInformation & Upload.InspectionResults; url: string };
+            console.log('Image generation result:', result);
+            this._createImage(result, { text: RTFCast(image_prompt), ai: 'dall-e-3', tags: new List<string>(['@ai']) });
+            return url
+                ? [
+                      {
+                          type: 'image_url',
+                          image_url: { url },
+                      },
+                  ]
+                : [
+                      {
+                          type: 'text',
+                          text: `An error occurred while generating image.`,
+                      },
+                  ];
+        } catch (error) {
+            console.log(error);
+            return [
+                {
+                    type: 'text',
+                    text: `An error occurred while generating image.`,
+                },
+            ];
+        }
+    }
+}