mono/packages/ui/src/modules/ai/imageTools.ts
2026-03-21 20:18:25 +01:00

131 lines
5.3 KiB
TypeScript

/**
* Image Tools for Chat Playground
*
* Provides generate_image tool that the AI can invoke to create images
* and embed them as markdown in its responses.
* Uses the same manual tool definition pattern as searchTools.ts.
*/
import { z } from 'zod';
import { createImage as createImageRouter } from '@/lib/image-router';
import { supabase } from '@/integrations/supabase/client';
import type { RunnableToolFunctionWithParse } from 'openai/lib/RunnableFunction';
type LogFunction = (level: 'info' | 'warn' | 'error' | 'debug', message: string, data?: any) => void;
const defaultLog: LogFunction = (level, message, data) => console.log(`[IMAGE-TOOLS][${level}] ${message}`, data ?? '');
// ── Upload helper ────────────────────────────────────────────────────────
const uploadToTempBucket = async (
imageData: ArrayBuffer,
prompt: string,
userId: string,
addLog: LogFunction,
): Promise<string | null> => {
try {
const ts = Date.now();
const slug = prompt.slice(0, 20).replace(/[^a-zA-Z0-9]/g, '-');
const fileName = `${userId}/${ts}-${slug}.png`;
const uint8 = new Uint8Array(imageData);
const { error } = await supabase.storage
.from('temp-images')
.upload(fileName, uint8, { contentType: 'image/png', cacheControl: '3600' });
if (error) {
addLog('error', 'Upload failed', error);
return null;
}
const { data: { publicUrl } } = supabase.storage
.from('temp-images')
.getPublicUrl(fileName);
addLog('info', 'Image uploaded', { fileName, publicUrl });
return publicUrl;
} catch (err) {
addLog('error', 'Upload error', err);
return null;
}
};
// ── Zod schemas ──────────────────────────────────────────────────────────
const generateImageSchema = z.object({
prompt: z.string(),
altText: z.string().optional(),
caption: z.string().optional(),
});
type GenerateImageArgs = z.infer<typeof generateImageSchema>;
// ── Tool: generate_image ─────────────────────────────────────────────────
export const createImageTool = (
userId: string,
addLog: LogFunction = defaultLog,
modelString?: string,
): RunnableToolFunctionWithParse<GenerateImageArgs> => ({
type: 'function',
function: {
name: 'generate_image',
description:
'Generate an image from a text prompt. The image is uploaded to storage and ' +
'returned as a markdown image tag you can embed in your response. ' +
'Use descriptive, detailed prompts for best results.',
parameters: {
type: 'object',
properties: {
prompt: {
type: 'string',
description: 'Detailed prompt describing the image to generate.',
},
altText: {
type: 'string',
description: 'Alt text for the image (defaults to prompt).',
},
caption: {
type: 'string',
description: 'Optional caption to display below the image.',
},
},
required: ['prompt'],
} as any,
parse(input: string): GenerateImageArgs {
return generateImageSchema.parse(JSON.parse(input));
},
function: async (args: GenerateImageArgs) => {
try {
addLog('info', '[IMAGE-TOOLS] generate_image called', { prompt: args.prompt.substring(0, 100), model: modelString });
const result = await createImageRouter(args.prompt, modelString);
if (!result) {
return JSON.stringify({
success: false,
error: 'Image generation failed',
markdown: `_Image generation failed for: "${args.prompt}"_`,
});
}
const publicUrl = await uploadToTempBucket(result.imageData, args.prompt, userId, addLog);
if (!publicUrl) {
return JSON.stringify({
success: false,
error: 'Upload failed',
markdown: `_Image upload failed for: "${args.prompt}"_`,
});
}
const alt = args.altText || args.prompt;
const cap = args.caption ? `\n*${args.caption}*` : '';
const markdown = `![${alt}](${publicUrl})${cap}`;
addLog('info', 'Image embedded', { url: publicUrl });
return JSON.stringify({ success: true, markdown, imageUrl: publicUrl });
} catch (err: any) {
addLog('error', 'generate_image failed', err);
return JSON.stringify({ success: false, error: err.message });
}
},
},
});