1 files changed, 22 insertions, 22 deletions

diff --git a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
index cf9e8cfc8..3db401b14 100644
--- a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
+++ b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
@@ -4,6 +4,7 @@ import { BaseTool } from './BaseTool';
 import { Observation } from '../types/types';
 import { ParametersType, ToolInfo } from '../types/tool_types';
 import { DocumentOptions } from '../../../../documents/Documents';
+import { ClientUtils } from '../../../../../ClientUtils';
 
 const imageCreationToolParams = [
     {
@@ -36,31 +37,30 @@ export class ImageCreationTool extends BaseTool<ImageCreationToolParamsType> {
         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', {
+            const { url } = await Networking.PostToServer('/generateImage', {
                 image_prompt,
             });
-            const id = uuidv4();
+            if (url) {
+                const result = await Networking.PostToServer('/uploadRemoteImage', { sources: [url] });
+                const source = ClientUtils.prepend(result[0].accessPaths.agnostic.client);
 
-            this._addLinkedDoc('image', image_path, {}, id);
-            return [
-                {
-                    type: 'image_url',
-                    image_url: { url: `data:image/jpeg;base64,${base64_data}` },
-                },
-            ];
+                const id = uuidv4();
+
+                this._addLinkedDoc('image', source, {}, id);
+                return [
+                    {
+                        type: 'image_url',
+                        image_url: { url },
+                    },
+                ];
+            } else {
+                return [
+                    {
+                        type: 'text',
+                        text: `An error occurred while generating image.`,
+                    },
+                ];
+            }
         } catch (error) {
             console.log(error);
             return [