Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions apps/web/client/src/app/_components/hero/create.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import { cn } from '@onlook/ui/utils';
import { compressImageInBrowser } from '@onlook/utility';
import localforage from 'localforage';
import { observer } from 'mobx-react-lite';
import { v4 as uuidv4 } from 'uuid';
import { AnimatePresence } from 'motion/react';
import { useRouter } from 'next/navigation';
import { useEffect, useRef, useState } from 'react';
Expand Down Expand Up @@ -200,6 +201,7 @@ export const Create = observer(({
content: base64,
displayName: file.name,
mimeType: file.type,
id: uuidv4(),
};
} catch (error) {
console.error('Error reading file:', error);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from '@onlook/ui/tooltip';
import { cn } from '@onlook/ui/utils';
import { compressImageInBrowser } from '@onlook/utility';
import { observer } from 'mobx-react-lite';
import { v4 as uuidv4 } from 'uuid';
import { useTranslations } from 'next-intl';
import { useEffect, useMemo, useRef, useState } from 'react';
import { validateImageLimit } from '../context-pills/helpers';
Expand Down Expand Up @@ -213,6 +214,7 @@ export const ChatInput = observer(({
content: base64URL,
mimeType: file.type,
displayName: displayName ?? file.name,
id: uuidv4(),
};
editorEngine.chat.context.addContexts([contextImage]);
};
Expand Down Expand Up @@ -265,6 +267,7 @@ export const ChatInput = observer(({
content: screenshotData,
mimeType: mimeType,
displayName: 'Screenshot',
id: uuidv4(),
};
editorEngine.chat.context.addContexts([contextImage]);
toast.success('Screenshot added to chat');
Expand Down
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import { FuzzyEditFileTool, SearchReplaceEditTool, SearchReplaceMultiEditFileTool, TerminalCommandTool, TypecheckTool, WebSearchTool, WriteFileTool } from '@onlook/ai';
import { FuzzyEditFileTool, SearchReplaceEditTool, SearchReplaceMultiEditFileTool, TerminalCommandTool, TypecheckTool, UploadImageTool, ViewImageTool, WebSearchTool, WriteFileTool } from '@onlook/ai';
import type { WebSearchResult } from '@onlook/models';
import { Icons } from '@onlook/ui/icons';
import type { ToolUIPart } from 'ai';
import stripAnsi from 'strip-ansi';
import { type z } from 'zod';
Expand Down Expand Up @@ -69,6 +70,62 @@ export const ToolCallDisplay = ({
}
}

if (toolName === ViewImageTool.toolName) {
const args = toolPart.input as z.infer<typeof ViewImageTool.parameters> | null;
const result = toolPart.output as { message: 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">View Image</span>
</div>
{args?.image_id && (
<div className="text-xs text-foreground-secondary">
Image ID: {args.image_id}
</div>
)}
{result?.message && (
<div className="text-xs text-green-600 mt-1">
{result.message}
</div>
)}
</div>
);
}

if (toolName === UploadImageTool.toolName) {
const args = toolPart.input as z.infer<typeof UploadImageTool.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_id && (
<div className="text-xs text-foreground-secondary">
Image ID: {args.image_id}
</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 === WriteFileTool.toolName) {
const args = toolPart.input as z.infer<typeof WriteFileTool.parameters> | null;
const filePath = args?.file_path;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
} from '@onlook/models';
import { toast } from '@onlook/ui/sonner';
import { useEffect, useRef, useState } from 'react';
import { v4 as uuidv4 } from 'uuid';
import { useTabActive } from '../_hooks/use-tab-active';

interface ProjectReadyState {
Expand Down Expand Up @@ -94,6 +95,7 @@ export const useStartProject = () => {
content: context.content,
mimeType: context.mimeType,
displayName: 'user image',
id: uuidv4(),
}));

const context: MessageContext[] = [...createContext, ...imageContexts];
Expand Down
2 changes: 1 addition & 1 deletion apps/web/client/src/components/tools/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -39,4 +39,4 @@ export async function handleToolCall(agentType: AgentType, toolCall: ToolCall<st
});
}

}
}
1 change: 1 addition & 0 deletions bun.lock
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,7 @@
"fg": "^0.0.3",
"gpt-tokenizer": "^3.0.1",
"marked": "^15.0.7",
"mime-lite": "^1.0.3",
"openai": "^4.103.0",
"uuid": "^11.1.0",
"zod": "^4.1.3",
Expand Down
1 change: 1 addition & 0 deletions packages/ai/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
"fg": "^0.0.3",
"gpt-tokenizer": "^3.0.1",
"marked": "^15.0.7",
"mime-lite": "^1.0.3",
"openai": "^4.103.0",
"zod": "^4.1.3",
"@onlook/ui": "*",
Expand Down
6 changes: 5 additions & 1 deletion packages/ai/src/agents/tool-lookup.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { ListFilesTool, ReadFileTool, BashReadTool, OnlookInstructionsTool, ReadStyleGuideTool, ListBranchesTool, ScrapeUrlTool, WebSearchTool, GlobTool, GrepTool, TypecheckTool, CheckErrorsTool, SearchReplaceEditTool, SearchReplaceMultiEditFileTool, FuzzyEditFileTool, WriteFileTool, BashEditTool, SandboxTool, TerminalCommandTool } from "../tools";
import { ListFilesTool, ReadFileTool, BashReadTool, OnlookInstructionsTool, ReadStyleGuideTool, ListBranchesTool, ScrapeUrlTool, WebSearchTool, GlobTool, GrepTool, TypecheckTool, CheckErrorsTool, SearchReplaceEditTool, SearchReplaceMultiEditFileTool, FuzzyEditFileTool, WriteFileTool, BashEditTool, SandboxTool, TerminalCommandTool, ViewImageTool, UploadImageTool } from "../tools";
import { UserAgentTool } from "./tools";

export const allTools = [
Expand All @@ -21,6 +21,8 @@ export const allTools = [
BashEditTool,
SandboxTool,
TerminalCommandTool,
ViewImageTool,
UploadImageTool,
// UserAgentTool,
];

Expand All @@ -38,6 +40,7 @@ export const readOnlyRootTools = [
GrepTool,
TypecheckTool,
CheckErrorsTool,
ViewImageTool,
// UserAgentTool,
]
const editOnlyRootTools = [
Expand All @@ -48,6 +51,7 @@ const editOnlyRootTools = [
BashEditTool,
SandboxTool,
TerminalCommandTool,
UploadImageTool,
// UserAgentTool
]

Expand Down
15 changes: 7 additions & 8 deletions packages/ai/src/prompt/provider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -112,17 +112,16 @@ export function getHydratedUserMessage(
.join('\n');
prompt += wrapXml('instruction', textContent);

userParts.push({ type: 'text', text: prompt });

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) => `- ID: ${img.id}, Name: "${img.displayName}", Type: ${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 with the image ID to analyze the image content\n- Use the "upload_image" tool with the image ID 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',
Expand Down
2 changes: 2 additions & 0 deletions packages/ai/src/tools/classes/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,5 +15,7 @@ export { SearchReplaceEditTool } from './search-replace-edit';
export { SearchReplaceMultiEditFileTool } from './search-replace-multi-edit';
export { TerminalCommandTool } from './terminal-command';
export { TypecheckTool } from './typecheck';
export { UploadImageTool } from './upload-image';
export { ViewImageTool } from './view-image';
export { WebSearchTool } from './web-search';
export { WriteFileTool } from './write-file';
90 changes: 90 additions & 0 deletions packages/ai/src/tools/classes/upload-image.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,90 @@
import { Icons } from '@onlook/ui/icons';
import type { EditorEngine } from '@onlook/web-client/src/components/store/editor/engine';
import { MessageContextType, type MessageContext } from '@onlook/models';
import mime from 'mime-lite';
import { v4 as uuidv4 } from 'uuid';
import { z } from 'zod';
import { ClientTool } from '../models/client';
import { BRANCH_ID_SCHEMA } from '../shared/type';

export class UploadImageTool extends ClientTool {
static readonly toolName = 'upload_image';
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.";
static readonly parameters = z.object({
image_id: z
.string()
.describe(
'The unique ID of the image from the available images list',
),
destination_path: z
.string()
.optional()
.describe('Destination path within the project. Defaults to "public/images" if not specified.'),
filename: z
.string()
.optional()
.describe('Custom filename (without extension). If not provided, a UUID will be generated'),
branchId: BRANCH_ID_SCHEMA,
});
static readonly icon = Icons.Image;

async handle(
args: z.infer<typeof UploadImageTool.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}`);
}

const context = editorEngine.chat.context.context;
const imageContext = context.find((ctx) =>
ctx.type === MessageContextType.IMAGE && ctx.id === args.image_id
);

if (!imageContext || imageContext.type !== MessageContextType.IMAGE) {
throw new Error(`No image found with ID: ${args.image_id}`);
}

const fullPath = await this.uploadImageToSandbox(imageContext, args, sandbox);
await editorEngine.image.scanImages();

return `Image "${imageContext.displayName}" uploaded successfully to ${fullPath}`;
} catch (error) {
throw new Error(`Cannot upload image: ${error}`);
}
}

getLabel(input?: z.infer<typeof UploadImageTool.parameters>): string {
if (input?.image_id) {
return 'Uploading image ' + input.image_id.substring(0, 20);
}
return 'Uploading image';
}

private async uploadImageToSandbox(
imageContext: Extract<MessageContext, { type: MessageContextType.IMAGE }>,
args: z.infer<typeof UploadImageTool.parameters>,
sandbox: any
): Promise<string> {
const mimeType = imageContext.mimeType;
const extension = mime.getExtension(mimeType) || 'png';
const filename = args.filename ? `${args.filename}.${extension}` : `${uuidv4()}.${extension}`;
const destinationPath = args.destination_path?.trim() || 'public/images';
const fullPath = `${destinationPath}/${filename}`;
const base64Data = imageContext.content.replace(/^data:image\/[a-zA-Z0-9+.-]+;base64,/, '');
const binaryData = this.base64ToUint8Array(base64Data);
await sandbox.writeBinaryFile(fullPath, binaryData);
return fullPath;
}

private 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;
}
}
51 changes: 51 additions & 0 deletions packages/ai/src/tools/classes/view-image.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,51 @@
import { Icons } from '@onlook/ui/icons';
import type { EditorEngine } from '@onlook/web-client/src/components/store/editor/engine';
import { MessageContextType } from '@onlook/models';
import { z } from 'zod';
import { ClientTool } from '../models/client';

export class ViewImageTool extends ClientTool {
static readonly toolName = 'view_image';
static readonly 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.";
static readonly parameters = z.object({
image_id: z
.string()
.describe(
'The unique ID of the image from the available images list',
),
});
static readonly icon = Icons.Image;

async handle(
args: z.infer<typeof ViewImageTool.parameters>,
editorEngine: EditorEngine
): Promise<{ image: { mimeType: string; data: string }; message: string }> {
try {
const context = editorEngine.chat.context.context;
const imageContext = context.find((ctx) =>
ctx.type === MessageContextType.IMAGE && ctx.id === args.image_id
);

if (!imageContext || imageContext.type !== MessageContextType.IMAGE) {
throw new Error(`No image found with ID: ${args.image_id}`);
}

return {
image: {
mimeType: imageContext.mimeType,
data: imageContext.content,
},
message: `Retrieved image "${imageContext.displayName}" for analysis.`,
};
} catch (error) {
throw new Error(`Cannot view image: ${error}`);
}
}

getLabel(input?: z.infer<typeof ViewImageTool.parameters>): string {
if (input?.image_id) {
return 'Viewing image ' + input.image_id.substring(0, 20);
}
return 'Viewing image';
}
}
2 changes: 1 addition & 1 deletion packages/ai/src/tools/toolset.ts
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,4 @@ export function getAvailableTools(agentType: AgentType, chatType: ChatType) {
export const TOOLS_MAP = new Map(allTools.map(toolClass => [toolClass.toolName, toolClass]));
const allToolSet = createToolSet(allTools);

export type ChatTools = InferUITools<typeof allToolSet>;
export type ChatTools = InferUITools<typeof allToolSet>;
1 change: 1 addition & 0 deletions packages/models/src/chat/message/context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ export type HighlightMessageContext = BaseMessageContext & {
export type ImageMessageContext = BaseMessageContext & {
type: MessageContextType.IMAGE;
mimeType: string;
id: string;
};

export type ErrorMessageContext = BaseMessageContext & {
Expand Down