1 files changed, 16 insertions, 20 deletions

diff --git a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
index 177552c5c..dc6140871 100644
--- a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
+++ b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
@@ -1,10 +1,10 @@
-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';
+import { Upload } from '../../../../../server/SharedMediaTypes';
 
 const imageCreationToolParams = [
     {
@@ -25,8 +25,8 @@ const imageCreationToolInfo: ToolInfo<ImageCreationToolParamsType> = {
 };
 
 export class ImageCreationTool extends BaseTool<ImageCreationToolParamsType> {
-    private _createImage: (result: any, options: DocumentOptions) => void;
-    constructor(createImage: (result: any, options: DocumentOptions) => void) {
+    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;
     }
@@ -42,23 +42,19 @@ export class ImageCreationTool extends BaseTool<ImageCreationToolParamsType> {
             });
             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.`,
-                    },
-                ];
-            }
+            return url
+                ? [
+                      {
+                          type: 'image_url',
+                          image_url: { url },
+                      },
+                  ]
+                : [
+                      {
+                          type: 'text',
+                          text: `An error occurred while generating image.`,
+                      },
+                  ];
         } catch (error) {
             console.log(error);
             return [