add image upload too to AI Chat

Author: iNerdStack

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/message-content/tool-call-display.tsx

@@ -8,6 +8,8 @@ import {
     TERMINAL_COMMAND_TOOL_NAME, TERMINAL_COMMAND_TOOL_PARAMETERS, TODO_WRITE_TOOL_NAME,
     type TODO_WRITE_TOOL_PARAMETERS,
     TYPECHECK_TOOL_NAME,
+    UPLOAD_IMAGE_TOOL_NAME,
+    type UPLOAD_IMAGE_TOOL_PARAMETERS,
     WEB_SEARCH_TOOL_NAME,
     type WEB_SEARCH_TOOL_PARAMETERS,
     WRITE_FILE_TOOL_NAME,
@@ -89,6 +91,39 @@ export const ToolCallDisplay = ({
         }
     }
 
+    if (toolName === UPLOAD_IMAGE_TOOL_NAME) {
+        const args = toolPart.input as z.infer<typeof UPLOAD_IMAGE_TOOL_PARAMETERS> | null;
+        const result = toolPart.output as string | null;
+        return (
+            <div className="flex flex-col gap-2 p-3 border rounded-lg bg-background-secondary">
+                <div className="flex items-center gap-2">
+                    <Icons.Image className="w-4 h-4" />
+                    <span className="text-sm font-medium">Upload Image</span>
+                </div>
+                {args?.image_reference && (
+                    <div className="text-xs text-foreground-secondary">
+                        Image: {args.image_reference}
+                    </div>
+                )}
+                {args?.destination_path && (
+                    <div className="text-xs text-foreground-secondary">
+                        Destination: {args.destination_path}
+                    </div>
+                )}
+                {args?.filename && (
+                    <div className="text-xs text-foreground-secondary">
+                        Filename: {args.filename}
+                    </div>
+                )}
+                {result && (
+                    <div className="text-xs text-green-600 mt-1">
+                        {result}
+                    </div>
+                )}
+            </div>
+        );
+    }
+
     if (toolName === WRITE_FILE_TOOL_NAME) {
         const args = toolPart.input as z.infer<typeof WRITE_FILE_TOOL_PARAMETERS> | null;
         const filePath = args?.file_path;

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-messages/message-content/tool-call-simple.tsx

@@ -28,6 +28,8 @@ import {
     TODO_WRITE_TOOL_NAME,
     type TODO_WRITE_TOOL_PARAMETERS,
     TYPECHECK_TOOL_NAME,
+    UPLOAD_IMAGE_TOOL_NAME,
+    type UPLOAD_IMAGE_TOOL_PARAMETERS,
     WEB_SEARCH_TOOL_NAME,
     type WEB_SEARCH_TOOL_PARAMETERS,
     WRITE_FILE_TOOL_NAME,
@@ -60,6 +62,7 @@ const TOOL_ICONS: Record<string, any> = {
     [TYPECHECK_TOOL_NAME]: Icons.MagnifyingGlass,
     [LIST_BRANCHES_TOOL_NAME]: Icons.Branch,
     [GLOB_TOOL_NAME]: Icons.MagnifyingGlass,
+    [UPLOAD_IMAGE_TOOL_NAME]: Icons.Image,
 } as const;
 
 export function ToolCallSimple({
@@ -217,6 +220,14 @@ const getLabel = (toolName: string, toolPart: ToolUIPart) => {
                 return 'Reading Onlook instructions';
             case TYPECHECK_TOOL_NAME:
                 return 'Checking types';
+            case UPLOAD_IMAGE_TOOL_NAME: {
+                const params = toolPart.input as z.infer<typeof UPLOAD_IMAGE_TOOL_PARAMETERS>;
+                if (params?.image_reference) {
+                    return 'Uploading image ' + truncateString(params.image_reference);
+                } else {
+                    return 'Uploading image';
+                }
+            }
             default:
                 return toolName?.replace(/[-_]/g, ' ').replace(/\b\w/g, c => c.toUpperCase());
         }

apps/web/client/src/components/tools/handlers/edit.ts

@@ -4,8 +4,11 @@ import {
     FUZZY_EDIT_FILE_TOOL_PARAMETERS,
     SEARCH_REPLACE_EDIT_FILE_TOOL_PARAMETERS,
     SEARCH_REPLACE_MULTI_EDIT_FILE_TOOL_PARAMETERS,
+    UPLOAD_IMAGE_TOOL_PARAMETERS,
+    VIEW_IMAGE_TOOL_PARAMETERS,
     WRITE_FILE_TOOL_PARAMETERS
 } from '@onlook/ai';
+import { MessageContextType } from '@onlook/models';
 import { z } from 'zod';
 
 export async function handleSearchReplaceEditFileTool(args: z.infer<typeof SEARCH_REPLACE_EDIT_FILE_TOOL_PARAMETERS>, editorEngine: EditorEngine): Promise<string> {
@@ -167,3 +170,168 @@ export async function handleFuzzyEditFileTool(
     }
     return 'File edited!';
 }
+
+export async function handleViewImageTool(args: z.infer<typeof VIEW_IMAGE_TOOL_PARAMETERS>, editorEngine: EditorEngine): Promise<{ image: { mimeType: string; data: string }; message: string }> {
+    try {
+        // Find the image in the chat context by reference
+        const context = editorEngine.chat.context.context;
+        const imageContext = context.find((ctx) => {
+            if (ctx.type !== MessageContextType.IMAGE) {
+                return false;
+            }
+            // Try to match by display name, index number, or description
+            const ref = args.image_reference.toLowerCase();
+            return ctx.displayName.toLowerCase().includes(ref) ||
+                   ref.includes(ctx.displayName.toLowerCase()) ||
+                   ref.match(/^\d+$/) && context.filter(c => c.type === MessageContextType.IMAGE)[parseInt(ref) - 1] === ctx;
+        });
+
+        if (!imageContext || imageContext.type !== MessageContextType.IMAGE) {
+            // Try to find by index number
+            const imageContexts = context.filter(ctx => ctx.type === MessageContextType.IMAGE);
+            const indexMatch = args.image_reference.match(/^\d+$/);
+            if (indexMatch) {
+                const index = parseInt(indexMatch[0]) - 1;
+                if (index >= 0 && index < imageContexts.length) {
+                    const foundImage = imageContexts[index];
+                    if (foundImage && foundImage.type === MessageContextType.IMAGE) {
+                        return {
+                            image: {
+                                mimeType: foundImage.mimeType,
+                                data: foundImage.content,
+                            },
+                            message: `Retrieved image "${foundImage.displayName}" for analysis.`,
+                        };
+                    }
+                }
+            }
+
+            throw new Error(`No image found matching reference: ${args.image_reference}`);
+        }
+
+        return {
+            image: {
+                mimeType: imageContext.mimeType,
+                data: imageContext.content,
+            },
+            message: `Retrieved image "${imageContext.displayName}" for analysis.`,
+        };
+    } catch (error) {
+        throw new Error(`Cannot view image: ${error}`);
+    }
+}
+
+export async function handleUploadImageTool(args: z.infer<typeof UPLOAD_IMAGE_TOOL_PARAMETERS>, editorEngine: EditorEngine): Promise<string> {
+    try {
+        const sandbox = editorEngine.branches.getSandboxById(args.branchId);
+        if (!sandbox) {
+            throw new Error(`Sandbox not found for branch ID: ${args.branchId}`);
+        }
+
+        // Find the image in the chat context by reference
+        const context = editorEngine.chat.context.context;
+        const imageContext = context.find((ctx) => {
+            if (ctx.type !== MessageContextType.IMAGE) {
+                return false;
+            }
+            // Try to match by display name or description
+            return ctx.displayName.toLowerCase().includes(args.image_reference.toLowerCase()) ||
+                   args.image_reference.toLowerCase().includes(ctx.displayName.toLowerCase());
+        });
+
+        if (!imageContext || imageContext.type !== MessageContextType.IMAGE) {
+            // Try to find the most recent image if no specific match
+            const recentImages = context.filter(ctx => ctx.type === MessageContextType.IMAGE);
+            if (recentImages.length === 0) {
+                throw new Error(`No image found matching reference: ${args.image_reference}`);
+            }
+
+            // Use the most recent image if no specific match
+            const mostRecentImage = recentImages[recentImages.length - 1];
+            if (!mostRecentImage || mostRecentImage.type !== MessageContextType.IMAGE) {
+                throw new Error(`No image found matching reference: ${args.image_reference}`);
+            }
+
+            console.warn(`No exact match for "${args.image_reference}", using most recent image: ${mostRecentImage.displayName}`);
+
+            // Extract MIME type and file extension
+            const mimeType = mostRecentImage.mimeType;
+            const extension = getExtensionFromMimeType(mimeType);
+
+            // Generate filename
+            const filename = args.filename ? `${args.filename}.${extension}` : `${generateUUID()}.${extension}`;
+
+            // Determine destination path
+            const destinationPath = args.destination_path || 'public/assets/images';
+            const fullPath = `${destinationPath}/${filename}`;
+
+            // Convert base64 to binary data
+            const base64Data = mostRecentImage.content.replace(/^data:image\/[a-z]+;base64,/, '');
+            const binaryData = base64ToUint8Array(base64Data);
+
+            // Upload to sandbox
+            await sandbox.writeBinaryFile(fullPath, binaryData);
+
+            // Refresh image scanning to update the UI
+            await editorEngine.image.scanImages();
+
+            return `Image "${mostRecentImage.displayName}" uploaded successfully to ${fullPath}`;
+        }
+
+        // Extract MIME type and file extension
+        const mimeType = imageContext.mimeType;
+        const extension = getExtensionFromMimeType(mimeType);
+
+        // Generate filename
+        const filename = args.filename ? `${args.filename}.${extension}` : `${generateUUID()}.${extension}`;
+
+        // Determine destination path
+        const destinationPath = args.destination_path || 'public/assets/images';
+        const fullPath = `${destinationPath}/${filename}`;
+
+        // Convert base64 to binary data
+        const base64Data = imageContext.content.replace(/^data:image\/[a-z]+;base64,/, '');
+        const binaryData = base64ToUint8Array(base64Data);
+
+        // Upload to sandbox
+        await sandbox.writeBinaryFile(fullPath, binaryData);
+
+        // Refresh image scanning to update the UI
+        await editorEngine.image.scanImages();
+
+        return `Image "${imageContext.displayName}" uploaded successfully to ${fullPath}`;
+    } catch (error) {
+        throw new Error(`Cannot upload image: ${error}`);
+    }
+}
+
+function getExtensionFromMimeType(mimeType: string): string {
+    const mimeToExt: Record<string, string> = {
+        'image/jpeg': 'jpg',
+        'image/jpg': 'jpg',
+        'image/png': 'png',
+        'image/gif': 'gif',
+        'image/webp': 'webp',
+        'image/svg+xml': 'svg',
+        'image/bmp': 'bmp',
+        'image/tiff': 'tiff',
+    };
+    return mimeToExt[mimeType.toLowerCase()] || 'png';
+}
+
+function base64ToUint8Array(base64: string): Uint8Array {
+    const binaryString = atob(base64);
+    const bytes = new Uint8Array(binaryString.length);
+    for (let i = 0; i < binaryString.length; i++) {
+        bytes[i] = binaryString.charCodeAt(i);
+    }
+    return bytes;
+}
+
+function generateUUID(): string {
+    return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, function(c) {
+        const r = Math.random() * 16 | 0;
+        const v = c === 'x' ? r : (r & 0x3 | 0x8);
+        return v.toString(16);
+    });
+}

apps/web/client/src/components/tools/tools.ts

@@ -32,6 +32,10 @@ import {
     TERMINAL_COMMAND_TOOL_PARAMETERS,
     TYPECHECK_TOOL_NAME,
     TYPECHECK_TOOL_PARAMETERS,
+    UPLOAD_IMAGE_TOOL_NAME,
+    UPLOAD_IMAGE_TOOL_PARAMETERS,
+    VIEW_IMAGE_TOOL_NAME,
+    VIEW_IMAGE_TOOL_PARAMETERS,
     WEB_SEARCH_TOOL_NAME,
     WEB_SEARCH_TOOL_PARAMETERS,
     WRITE_FILE_TOOL_NAME,
@@ -56,6 +60,8 @@ import {
     handleSearchReplaceMultiEditFileTool,
     handleTerminalCommandTool,
     handleTypecheckTool,
+    handleUploadImageTool,
+    handleViewImageTool,
     handleWebSearchTool,
     handleWriteFileTool
 } from './handlers';
@@ -181,6 +187,18 @@ const TOOL_HANDLERS: ClientToolMap = {
         handler: async (args: z.infer<typeof CHECK_ERRORS_TOOL_PARAMETERS>, editorEngine: EditorEngine) =>
             handleCheckErrors(args, editorEngine),
     },
+    [VIEW_IMAGE_TOOL_NAME]: {
+        name: VIEW_IMAGE_TOOL_NAME,
+        inputSchema: VIEW_IMAGE_TOOL_PARAMETERS,
+        handler: async (args: z.infer<typeof VIEW_IMAGE_TOOL_PARAMETERS>, editorEngine: EditorEngine) =>
+            handleViewImageTool(args, editorEngine),
+    },
+    [UPLOAD_IMAGE_TOOL_NAME]: {
+        name: UPLOAD_IMAGE_TOOL_NAME,
+        inputSchema: UPLOAD_IMAGE_TOOL_PARAMETERS,
+        handler: async (args: z.infer<typeof UPLOAD_IMAGE_TOOL_PARAMETERS>, editorEngine: EditorEngine) =>
+            handleUploadImageTool(args, editorEngine),
+    },
 };
 
 export async function handleToolCall(toolCall: ToolCall<string, unknown>, editorEngine: EditorEngine, addToolResult: (toolResult: { tool: string, toolCallId: string, output: any }) => Promise<void>) {

packages/ai/src/prompt/provider.ts

@@ -6,7 +6,6 @@ import type {
     MessageContext,
     ProjectMessageContext,
 } from '@onlook/models';
-import type { FileUIPart } from 'ai';
 import { ASK_MODE_SYSTEM_PROMPT } from './ask';
 import { CONTEXT_PROMPTS } from './context';
 import { CREATE_NEW_PAGE_SYSTEM_PROMPT } from './create';
@@ -111,17 +110,18 @@ export function getHydratedUserMessage(
         .join('\n');
     prompt += wrapXml('instruction', textContent);
 
-    userParts.push({ type: 'text', text: prompt });
-
+    // Add image references to prompt (but don't send image data yet)
+    // AI will decide whether to view or upload them using tools
     if (images.length > 0) {
-        const attachments: FileUIPart[] = images.map((i) => ({
-            type: 'file',
-            mediaType: i.mimeType,
-            url: i.content,
-        }));
-        userParts = userParts.concat(attachments);
+        const imageList = images
+            .map((img, idx) => `${idx + 1}. "${img.displayName}" (${img.mimeType})`)
+            .join('\n');
+        const imagesPrompt = `The user has attached ${images.length} image(s) to this message:\n${imageList}\n\nYou can:\n- Use the "view_image" tool to analyze the image content\n- Use the "upload_image" tool to save it to the project\n\nDetermine the appropriate action based on the user's request.`;
+        prompt += wrapXml('available-images', imagesPrompt);
     }
 
+    userParts.push({ type: 'text', text: prompt });
+
     return {
         id,
         role: 'user',

packages/ai/src/tools/tools/edit.ts

@@ -73,3 +73,40 @@ export const fuzzyEditFileTool = tool({
         'Edit the contents of a file with fuzzy matching instead of search and replace. This should be used as a fallback when the search and replace tool fails. It calls another agent to do the actual editing.',
     inputSchema: FUZZY_EDIT_FILE_TOOL_PARAMETERS,
 });
+
+export const VIEW_IMAGE_TOOL_NAME = 'view_image';
+export const VIEW_IMAGE_TOOL_PARAMETERS = z.object({
+    image_reference: z
+        .string()
+        .describe(
+            'Reference to an image in the chat context (use the display name or index number)',
+        ),
+});
+export const viewImageTool = tool({
+    description:
+        "Retrieves and views an image from the chat context for analysis. Use this tool when the user asks you to analyze, describe, or work with an image they've attached. The image data will be returned so you can see and analyze its contents. This does NOT save the image to the project.",
+    inputSchema: VIEW_IMAGE_TOOL_PARAMETERS,
+});
+
+export const UPLOAD_IMAGE_TOOL_NAME = 'upload_image';
+export const UPLOAD_IMAGE_TOOL_PARAMETERS = z.object({
+    image_reference: z
+        .string()
+        .describe(
+            'Reference to an image in the chat context (use the display name or index number)',
+        ),
+    destination_path: z
+        .string()
+        .optional()
+        .describe('Destination path within the project (default: "public/assets/images")'),
+    filename: z
+        .string()
+        .optional()
+        .describe('Custom filename (without extension). If not provided, a UUID will be generated'),
+    branchId: BRANCH_ID_SCHEMA,
+});
+export const uploadImageTool = tool({
+    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 the project's public directory and can be referenced in code. After uploading, you can use the file path in your code changes.",
+    inputSchema: UPLOAD_IMAGE_TOOL_PARAMETERS,
+});

packages/ai/src/tools/toolset.ts

@@ -33,6 +33,10 @@ import {
     terminalCommandTool,
     TYPECHECK_TOOL_NAME,
     typecheckTool,
+    UPLOAD_IMAGE_TOOL_NAME,
+    uploadImageTool,
+    VIEW_IMAGE_TOOL_NAME,
+    viewImageTool,
     WEB_SEARCH_TOOL_NAME,
     webSearchTool,
     WRITE_FILE_TOOL_NAME,
@@ -62,6 +66,8 @@ export const BUILD_TOOL_SET: ToolSet = {
     [SANDBOX_TOOL_NAME]: sandboxTool,
     [TERMINAL_COMMAND_TOOL_NAME]: terminalCommandTool,
     [TYPECHECK_TOOL_NAME]: typecheckTool,
+    [VIEW_IMAGE_TOOL_NAME]: viewImageTool,
+    [UPLOAD_IMAGE_TOOL_NAME]: uploadImageTool,
 };
 
 export type ChatTools = InferUITools<typeof BUILD_TOOL_SET>;