|
| 1 | +import { Icons } from '@onlook/ui/icons'; |
| 2 | +import type { EditorEngine } from '@onlook/web-client/src/components/store/editor/engine'; |
| 3 | +import type { SandboxManager } from '@onlook/web-client/src/components/store/editor/sandbox'; |
| 4 | +import { MessageContextType, type ImageMessageContext } from '@onlook/models'; |
| 5 | +import mime from 'mime-lite'; |
| 6 | +import { v4 as uuidv4 } from 'uuid'; |
| 7 | +import { z } from 'zod'; |
| 8 | +import { ClientTool } from '../models/client'; |
| 9 | +import { BRANCH_ID_SCHEMA } from '../shared/type'; |
| 10 | + |
| 11 | +export class UploadImageTool extends ClientTool { |
| 12 | + static readonly toolName = 'upload_image'; |
| 13 | + static readonly description = "Uploads an image from the chat context to the project's file system. Use this tool when the user asks you to save, add, or upload an image to their project. The image will be stored in public/images/ directory by default and can be referenced in code. After uploading, you can use the file path in your code changes."; |
| 14 | + static readonly parameters = z.object({ |
| 15 | + image_id: z |
| 16 | + .string() |
| 17 | + .describe( |
| 18 | + 'The unique ID of the image from the available images list', |
| 19 | + ), |
| 20 | + destination_path: z |
| 21 | + .string() |
| 22 | + .optional() |
| 23 | + .describe('Destination path within the project. Defaults to "public/images" if not specified.'), |
| 24 | + filename: z |
| 25 | + .string() |
| 26 | + .optional() |
| 27 | + .describe('Custom filename (without extension). If not provided, a UUID will be generated'), |
| 28 | + branchId: BRANCH_ID_SCHEMA, |
| 29 | + }); |
| 30 | + static readonly icon = Icons.Image; |
| 31 | + |
| 32 | + async handle( |
| 33 | + args: z.infer<typeof UploadImageTool.parameters>, |
| 34 | + editorEngine: EditorEngine |
| 35 | + ): Promise<string> { |
| 36 | + try { |
| 37 | + const sandbox = editorEngine.branches.getSandboxById(args.branchId); |
| 38 | + if (!sandbox) { |
| 39 | + throw new Error(`Sandbox not found for branch ID: ${args.branchId}`); |
| 40 | + } |
| 41 | + |
| 42 | + // Get the current conversation ID |
| 43 | + const conversationId = editorEngine.chat.getCurrentConversationId(); |
| 44 | + if (!conversationId) { |
| 45 | + throw new Error('No active conversation found'); |
| 46 | + } |
| 47 | + |
| 48 | + // Fetch all messages from the conversation |
| 49 | + const messages = await editorEngine.api.getConversationMessages(conversationId); |
| 50 | + |
| 51 | + // Find the image in message contexts by ID, searching from most recent to oldest |
| 52 | + let imageContext: ImageMessageContext | null = null; |
| 53 | + for (let i = messages.length - 1; i >= 0; i--) { |
| 54 | + const message = messages[i]; |
| 55 | + if (!message || !message.metadata?.context) continue; |
| 56 | + |
| 57 | + const contexts = message.metadata.context; |
| 58 | + const imageContexts = contexts.filter( |
| 59 | + (ctx) => ctx.type === MessageContextType.IMAGE |
| 60 | + ); |
| 61 | + |
| 62 | + // Find image by ID |
| 63 | + const match = imageContexts.find((ctx) => ctx.id === args.image_id); |
| 64 | + |
| 65 | + if (match) { |
| 66 | + imageContext = match; |
| 67 | + break; |
| 68 | + } |
| 69 | + } |
| 70 | + |
| 71 | + if (!imageContext) { |
| 72 | + throw new Error(`No image found with ID: ${args.image_id}`); |
| 73 | + } |
| 74 | + |
| 75 | + // Upload the image to the sandbox |
| 76 | + const fullPath = await this.uploadImageToSandbox(imageContext, args, sandbox); |
| 77 | + |
| 78 | + return `Image "${imageContext.displayName}" uploaded successfully to ${fullPath}`; |
| 79 | + } catch (error) { |
| 80 | + throw new Error(`Cannot upload image: ${error}`); |
| 81 | + } |
| 82 | + } |
| 83 | + |
| 84 | + static getLabel(input?: z.infer<typeof UploadImageTool.parameters>): string { |
| 85 | + if (input?.filename) { |
| 86 | + return 'Uploading image ' + input.filename.substring(0, 20); |
| 87 | + } |
| 88 | + return 'Uploading image'; |
| 89 | + } |
| 90 | + |
| 91 | + private async uploadImageToSandbox( |
| 92 | + imageContext: ImageMessageContext, |
| 93 | + args: z.infer<typeof UploadImageTool.parameters>, |
| 94 | + sandbox: SandboxManager |
| 95 | + ): Promise<string> { |
| 96 | + const mimeType = imageContext.mimeType; |
| 97 | + const extension = mime.getExtension(mimeType) || 'png'; |
| 98 | + const filename = args.filename ? `${args.filename}.${extension}` : `${uuidv4()}.${extension}`; |
| 99 | + const destinationPath = args.destination_path?.trim() || 'public/images'; |
| 100 | + const fullPath = `${destinationPath}/${filename}`; |
| 101 | + |
| 102 | + // Extract base64 data from the content (remove data URL prefix if present) |
| 103 | + const base64Data = imageContext.content.replace(/^data:image\/[a-zA-Z0-9+.-]+;base64,/, ''); |
| 104 | + const binaryData = this.base64ToUint8Array(base64Data); |
| 105 | + |
| 106 | + await sandbox.writeFile(fullPath, binaryData); |
| 107 | + return fullPath; |
| 108 | + } |
| 109 | + |
| 110 | + private base64ToUint8Array(base64: string): Uint8Array { |
| 111 | + const binaryString = atob(base64); |
| 112 | + const bytes = new Uint8Array(binaryString.length); |
| 113 | + for (let i = 0; i < binaryString.length; i++) { |
| 114 | + bytes[i] = binaryString.charCodeAt(i); |
| 115 | + } |
| 116 | + return bytes; |
| 117 | + } |
| 118 | +} |
0 commit comments