Skip to content

Commit dc2f4a2

Browse files
committed
more clean up
1 parent 578893f commit dc2f4a2

File tree

10 files changed

+39
-87
lines changed

10 files changed

+39
-87
lines changed

apps/web/client/src/app/_components/hero/create.tsx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { cn } from '@onlook/ui/utils';
1616
import { compressImageInBrowser } from '@onlook/utility';
1717
import localforage from 'localforage';
1818
import { observer } from 'mobx-react-lite';
19+
import { v4 as uuidv4 } from 'uuid';
1920
import { AnimatePresence } from 'motion/react';
2021
import { useRouter } from 'next/navigation';
2122
import { useEffect, useRef, useState } from 'react';
@@ -200,6 +201,7 @@ export const Create = observer(({
200201
content: base64,
201202
displayName: file.name,
202203
mimeType: file.type,
204+
id: uuidv4(),
203205
};
204206
} catch (error) {
205207
console.error('Error reading file:', error);

apps/web/client/src/app/project/[id]/_components/right-panel/chat-tab/chat-input/index.tsx

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { Tooltip, TooltipContent, TooltipTrigger } from '@onlook/ui/tooltip';
1616
import { cn } from '@onlook/ui/utils';
1717
import { compressImageInBrowser } from '@onlook/utility';
1818
import { observer } from 'mobx-react-lite';
19+
import { v4 as uuidv4 } from 'uuid';
1920
import { useTranslations } from 'next-intl';
2021
import { useEffect, useMemo, useRef, useState } from 'react';
2122
import { validateImageLimit } from '../context-pills/helpers';
@@ -213,6 +214,7 @@ export const ChatInput = observer(({
213214
content: base64URL,
214215
mimeType: file.type,
215216
displayName: displayName ?? file.name,
217+
id: uuidv4(),
216218
};
217219
editorEngine.chat.context.addContexts([contextImage]);
218220
};
@@ -265,6 +267,7 @@ export const ChatInput = observer(({
265267
content: screenshotData,
266268
mimeType: mimeType,
267269
displayName: 'Screenshot',
270+
id: uuidv4(),
268271
};
269272
editorEngine.chat.context.addContexts([contextImage]);
270273
toast.success('Screenshot added to chat');

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,9 +79,9 @@ export const ToolCallDisplay = ({
7979
<Icons.Image className="w-4 h-4" />
8080
<span className="text-sm font-medium">View Image</span>
8181
</div>
82-
{args?.image_reference && (
82+
{args?.image_id && (
8383
<div className="text-xs text-foreground-secondary">
84-
Image: {args.image_reference}
84+
Image ID: {args.image_id}
8585
</div>
8686
)}
8787
{result?.message && (
@@ -102,9 +102,9 @@ export const ToolCallDisplay = ({
102102
<Icons.Image className="w-4 h-4" />
103103
<span className="text-sm font-medium">Upload Image</span>
104104
</div>
105-
{args?.image_reference && (
105+
{args?.image_id && (
106106
<div className="text-xs text-foreground-secondary">
107-
Image: {args.image_reference}
107+
Image ID: {args.image_id}
108108
</div>
109109
)}
110110
{args?.destination_path && (

apps/web/client/src/app/project/[id]/_hooks/use-start-project.tsx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import {
1313
} from '@onlook/models';
1414
import { toast } from '@onlook/ui/sonner';
1515
import { useEffect, useRef, useState } from 'react';
16+
import { v4 as uuidv4 } from 'uuid';
1617
import { useTabActive } from '../_hooks/use-tab-active';
1718

1819
interface ProjectReadyState {
@@ -94,6 +95,7 @@ export const useStartProject = () => {
9495
content: context.content,
9596
mimeType: context.mimeType,
9697
displayName: 'user image',
98+
id: uuidv4(),
9799
}));
98100

99101
const context: MessageContext[] = [...createContext, ...imageContexts];

bun.lock

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@
221221
"fg": "^0.0.3",
222222
"gpt-tokenizer": "^3.0.1",
223223
"marked": "^15.0.7",
224+
"mime-lite": "^1.0.3",
224225
"openai": "^4.103.0",
225226
"uuid": "^11.1.0",
226227
"zod": "^4.1.3",

packages/ai/package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
"fg": "^0.0.3",
4747
"gpt-tokenizer": "^3.0.1",
4848
"marked": "^15.0.7",
49+
"mime-lite": "^1.0.3",
4950
"openai": "^4.103.0",
5051
"zod": "^4.1.3",
5152
"@onlook/ui": "*",

packages/ai/src/prompt/provider.ts

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -112,13 +112,11 @@ export function getHydratedUserMessage(
112112
.join('\n');
113113
prompt += wrapXml('instruction', textContent);
114114

115-
// Add image references to prompt (but doesnt send image data yet)
116-
// AI will decide whether to view or upload them using tools
117115
if (images.length > 0) {
118116
const imageList = images
119-
.map((img, idx) => `${idx + 1}. "${img.displayName}" (${img.mimeType})`)
117+
.map((img) => `- ID: ${img.id}, Name: "${img.displayName}", Type: ${img.mimeType}`)
120118
.join('\n');
121-
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.`;
119+
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.`;
122120
prompt += wrapXml('available-images', imagesPrompt);
123121
}
124122

packages/ai/src/tools/classes/upload-image.ts

Lines changed: 15 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -1,24 +1,25 @@
11
import { Icons } from '@onlook/ui/icons';
22
import type { EditorEngine } from '@onlook/web-client/src/components/store/editor/engine';
3-
import { MessageContextType } from '@onlook/models';
3+
import { MessageContextType, type MessageContext } from '@onlook/models';
4+
import mime from 'mime-lite';
45
import { v4 as uuidv4 } from 'uuid';
56
import { z } from 'zod';
67
import { ClientTool } from '../models/client';
78
import { BRANCH_ID_SCHEMA } from '../shared/type';
89

910
export class UploadImageTool extends ClientTool {
1011
static readonly toolName = 'upload_image';
11-
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 the project's public directory and can be referenced in code. After uploading, you can use the file path in your code changes.";
12+
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.";
1213
static readonly parameters = z.object({
13-
image_reference: z
14+
image_id: z
1415
.string()
1516
.describe(
16-
'Reference to an image in the chat context (use the display name or index number)',
17+
'The unique ID of the image from the available images list',
1718
),
1819
destination_path: z
1920
.string()
2021
.optional()
21-
.describe('Destination path within the project (default: "public/assets/images")'),
22+
.describe('Destination path within the project. Defaults to "public/images" if not specified.'),
2223
filename: z
2324
.string()
2425
.optional()
@@ -38,31 +39,12 @@ export class UploadImageTool extends ClientTool {
3839
}
3940

4041
const context = editorEngine.chat.context.context;
41-
const imageContext = context.find((ctx) => {
42-
if (ctx.type !== MessageContextType.IMAGE) {
43-
return false;
44-
}
45-
return ctx.displayName.toLowerCase().includes(args.image_reference.toLowerCase()) ||
46-
args.image_reference.toLowerCase().includes(ctx.displayName.toLowerCase());
47-
});
42+
const imageContext = context.find((ctx) =>
43+
ctx.type === MessageContextType.IMAGE && ctx.id === args.image_id
44+
);
4845

4946
if (!imageContext || imageContext.type !== MessageContextType.IMAGE) {
50-
const recentImages = context.filter(ctx => ctx.type === MessageContextType.IMAGE);
51-
if (recentImages.length === 0) {
52-
throw new Error(`No image found matching reference: ${args.image_reference}`);
53-
}
54-
55-
const mostRecentImage = recentImages[recentImages.length - 1];
56-
if (!mostRecentImage || mostRecentImage.type !== MessageContextType.IMAGE) {
57-
throw new Error(`No image found matching reference: ${args.image_reference}`);
58-
}
59-
60-
console.warn(`No exact match for "${args.image_reference}", using most recent image: ${mostRecentImage.displayName}`);
61-
62-
const fullPath = await this.uploadImageToSandbox(mostRecentImage, args, sandbox);
63-
await editorEngine.image.scanImages();
64-
65-
return `Image "${mostRecentImage.displayName}" uploaded successfully to ${fullPath}`;
47+
throw new Error(`No image found with ID: ${args.image_id}`);
6648
}
6749

6850
const fullPath = await this.uploadImageToSandbox(imageContext, args, sandbox);
@@ -75,42 +57,28 @@ export class UploadImageTool extends ClientTool {
7557
}
7658

7759
getLabel(input?: z.infer<typeof UploadImageTool.parameters>): string {
78-
if (input?.image_reference) {
79-
return 'Uploading image ' + input.image_reference.substring(0, 20);
60+
if (input?.image_id) {
61+
return 'Uploading image ' + input.image_id.substring(0, 20);
8062
}
8163
return 'Uploading image';
8264
}
8365

8466
private async uploadImageToSandbox(
85-
imageContext: Extract<import('@onlook/models').MessageContext, { type: MessageContextType.IMAGE }>,
67+
imageContext: Extract<MessageContext, { type: MessageContextType.IMAGE }>,
8668
args: z.infer<typeof UploadImageTool.parameters>,
8769
sandbox: any
8870
): Promise<string> {
8971
const mimeType = imageContext.mimeType;
90-
const extension = this.getExtensionFromMimeType(mimeType);
72+
const extension = mime.getExtension(mimeType) || 'png';
9173
const filename = args.filename ? `${args.filename}.${extension}` : `${uuidv4()}.${extension}`;
92-
const destinationPath = args.destination_path || 'public/assets/images';
74+
const destinationPath = args.destination_path?.trim() || 'public/images';
9375
const fullPath = `${destinationPath}/${filename}`;
9476
const base64Data = imageContext.content.replace(/^data:image\/[a-zA-Z0-9+.-]+;base64,/, '');
9577
const binaryData = this.base64ToUint8Array(base64Data);
9678
await sandbox.writeBinaryFile(fullPath, binaryData);
9779
return fullPath;
9880
}
9981

100-
private getExtensionFromMimeType(mimeType: string): string {
101-
const mimeToExt: Record<string, string> = {
102-
'image/jpeg': 'jpg',
103-
'image/jpg': 'jpg',
104-
'image/png': 'png',
105-
'image/gif': 'gif',
106-
'image/webp': 'webp',
107-
'image/svg+xml': 'svg',
108-
'image/bmp': 'bmp',
109-
'image/tiff': 'tiff',
110-
};
111-
return mimeToExt[mimeType.toLowerCase()] || 'png';
112-
}
113-
11482
private base64ToUint8Array(base64: string): Uint8Array {
11583
const binaryString = atob(base64);
11684
const bytes = new Uint8Array(binaryString.length);

packages/ai/src/tools/classes/view-image.ts

Lines changed: 8 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,10 @@ export class ViewImageTool extends ClientTool {
88
static readonly toolName = 'view_image';
99
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.";
1010
static readonly parameters = z.object({
11-
image_reference: z
11+
image_id: z
1212
.string()
1313
.describe(
14-
'Reference to an image in the chat context (use the display name or index number)',
14+
'The unique ID of the image from the available images list',
1515
),
1616
});
1717
static readonly icon = Icons.Image;
@@ -22,36 +22,12 @@ export class ViewImageTool extends ClientTool {
2222
): Promise<{ image: { mimeType: string; data: string }; message: string }> {
2323
try {
2424
const context = editorEngine.chat.context.context;
25-
const imageContext = context.find((ctx) => {
26-
if (ctx.type !== MessageContextType.IMAGE) {
27-
return false;
28-
}
29-
const ref = args.image_reference.toLowerCase();
30-
return ctx.displayName.toLowerCase().includes(ref) ||
31-
ref.includes(ctx.displayName.toLowerCase()) ||
32-
ref.match(/^\d+$/) && context.filter(c => c.type === MessageContextType.IMAGE)[parseInt(ref) - 1] === ctx;
33-
});
25+
const imageContext = context.find((ctx) =>
26+
ctx.type === MessageContextType.IMAGE && ctx.id === args.image_id
27+
);
3428

3529
if (!imageContext || imageContext.type !== MessageContextType.IMAGE) {
36-
const imageContexts = context.filter(ctx => ctx.type === MessageContextType.IMAGE);
37-
const indexMatch = args.image_reference.match(/^\d+$/);
38-
if (indexMatch) {
39-
const index = parseInt(indexMatch[0]) - 1;
40-
if (index >= 0 && index < imageContexts.length) {
41-
const foundImage = imageContexts[index];
42-
if (foundImage && foundImage.type === MessageContextType.IMAGE) {
43-
return {
44-
image: {
45-
mimeType: foundImage.mimeType,
46-
data: foundImage.content,
47-
},
48-
message: `Retrieved image "${foundImage.displayName}" for analysis.`,
49-
};
50-
}
51-
}
52-
}
53-
54-
throw new Error(`No image found matching reference: ${args.image_reference}`);
30+
throw new Error(`No image found with ID: ${args.image_id}`);
5531
}
5632

5733
return {
@@ -67,8 +43,8 @@ export class ViewImageTool extends ClientTool {
6743
}
6844

6945
getLabel(input?: z.infer<typeof ViewImageTool.parameters>): string {
70-
if (input?.image_reference) {
71-
return 'Viewing image ' + input.image_reference.substring(0, 20);
46+
if (input?.image_id) {
47+
return 'Viewing image ' + input.image_id.substring(0, 20);
7248
}
7349
return 'Viewing image';
7450
}

packages/models/src/chat/message/context.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ export type HighlightMessageContext = BaseMessageContext & {
3838
export type ImageMessageContext = BaseMessageContext & {
3939
type: MessageContextType.IMAGE;
4040
mimeType: string;
41+
id: string;
4142
};
4243

4344
export type ErrorMessageContext = BaseMessageContext & {

0 commit comments

Comments
 (0)