aboutsummaryrefslogtreecommitdiff
path: root/src/client/views/nodes/chatbot/tools/ImageCreationTool.ts
blob: c5b1e028b146dd459debc58536cecafe02d09281 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
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.`,
                },
            ];
        }
    }
}