diff --git a/packages/docs/src/content/docs/reference/config-and-env.md b/packages/docs/src/content/docs/reference/config-and-env.md index a51cf1ae7c..1d370b21d4 100644 --- a/packages/docs/src/content/docs/reference/config-and-env.md +++ b/packages/docs/src/content/docs/reference/config-and-env.md @@ -319,7 +319,7 @@ const app = await createApp({ guardianModelId: "openai/gpt-6-luna", embeddingModelId: "openai/text-embedding-3-small", webSearchModelId: "openai/gpt-6-luna", - imageGenerationModelId: "google/gemini-3-pro-image", + imageGenerationModelId: "openai/gpt-image-2", visionModelId: "openai/gpt-5.6-sol", }); ``` diff --git a/packages/junior-evals/src/fixture/gateway.ts b/packages/junior-evals/src/fixture/gateway.ts index 815e4f87ed..c37a405411 100644 --- a/packages/junior-evals/src/fixture/gateway.ts +++ b/packages/junior-evals/src/fixture/gateway.ts @@ -16,22 +16,14 @@ import { mswServer } from "@junior-tests/msw/server"; import { answerWebSearch } from "./web"; const GATEWAY_MESSAGES_URL = "https://ai-gateway.vercel.sh/v1/messages"; -const GATEWAY_CHAT_COMPLETIONS_URL = - "https://ai-gateway.vercel.sh/v1/chat/completions"; +// The AI SDK `generateImage()` call of the gateway provider. +const GATEWAY_IMAGE_MODEL_URL = + "https://ai-gateway.vercel.sh/v4/ai/image-model"; // A valid 1x1 PNG. The agent can view it, so the model must decode it. -const STUB_PNG_DATA_URL = - "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII="; +const STUB_PNG_BASE64 = + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII="; -/** Whether a request asks the gateway to generate an image. */ -async function isImageGeneration(request: Request): Promise { - if (!request.url.startsWith(GATEWAY_CHAT_COMPLETIONS_URL)) return false; - const payload = (await request.clone().json()) as { modalities?: unknown }; - return ( - Array.isArray(payload.modalities) && payload.modalities.includes("image") - ); -} - -let generatedImageDataUrl = STUB_PNG_DATA_URL; +let generatedImageBase64 = STUB_PNG_BASE64; /** * Set the PNG that image generation returns in this test. Use it when the @@ -39,9 +31,9 @@ let generatedImageDataUrl = STUB_PNG_DATA_URL; * look like the result, so the agent may try again. */ export function mockGeneratedImage(png: Buffer): void { - generatedImageDataUrl = `data:image/png;base64,${png.toString("base64")}`; + generatedImageBase64 = png.toString("base64"); onTestFinished(() => { - generatedImageDataUrl = STUB_PNG_DATA_URL; + generatedImageBase64 = STUB_PNG_BASE64; }); } @@ -135,16 +127,8 @@ export function installGatewayObserver(): GatewayObserver { counts[endpoint] = (counts[endpoint] ?? 0) + 1; const searchResponse = await answerWebSearch(request); if (searchResponse) return searchResponse; - if (await isImageGeneration(request)) { - return HttpResponse.json({ - choices: [ - { - message: { - images: [{ image_url: { url: generatedImageDataUrl } }], - }, - }, - ], - }); + if (request.url === GATEWAY_IMAGE_MODEL_URL) { + return HttpResponse.json({ images: [generatedImageBase64] }); } const hook = progressHook; if ( diff --git a/packages/junior/scripts/tool-error-classification-baseline.txt b/packages/junior/scripts/tool-error-classification-baseline.txt index 2f5bcd6fa9..ead5a7ddc4 100644 --- a/packages/junior/scripts/tool-error-classification-baseline.txt +++ b/packages/junior/scripts/tool-error-classification-baseline.txt @@ -57,5 +57,5 @@ packages/junior/src/chat/tools/user-lookup.ts throw new Error( `Multiple Slack u # Tool-call identity comes from the runtime, not model arguments. packages/junior/src/chat/tools/watch-timer.ts throw new Error("Timer creation requires a tool-call identity"); packages/junior/src/chat/tools/web/fetch-content.ts throw new Error(message); -packages/junior/src/chat/tools/web/image-generate.ts throw new Error( parseImageGenerationError(response.status, text, model), ); +packages/junior/src/chat/tools/web/image-generate.ts throw new Error( `image generation failed: configured model "${model}" is not an image generation model. Set AI_IMAGE_MODEL to an AI Gateway image model.`, ); packages/junior/src/chat/tools/user-customization.ts throw new Error("User customization target user was not found."); diff --git a/packages/junior/src/chat/config.ts b/packages/junior/src/chat/config.ts index 8a5376d0ab..5e543b6203 100644 --- a/packages/junior/src/chat/config.ts +++ b/packages/junior/src/chat/config.ts @@ -268,7 +268,7 @@ const DEFAULT_HANDOFF_MODEL_ID = resolveGatewayModel( const DEFAULT_WEB_SEARCH_MODEL_ID = resolveGatewayModel("openai/gpt-6-luna").id; const DEFAULT_DECISION_MODEL_ID = "typesafe-ai/jev"; const DEFAULT_EMBEDDING_MODEL_ID = "openai/text-embedding-3-small"; -const DEFAULT_IMAGE_GENERATION_MODEL_ID = "google/gemini-3-pro-image"; +const DEFAULT_IMAGE_GENERATION_MODEL_ID = "openai/gpt-image-2"; const DEFAULT_VISION_MODEL_ID = resolveGatewayModel("openai/gpt-5.6-sol").id; function validateGatewayModelId(raw: string | undefined): string | undefined { diff --git a/packages/junior/src/chat/tools/web/image-generate.ts b/packages/junior/src/chat/tools/web/image-generate.ts index ab44271857..027cc2da18 100644 --- a/packages/junior/src/chat/tools/web/image-generate.ts +++ b/packages/junior/src/chat/tools/web/image-generate.ts @@ -2,13 +2,14 @@ import { z } from "zod"; import { juniorToolOutputSchema } from "@/chat/tool-support/structured-result"; import { zodTool } from "@/chat/tool-support/zod-tool"; import type { ImageGenerateToolDeps, ToolHooks } from "@/chat/tools/types"; +import { createGatewayProvider } from "@ai-sdk/gateway"; +import { generateImage, NoImageGeneratedError, type GeneratedFile } from "ai"; import { botConfig } from "@/chat/config"; import { completeText, getGatewayApiKey } from "@/chat/pi/client"; import { JUNIOR_PERSONALITY } from "@/chat/prompt"; import { logInfo, logWarn } from "@/chat/logging"; import { generatedArtifactFileRefSchema } from "@/chat/tools/sandbox/file-uploads"; import { readSandboxImage } from "@/chat/tools/sandbox/image-files"; -import type { InputImageType } from "@/chat/attachments/media"; import { ToolInputError } from "@/chat/tools/execution/tool-input-error"; import type { SandboxWorkspace } from "@/chat/sandbox/workspace"; @@ -21,7 +22,8 @@ const MAX_SOURCE_IMAGE_BYTES = 25 * 1024 * 1024; // Longest edges to try when an image must be resized, largest first. const RESIZE_LONGEST_EDGES = [2048, 1536, 1024, 768]; -type InputImage = { path: string; data: Buffer; mimeType: InputImageType }; +// The AI SDK reads the media type from the image bytes. +type InputImage = { path: string; data: Buffer }; const resizedImageSchema = z .object({ @@ -87,7 +89,6 @@ async function shrinkImage( const { default: sharp } = await import("sharp"); const { hasAlpha } = await sharp(image.data).metadata(); // JPEG is smallest for photos. WebP keeps transparency. - const mimeType = hasAlpha ? "image/webp" : "image/jpeg"; for (const edge of RESIZE_LONGEST_EDGES) { const resized = sharp(image.data).rotate().resize({ width: edge, @@ -100,7 +101,7 @@ async function shrinkImage( ).toBuffer({ resolveWithObject: true }); if (data.byteLength <= maxBytes) { return { - image: { path: image.path, data, mimeType }, + image: { path: image.path, data }, width: info.width, height: info.height, }; @@ -159,26 +160,6 @@ function extensionForMediaType(mediaType: string): string { return "bin"; } -function parseImageGenerationError( - status: number, - body: string, - model: string, -): string { - if (!body) return `image generation failed: ${status}`; - - try { - const payload = JSON.parse(body) as { error?: { message?: string } }; - const message = payload.error?.message?.trim(); - if (!message) return `image generation failed: ${status} ${body}`; - if (message.includes("not an image model")) { - return `image generation failed: configured model "${model}" is not an image generation model. Set AI_IMAGE_MODEL to a compatible image model.`; - } - return `image generation failed: ${status} ${message}`; - } catch { - return `image generation failed: ${status} ${body}`; - } -} - /** Create the image tool with sandbox artifact guidance matched to file-send capability. */ export function createImageGenerateTool( hooks: Required>, @@ -216,10 +197,10 @@ export function createImageGenerateTool( }), outputSchema: imageGenerateOutputSchema, execute: async ({ prompt, images = [] }) => { - const fetchImpl = deps.fetch ?? fetch; - // Raw fetch does not resolve AI Gateway auth on its own, so this path - // has to turn the shared gateway credential into a bearer token. - const apiKey = await getGatewayApiKey(); + const gateway = createGatewayProvider({ + apiKey: await getGatewayApiKey(), + ...(deps.fetch ? { fetch: deps.fetch } : undefined), + }); const model = options.modelId; const sourceImages = await Promise.all( images.map(async (path) => ({ @@ -245,83 +226,40 @@ export function createImageGenerateTool( // Enrichment restyles the whole image, so it would undo an edit. const enrichedPrompt = inputImages.length > 0 ? prompt : await enrichImagePrompt(prompt); - const content = - inputImages.length > 0 - ? [ - { type: "text", text: enrichedPrompt }, - ...inputImages.map((image) => ({ - type: "image_url", - image_url: { - url: `data:${image.mimeType};base64,${image.data.toString("base64")}`, - }, - })), - ] - : enrichedPrompt; - const response = await fetchImpl( - "https://ai-gateway.vercel.sh/v1/chat/completions", - { - method: "POST", - headers: { - "content-type": "application/json", - authorization: `Bearer ${apiKey}`, - }, - body: JSON.stringify({ - model, - messages: [{ role: "user", content }], - modalities: ["image"], - }), - }, - ); - - if (!response.ok) { - const text = await response.text(); - throw new Error( - parseImageGenerationError(response.status, text, model), - ); - } - - const payload = (await response.json()) as { - choices?: Array<{ - message?: { - images?: Array<{ - image_url?: { - url?: string; - }; - }>; - }; - }>; - }; - const uploads: Array<{ - data: Buffer; - filename: string; - mimeType: string; - }> = []; - const generatedImages = payload.choices?.[0]?.message?.images ?? []; - for (const [index, image] of generatedImages.entries()) { - let bytes: Buffer | null = null; - let mimeType = "image/png"; - const url = image.image_url?.url; - - if (typeof url === "string" && url.startsWith("data:")) { - const match = url.match(/^data:([^;,]+);base64,(.+)$/); - if (!match) continue; - mimeType = match[1] ?? mimeType; - bytes = Buffer.from(match[2] ?? "", "base64"); - } else if (typeof url === "string" && url.length > 0) { - const fetched = await fetchImpl(url); - if (!fetched.ok) continue; - mimeType = fetched.headers.get("content-type") ?? mimeType; - bytes = Buffer.from(await fetched.arrayBuffer()); + let generatedImages: GeneratedFile[]; + try { + ({ images: generatedImages } = await generateImage({ + model: gateway.imageModel(model), + // Input images make this an edit of those images. + prompt: + inputImages.length > 0 + ? { + text: enrichedPrompt, + images: inputImages.map((image) => image.data), + } + : enrichedPrompt, + n: 1, + })); + } catch (error) { + if (NoImageGeneratedError.isInstance(error)) { + generatedImages = []; + } else if ( + error instanceof Error && + error.message.includes("not an image model") + ) { + // The model id comes from app config, so the agent cannot fix it. + throw new Error( + `image generation failed: configured model "${model}" is not an image generation model. Set AI_IMAGE_MODEL to an AI Gateway image model.`, + ); + } else { + throw error; } - - if (!bytes) continue; - const extension = extensionForMediaType(mimeType); - uploads.push({ - data: bytes, - filename: `generated-image-${Date.now()}-${index + 1}.${extension}`, - mimeType, - }); } + const uploads = generatedImages.map((image, index) => ({ + data: Buffer.from(image.uint8Array), + filename: `generated-image-${Date.now()}-${index + 1}.${extensionForMediaType(image.mediaType)}`, + mimeType: image.mediaType, + })); if (uploads.length > 0) { const artifactRefs = await hooks.writeGeneratedArtifacts(uploads); diff --git a/packages/junior/tests/component/config/chat-config.test.ts b/packages/junior/tests/component/config/chat-config.test.ts index a783a5b0d1..867bf9861c 100644 --- a/packages/junior/tests/component/config/chat-config.test.ts +++ b/packages/junior/tests/component/config/chat-config.test.ts @@ -286,7 +286,7 @@ describe("chat config", () => { process.env.AI_IMAGE_MODEL = " "; const { botConfig } = await loadConfig(); - expect(botConfig.imageGenerationModelId).toBe("google/gemini-3-pro-image"); + expect(botConfig.imageGenerationModelId).toBe("openai/gpt-image-2"); }); it("uses AI_IMAGE_MODEL when configured", async () => { diff --git a/packages/junior/tests/unit/web/image-generate.test.ts b/packages/junior/tests/unit/web/image-generate.test.ts index e24af4af18..18a005d6af 100644 --- a/packages/junior/tests/unit/web/image-generate.test.ts +++ b/packages/junior/tests/unit/web/image-generate.test.ts @@ -30,49 +30,30 @@ function sandbox(files: Record = {}): SandboxWorkspace { } const defaultImageOptions = { - modelId: "google/gemini-3-pro-image", + modelId: "openai/gpt-image-2", sandbox: sandbox(), }; -function getRequestBody(fetchMock: ReturnType) { - const request = fetchMock.mock.calls[0]!; - expect(request).toBeDefined(); - expect(request[1]).toBeDefined(); - return JSON.parse((request[1] as RequestInit).body as string); -} +const IMAGE_MODEL_URL = "https://ai-gateway.vercel.sh/v4/ai/image-model"; -function createJsonResponse(payload: unknown) { - return { - ok: true, - status: 200, - json: async () => payload, - } as Response; +function imageResponse(images = [PNG_BYTES]) { + return Response.json({ + images: images.map((image) => image.toString("base64")), + }); } -function createErrorResponse(status: number, body: string) { - return { - ok: false, - status, - text: async () => body, - } as Response; +/** Fake AI Gateway that answers each image request in order. */ +function gatewayFetch(responses: Response[]) { + const queue = [...responses]; + return vi.fn( + async (_url: string | URL | Request, _init?: RequestInit) => queue.shift()!, + ); } -function imagePayload() { - return { - choices: [ - { - message: { - images: [ - { - image_url: { - url: `data:image/png;base64,${Buffer.from("img").toString("base64")}`, - }, - }, - ], - }, - }, - ], - }; +function requestBody(fetchMock: ReturnType, index = 0) { + const request = fetchMock.mock.calls[index]!; + expect(String(request[0])).toBe(IMAGE_MODEL_URL); + return JSON.parse(request[1]!.body as string); } function writeGeneratedArtifacts(files: FileUpload[]) { @@ -88,17 +69,13 @@ describe("createImageGenerateTool", () => { afterEach(() => { delete process.env.AI_GATEWAY_API_KEY; delete process.env.VERCEL_OIDC_TOKEN; - vi.unstubAllGlobals(); vi.clearAllMocks(); }); it("uses the configured image generation model", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; mockCompleteText.mockResolvedValueOnce({ text: "enriched prompt" } as any); - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([imageResponse()]); vi.spyOn(Date, "now").mockReturnValue(1_737_000_000_000); const uploads: FileUpload[] = []; @@ -108,71 +85,46 @@ describe("createImageGenerateTool", () => { uploads.push(...files); return writeGeneratedArtifacts(files); }, - } as any, - { modelId: "openai/dall-e-3", sandbox: sandbox() }, + }, + { modelId: "openai/gpt-image-1.5", sandbox: sandbox() }, + { fetch: fetchMock }, ); - if (typeof tool.execute !== "function") { - throw new Error("imageGenerate execute function missing"); - } - const result = await tool.execute({ prompt: "test prompt" }, {} as any); + const result = await tool.execute!({ prompt: "test prompt" }, {} as any); expect(fetchMock).toHaveBeenCalledTimes(1); - const request = fetchMock.mock.calls[0]!; - expect(request).toBeDefined(); - expect(request[0]).toBe("https://ai-gateway.vercel.sh/v1/chat/completions"); - expect(getRequestBody(fetchMock)).toMatchObject({ - model: "openai/dall-e-3", - messages: [{ role: "user", content: "enriched prompt" }], - modalities: ["image"], + expect(requestBody(fetchMock)).toMatchObject({ + prompt: "enriched prompt", + n: 1, }); + const headers = new Headers(fetchMock.mock.calls[0]![1]!.headers); + expect(headers.get("ai-model-id")).toBe("openai/gpt-image-1.5"); + expect(headers.get("authorization")).toBe("Bearer test-key"); expect(result).toMatchObject({ - model: "openai/dall-e-3", + model: "openai/gpt-image-1.5", image_count: 1, - delivery: expect.stringContaining("no file-send tool"), - }); - expect(JSON.stringify(result)).not.toContain("sendFiles"); - const generated = result as { - images: Array<{ - bytes: number; - filename: string; - mimeType: string; - path: string; - }>; - }; - expect(generated).toMatchObject({ images: [ { - bytes: 3, + bytes: PNG_BYTES.byteLength, filename: "generated-image-1737000000000-1.png", mimeType: "image/png", path: "/tmp/junior/artifacts/generated-image-1737000000000-1.png", }, ], + delivery: expect.stringContaining("no file-send tool"), }); - expect(JSON.stringify(generated)).not.toContain("attachment_path"); - expect(JSON.stringify(generated)).not.toContain("media_type"); - expect(uploads[0]?.filename).toContain("generated-image-1737000000000-1"); - expect(uploads[0]?.data).toEqual(Buffer.from("img")); + expect(JSON.stringify(result)).not.toContain("sendFiles"); + expect(uploads[0]?.data).toEqual(PNG_BYTES); }); it("recommends sendFiles when active file-send support is available", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; mockCompleteText.mockResolvedValueOnce({ text: "enriched prompt" } as any); - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); - vi.spyOn(Date, "now").mockReturnValue(1_737_000_000_000); const tool = createImageGenerateTool( - { - writeGeneratedArtifacts, - }, - { - ...defaultImageOptions, - canSendFilesToActiveConversation: true, - }, + { writeGeneratedArtifacts }, + { ...defaultImageOptions, canSendFilesToActiveConversation: true }, + { fetch: gatewayFetch([imageResponse()]) }, ); const result = await tool.execute!({ prompt: "test prompt" }, {} as any); @@ -187,28 +139,26 @@ describe("createImageGenerateTool", () => { it("returns an actionable error when model is not image-capable", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; mockCompleteText.mockResolvedValueOnce({ text: "enriched prompt" } as any); - const fetchMock = vi.fn().mockResolvedValueOnce( - createErrorResponse( - 400, - JSON.stringify({ + const fetchMock = gatewayFetch([ + Response.json( + { error: { + type: "invalid_request_error", message: "Model 'google/gemini-3-pro-image' is a language model, not an image model. Use the language generation API instead.", }, - }), + }, + { status: 400 }, ), - ); - vi.stubGlobal("fetch", fetchMock); + ]); const tool = createImageGenerateTool( - { writeGeneratedArtifacts } as any, - defaultImageOptions, + { writeGeneratedArtifacts }, + { modelId: "google/gemini-3-pro-image", sandbox: sandbox() }, + { fetch: fetchMock }, ); - if (typeof tool.execute !== "function") { - throw new Error("imageGenerate execute function missing"); - } await expect( - tool.execute({ prompt: "person in a forest" }, {} as any), + tool.execute!({ prompt: "person in a forest" }, {} as any), ).rejects.toThrow( 'configured model "google/gemini-3-pro-image" is not an image generation model', ); @@ -219,19 +169,16 @@ describe("createImageGenerateTool", () => { mockCompleteText.mockResolvedValueOnce({ text: "a dark, high-contrast dog with glowing eyes", } as any); - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([imageResponse()]); const tool = createImageGenerateTool( - { writeGeneratedArtifacts } as any, + { writeGeneratedArtifacts }, defaultImageOptions, + { fetch: fetchMock }, ); const result = await tool.execute!({ prompt: "draw a dog" }, {} as any); - const body = getRequestBody(fetchMock); - expect(body.messages[0].content).toBe( + expect(requestBody(fetchMock).prompt).toBe( "a dark, high-contrast dog with glowing eyes", ); expect(result).toMatchObject({ @@ -243,53 +190,38 @@ describe("createImageGenerateTool", () => { it("falls back to raw prompt when enrichment returns empty text", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; mockCompleteText.mockResolvedValueOnce({ text: " " } as any); - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([imageResponse()]); const tool = createImageGenerateTool( - { writeGeneratedArtifacts } as any, + { writeGeneratedArtifacts }, defaultImageOptions, + { fetch: fetchMock }, ); const result = await tool.execute!({ prompt: "draw a dog" }, {} as any); - const body = getRequestBody(fetchMock); - expect(body.messages[0].content).toBe("draw a dog"); - expect(result).toMatchObject({ - prompt: "draw a dog", - enrichedPrompt: "draw a dog", - }); + expect(requestBody(fetchMock).prompt).toBe("draw a dog"); + expect(result).toMatchObject({ enrichedPrompt: "draw a dog" }); }); it("falls back to raw prompt when enrichment fails", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; mockCompleteText.mockRejectedValueOnce(new Error("LLM unavailable")); - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([imageResponse()]); const tool = createImageGenerateTool( - { writeGeneratedArtifacts } as any, + { writeGeneratedArtifacts }, defaultImageOptions, + { fetch: fetchMock }, ); const result = await tool.execute!({ prompt: "draw a dog" }, {} as any); - const body = getRequestBody(fetchMock); - expect(body.messages[0].content).toBe("draw a dog"); - expect(result).toMatchObject({ - prompt: "draw a dog", - enrichedPrompt: "draw a dog", - }); + expect(requestBody(fetchMock).prompt).toBe("draw a dog"); + expect(result).toMatchObject({ enrichedPrompt: "draw a dog" }); }); it("sends input images with the unchanged prompt", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([imageResponse()]); const tool = createImageGenerateTool( { writeGeneratedArtifacts }, @@ -297,6 +229,7 @@ describe("createImageGenerateTool", () => { ...defaultImageOptions, sandbox: sandbox({ "/vercel/sandbox/photo.png": PNG_BYTES }), }, + { fetch: fetchMock }, ); const result = await tool.execute!( { prompt: "change the sign text", images: ["photo.png"] }, @@ -304,19 +237,10 @@ describe("createImageGenerateTool", () => { ); expect(mockCompleteText).not.toHaveBeenCalled(); - expect(getRequestBody(fetchMock).messages).toEqual([ - { - role: "user", - content: [ - { type: "text", text: "change the sign text" }, - { - type: "image_url", - image_url: { - url: `data:image/png;base64,${PNG_BYTES.toString("base64")}`, - }, - }, - ], - }, + const body = requestBody(fetchMock); + expect(body.prompt).toBe("change the sign text"); + expect(body.files).toEqual([ + expect.objectContaining({ data: PNG_BYTES.toString("base64") }), ]); expect(result).toMatchObject({ enrichedPrompt: "change the sign text", @@ -324,12 +248,31 @@ describe("createImageGenerateTool", () => { }); }); + it("returns no images when the model generates none", async () => { + process.env.AI_GATEWAY_API_KEY = "test-key"; + mockCompleteText.mockResolvedValueOnce({ text: "enriched prompt" } as any); + + const tool = createImageGenerateTool( + { writeGeneratedArtifacts }, + defaultImageOptions, + { + fetch: gatewayFetch([ + Response.json({ images: [], isRetryable: false }), + ]), + }, + ); + const result = await tool.execute!({ prompt: "draw a dog" }, {} as any); + + expect(result).toMatchObject({ + image_count: 0, + images: [], + delivery: "No images were generated.", + }); + }); + it("resizes input images over the request size limit", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; - const fetchMock = vi - .fn() - .mockResolvedValueOnce(createJsonResponse(imagePayload())); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([imageResponse()]); // Random pixels do not compress, so this PNG is about 9 MB. const largePng = await sharp({ create: { @@ -355,6 +298,7 @@ describe("createImageGenerateTool", () => { "/vercel/sandbox/badge.png": PNG_BYTES, }), }, + { fetch: fetchMock }, ); const result = await tool.execute!( { @@ -364,21 +308,16 @@ describe("createImageGenerateTool", () => { {} as any, ); - const [, sentPhoto, sentLogo] = - getRequestBody(fetchMock).messages[0].content; - const photoBytes = Buffer.from( - sentPhoto.image_url.url.replace("data:image/jpeg;base64,", ""), - "base64", - ); + const [sentPhoto, sentLogo] = requestBody(fetchMock).files; + expect(sentPhoto.mediaType).toBe("image/jpeg"); + const photoBytes = Buffer.from(sentPhoto.data, "base64"); // The small images leave their unused bytes to the photo, so it keeps // the largest size. An equal share would shrink it further. expect(await sharp(photoBytes).metadata()).toMatchObject({ format: "jpeg", width: 2048, }); - expect(sentLogo.image_url.url).toBe( - `data:image/png;base64,${PNG_BYTES.toString("base64")}`, - ); + expect(sentLogo.data).toBe(PNG_BYTES.toString("base64")); expect(result).toMatchObject({ resized_images: [ { @@ -393,8 +332,7 @@ describe("createImageGenerateTool", () => { it("rejects input images too large to read before calling the gateway", async () => { process.env.AI_GATEWAY_API_KEY = "test-key"; - const fetchMock = vi.fn(); - vi.stubGlobal("fetch", fetchMock); + const fetchMock = gatewayFetch([]); const tool = createImageGenerateTool( { writeGeneratedArtifacts }, @@ -407,6 +345,7 @@ describe("createImageGenerateTool", () => { ]), }), }, + { fetch: fetchMock }, ); await expect(