image generation works better

1 files changed, 8 insertions, 8 deletions

diff --git a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
index 3db401b14..ba1aa987a 100644
--- a/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
+++ b/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
@@ -5,6 +5,8 @@ import { Observation } from '../types/types';
 import { ParametersType, ToolInfo } from '../types/tool_types';
 import { DocumentOptions } from '../../../../documents/Documents';
 import { ClientUtils } from '../../../../../ClientUtils';
+import { DashUploadUtils } from '../../../../../server/DashUploadUtils';
+import { RTFCast, StrCast } from '../../../../../fields/Types';
 
 const imageCreationToolParams = [
     {
@@ -25,10 +27,10 @@ const imageCreationToolInfo: ToolInfo<ImageCreationToolParamsType> = {
 };
 
 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) {
+    private _createImage: (result: any, options: DocumentOptions) => void;
+    constructor(createImage: (result: any, options: DocumentOptions) => void) {
         super(imageCreationToolInfo);
-        this._addLinkedDoc = addLinkedDoc;
+        this._createImage = createImage;
     }
 
     async execute(args: ParametersType<ImageCreationToolParamsType>): Promise<Observation[]> {
@@ -37,16 +39,14 @@ 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 {
-            const { url } = await Networking.PostToServer('/generateImage', {
+            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 result = await Networking.PostToServer('/uploadRemoteImage', { sources: [url] });
-                const source = ClientUtils.prepend(result[0].accessPaths.agnostic.client);
-
                 const id = uuidv4();
 
-                this._addLinkedDoc('image', source, {}, id);
                 return [
                     {
                         type: 'image_url',