1 files changed, 72 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..177552c5c
--- /dev/null
+++ b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
@@ -0,0 +1,72 @@
+import { v4 as uuidv4 } from 'uuid';
+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';
+
+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: any, options: DocumentOptions) => void;
+    constructor(createImage: (result: any, 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,
+            });
+            console.log('Image generation result:', result);
+            this._createImage(result, { text: RTFCast(image_prompt) });
+            if (url) {
+                const id = uuidv4();
+
+                return [
+                    {
+                        type: 'image_url',
+                        image_url: { url },
+                    },
+                ];
+            } else {
+                return [
+                    {
+                        type: 'text',
+                        text: `An error occurred while generating image.`,
+                    },
+                ];
+            }
+        } catch (error) {
+            console.log(error);
+            return [
+                {
+                    type: 'text',
+                    text: `An error occurred while generating image.`,
+                },
+            ];
+        }
+    }
+}