Skip to content

Commit ae0f3ba

Browse files
NickB03claude
NickB03
and
claude
committed
feat: add image generation tool with Gemini Flash model
Adds a generateImage tool that creates and edits images using Google Gemini 2.5 Flash Image, with Supabase storage for persistence. - Tool definition with prompt, aspectRatio, and sourceImageUrl params - Server-side upload to Supabase storage (per-user/chat paths) - Compact chat thumbnail (200px max) with always-visible download - Lightbox preview with download, Escape-to-close, scroll lock - Blob-based download to handle cross-origin Supabase URLs - Tool UI registry integration and message persistence mapping - Chat/research mode prompt routing for image intent - Schema validation and test coverage Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 5ad9205 commit ae0f3ba

14 files changed

Lines changed: 635 additions & 9 deletions

File tree

‎components/tool-ui/generate-image/__tests__/schema.test.ts‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
import { describe, expect, it } from 'vitest'
2+
3+
import {
4+
parseSerializableGenerateImage,
5+
safeParseSerializableGenerateImage
6+
} from '../schema'
7+
8+
describe('safeParseSerializableGenerateImage', () => {
9+
it('parses valid output with all fields', () => {
10+
const result = safeParseSerializableGenerateImage({
11+
imageUrl: 'https://example.com/image.png',
12+
filename: 'generated-123.png',
13+
mediaType: 'image/png',
14+
description: 'a sunset',
15+
aspectRatio: '16:9'
16+
})
17+
expect(result).toEqual({
18+
imageUrl: 'https://example.com/image.png',
19+
filename: 'generated-123.png',
20+
mediaType: 'image/png',
21+
description: 'a sunset',
22+
aspectRatio: '16:9'
23+
})
24+
})
25+
26+
it('parses valid output without optional fields', () => {
27+
const result = safeParseSerializableGenerateImage({
28+
imageUrl: 'https://example.com/image.png',
29+
filename: 'generated-123.png',
30+
mediaType: 'image/png',
31+
description: 'a sunset'
32+
})
33+
expect(result).not.toBeNull()
34+
expect(result?.aspectRatio).toBeUndefined()
35+
})
36+
37+
it('returns null for error output', () => {
38+
const result = safeParseSerializableGenerateImage({
39+
error: 'something failed'
40+
})
41+
expect(result).toBeNull()
42+
})
43+
44+
it('returns null for invalid input', () => {
45+
expect(safeParseSerializableGenerateImage(null)).toBeNull()
46+
expect(safeParseSerializableGenerateImage({})).toBeNull()
47+
expect(safeParseSerializableGenerateImage('string')).toBeNull()
48+
})
49+
})

‎components/tool-ui/generate-image/generate-image.tsx‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
'use client'
2+
3+
import { useCallback, useEffect, useState } from 'react'
4+
5+
import { Download, X } from 'lucide-react'
6+
7+
import { Button } from '@/components/ui/button'
8+
9+
import type { GenerateImageProps } from './schema'
10+
11+
export function GenerateImage({
12+
imageUrl,
13+
filename,
14+
description,
15+
aspectRatio
16+
}: GenerateImageProps) {
17+
const [expanded, setExpanded] = useState(false)
18+
19+
useEffect(() => {
20+
if (!expanded) return
21+
const onKey = (e: KeyboardEvent) => {
22+
if (e.key === 'Escape') setExpanded(false)
23+
}
24+
document.addEventListener('keydown', onKey)
25+
document.body.style.overflow = 'hidden'
26+
return () => {
27+
document.removeEventListener('keydown', onKey)
28+
document.body.style.overflow = ''
29+
}
30+
}, [expanded])
31+
32+
const handleDownload = useCallback(
33+
async (e: React.MouseEvent) => {
34+
e.stopPropagation()
35+
try {
36+
const res = await fetch(imageUrl)
37+
const blob = await res.blob()
38+
const url = URL.createObjectURL(blob)
39+
const a = document.createElement('a')
40+
a.href = url
41+
a.download = filename
42+
document.body.appendChild(a)
43+
a.click()
44+
document.body.removeChild(a)
45+
URL.revokeObjectURL(url)
46+
} catch {
47+
window.open(imageUrl, '_blank')
48+
}
49+
},
50+
[imageUrl, filename]
51+
)
52+
53+
return (
54+
<>
55+
<figure className="group relative my-3 w-fit max-w-full overflow-hidden rounded-xl border border-border/50 bg-muted/30">
56+
<button
57+
type="button"
58+
onClick={() => setExpanded(true)}
59+
className="block cursor-zoom-in"
60+
>
61+
{/* eslint-disable-next-line @next/next/no-img-element -- dynamic external URL */}
62+
<img
63+
src={imageUrl}
64+
alt={description}
65+
className="max-h-[200px] w-auto max-w-full rounded-t-xl object-contain"
66+
loading="lazy"
67+
/>
68+
</button>
69+
<figcaption className="flex items-center justify-between gap-2 px-3 py-2 text-xs text-muted-foreground">
70+
<span className="line-clamp-1">{description}</span>
71+
<div className="flex shrink-0 items-center gap-1">
72+
{aspectRatio && (
73+
<span className="rounded bg-muted px-1.5 py-0.5 text-[10px] font-medium">
74+
{aspectRatio}
75+
</span>
76+
)}
77+
<Button
78+
variant="ghost"
79+
size="icon"
80+
className="size-6"
81+
onClick={handleDownload}
82+
>
83+
<Download className="size-3" />
84+
</Button>
85+
</div>
86+
</figcaption>
87+
</figure>
88+
89+
{expanded && (
90+
<div
91+
className="fixed inset-0 z-50 flex items-center justify-center bg-black/80 p-4 sm:p-8"
92+
onClick={() => setExpanded(false)}
93+
role="dialog"
94+
aria-modal="true"
95+
aria-label={description}
96+
>
97+
<div className="absolute right-3 top-3 flex items-center gap-1 sm:right-4 sm:top-4 sm:gap-2">
98+
<Button
99+
variant="ghost"
100+
size="icon"
101+
className="size-8 text-white hover:bg-white/10 sm:size-9"
102+
onClick={handleDownload}
103+
>
104+
<Download className="size-4 sm:size-5" />
105+
</Button>
106+
<Button
107+
variant="ghost"
108+
size="icon"
109+
className="size-8 text-white hover:bg-white/10 sm:size-9"
110+
onClick={() => setExpanded(false)}
111+
>
112+
<X className="size-4 sm:size-5" />
113+
</Button>
114+
</div>
115+
{/* eslint-disable-next-line @next/next/no-img-element -- lightbox, dynamic URL */}
116+
<img
117+
src={imageUrl}
118+
alt={description}
119+
className="max-h-[85vh] max-w-[90vw] rounded-lg object-contain"
120+
onClick={e => e.stopPropagation()}
121+
/>
122+
</div>
123+
)}
124+
</>
125+
)
126+
}

‎components/tool-ui/generate-image/index.tsx‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
export { GenerateImage } from './generate-image'
2+
export {
3+
type GenerateImageProps,
4+
parseSerializableGenerateImage,
5+
safeParseSerializableGenerateImage
6+
} from './schema'

‎components/tool-ui/generate-image/schema.ts‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
import { z } from 'zod'
2+
3+
import { defineToolUiContract } from '../shared/contract'
4+
5+
export const GenerateImagePropsSchema = z.object({
6+
imageUrl: z.string().url(),
7+
filename: z.string().min(1),
8+
mediaType: z.string().min(1),
9+
description: z.string().min(1),
10+
aspectRatio: z.string().optional()
11+
})
12+
13+
export type GenerateImageProps = z.infer<typeof GenerateImagePropsSchema>
14+
15+
export const SerializableGenerateImageSchema = GenerateImagePropsSchema
16+
17+
export type SerializableGenerateImage = z.infer<
18+
typeof SerializableGenerateImageSchema
19+
>
20+
21+
const SerializableGenerateImageContract = defineToolUiContract(
22+
'GenerateImage',
23+
SerializableGenerateImageSchema
24+
)
25+
26+
export const parseSerializableGenerateImage: (
27+
input: unknown
28+
) => SerializableGenerateImage = SerializableGenerateImageContract.parse
29+
30+
export const safeParseSerializableGenerateImage: (
31+
input: unknown
32+
) => SerializableGenerateImage | null =
33+
SerializableGenerateImageContract.safeParse

‎components/tool-ui/registry.tsx‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ import { CitationList } from './citation/citation-list'
1010
import { safeParseSerializableCitation } from './citation/schema'
1111
import { DataTable } from './data-table/data-table'
1212
import { safeParseSerializableDataTable } from './data-table/schema'
13+
import { GenerateImage } from './generate-image/generate-image'
14+
import { safeParseSerializableGenerateImage } from './generate-image/schema'
1315
import { LinkPreview } from './link-preview/link-preview'
1416
import { safeParseSerializableLinkPreview } from './link-preview/schema'
1517
import { OptionList } from './option-list/option-list'
@@ -152,6 +154,18 @@ const entries: ToolUIEntry[] = [
152154
)
153155
}
154156
},
157+
{
158+
name: 'generateImage',
159+
tryRender: output => {
160+
const parsed = safeParseSerializableGenerateImage(output)
161+
if (!parsed) return null
162+
return (
163+
<ToolErrorBoundary toolName="GenerateImage">
164+
<GenerateImage {...parsed} />
165+
</ToolErrorBoundary>
166+
)
167+
}
168+
},
155169
{
156170
name: 'canvasArtifactCard',
157171
tryRender: output => tryRenderCanvasArtifactCard(output)

‎lib/agents/prompts/search-mode-prompts.ts‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -114,6 +114,30 @@ You can create and update interactive frontend web artifacts using the tools bel
114114
${ARTIFACT_INTAKE_PROTOCOL}`
115115
}
116116

117+
function getImageGenerationPrompt(): string {
118+
return `
119+
IMAGE GENERATION:
120+
You have a \`generateImage\` tool that creates or edits images using an AI image model.
121+
122+
**When to use:**
123+
- The user asks you to create, generate, draw, illustrate, or visualize an image
124+
- The user wants a visual representation of something (diagram, mockup, concept art, photo, etc.)
125+
- The user asks to modify or edit a previously generated image
126+
127+
**How to use:**
128+
- Provide a detailed, descriptive prompt — the more specific, the better the result
129+
- Include details about: subject, style, composition, lighting, colors, mood, perspective
130+
- Set aspectRatio when the user specifies a format or when the content has a natural shape (landscape → 16:9, portrait → 9:16, square → 1:1)
131+
- For image editing: pass the sourceImageUrl of a previously generated image along with edit instructions in the prompt
132+
133+
**Important:**
134+
- Do NOT search the web before generating an image unless the user needs reference information
135+
- Generate the image directly when the request is clear
136+
- After generating, continue your response naturally — reference the image in your text
137+
- If the user asks to modify a generated image, use the same tool with the sourceImageUrl parameter
138+
`
139+
}
140+
117141
export function getChatModePrompt(): string {
118142
const hasGeneralProvider = isGeneralSearchProviderAvailable()
119143

@@ -125,6 +149,7 @@ You are a fast, efficient AI assistant optimized for quick responses. You have a
125149
**INTENT ROUTING (check FIRST before anything else):**
126150
Before starting any search or research, determine the user's primary intent:
127151
- **BUILD/CREATE request** — the user wants you to build, create, make, generate, or design an interactive app, widget, dashboard, tracker, tool, calculator, visualization, game, demo, timer, or chart → **Skip search entirely.** Go directly to the CANVAS ARTIFACTS section below. CALL the \`createCanvasArtifact\` tool immediately for specific requests, or run the Artifact Intake Protocol for broad/open requests. Do NOT search the web first — the user wants you to write code, not find information.
152+
- **IMAGE request** — the user wants you to generate, draw, create, illustrate, or visualize an image/picture/photo/illustration → **Call \`generateImage\` tool directly.** Do NOT search first unless the user needs factual reference.
128153
- **MODIFY/UPDATE request** — the user wants to change, fix, improve, or add to an existing artifact → **Skip search.** If the artifact source code is not in the conversation context, CALL \`readCanvasArtifact\` first. Then CALL \`updateCanvasArtifact\` with the full replacement file set.
129154
- **RESEARCH-THEN-BUILD request** — the user wants to learn about a topic AND build something based on the findings (e.g., "research React dashboard best practices and then build me one") → Perform the research phase first (search, gather information), then proceed to canvas artifact tools to build the artifact.
130155
- **FACTUAL/CURRENT-DATA ARTIFACT request** — the user wants to build an artifact that depends on specific entities, freshness, dates, statistics, or other current facts → run a short search phase first, then build the artifact.
@@ -330,7 +355,9 @@ Call the displayTable tool with the comparison data, then continue:
330355
331356
End with a synthesizing conclusion that ties the main points together into a clear overall picture.
332357
333-
${getCanvasArtifactsPrompt()}`
358+
${getCanvasArtifactsPrompt()}
359+
360+
${getImageGenerationPrompt()}`
334361
}
335362

336363
export function getResearchModePrompt(): string {
@@ -660,7 +687,9 @@ Flexible example:
660687
661688
Conclude with a brief synthesis that ties together the main insights into a clear overall understanding.
662689
663-
${getCanvasArtifactsPrompt()}`
690+
${getCanvasArtifactsPrompt()}
691+
692+
${getImageGenerationPrompt()}`
664693
}
665694

666695
// Export static prompts for backward compatibility

‎lib/agents/researcher.ts‎

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import { displayQuestionWizardTool } from '../tools/display-question-wizard'
2121
import { displayTableTool } from '../tools/display-table'
2222
import { displayTimelineTool } from '../tools/display-timeline'
2323
import { fetchTool } from '../tools/fetch'
24+
import { createGenerateImageTool } from '../tools/generate-image'
2425
import { readCanvasArtifactTool } from '../tools/read-canvas-artifact'
2526
import { createSearchTool } from '../tools/search'
2627
import { createTodoTools } from '../tools/todo'
@@ -91,7 +92,8 @@ export function createResearcher({
9192
modelType,
9293
telemetryEnabled,
9394
experimentalContext,
94-
canvasToolContext
95+
canvasToolContext,
96+
imageToolContext
9597
}: {
9698
model: string
9799
modelConfig?: Model
@@ -102,6 +104,7 @@ export function createResearcher({
102104
telemetryEnabled?: boolean
103105
experimentalContext?: unknown
104106
canvasToolContext?: CanvasToolContext
107+
imageToolContext?: { userId: string; chatId: string }
105108
}) {
106109
try {
107110
const currentDate = new Date().toLocaleString()
@@ -210,6 +213,17 @@ export function createResearcher({
210213
)
211214
}
212215

216+
// Build image generation tool when context is available
217+
const imageTools = imageToolContext
218+
? {
219+
generateImage: createGenerateImageTool(imageToolContext)
220+
}
221+
: {}
222+
223+
if (imageToolContext) {
224+
activeToolsList.push('generateImage' as keyof ResearcherTools)
225+
}
226+
213227
// Build tools object with proper typing
214228
const tools: ResearcherTools = {
215229
search: searchTool,
@@ -224,7 +238,8 @@ export function createResearcher({
224238
displayCallout: displayCalloutTool,
225239
displayTimeline: displayTimelineTool,
226240
...todoTools,
227-
...canvasTools
241+
...canvasTools,
242+
...imageTools
228243
} as ResearcherTools
229244

230245
// Create ToolLoopAgent with all configuration

‎lib/streaming/create-chat-stream-response.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,8 @@ export async function createChatStreamResponse(
206206
parentTraceId,
207207
searchMode,
208208
modelType,
209-
canvasToolContext
209+
canvasToolContext,
210+
imageToolContext: { userId, chatId }
210211
})
211212

212213
// For OpenAI models, strip reasoning parts from UIMessages before conversion

‎lib/supabase/__tests__/server-storage.test.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
import { describe, expect, it } from 'vitest'
2+
3+
import { buildGeneratedImagePath } from '../server-storage'
4+
5+
describe('buildGeneratedImagePath', () => {
6+
it('constructs path from userId, chatId, and extension', () => {
7+
const path = buildGeneratedImagePath('user-1', 'chat-1', 'image/png')
8+
expect(path).toMatch(/^user-1\/chats\/chat-1\/generated-\d+\.png$/)
9+
})
10+
11+
it('extracts extension from mediaType', () => {
12+
const path = buildGeneratedImagePath('u', 'c', 'image/webp')
13+
expect(path.endsWith('.webp')).toBe(true)
14+
})
15+
16+
it('defaults extension from mediaType subtype', () => {
17+
const path = buildGeneratedImagePath('u', 'c', 'image/jpeg')
18+
expect(path.endsWith('.jpeg')).toBe(true)
19+
})
20+
})

0 commit comments

Comments
 (0)