Skip to content
Merged
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: 1 addition & 1 deletion packages/docs/src/content/docs/reference/config-and-env.md
Original file line number Diff line number Diff line change
Expand Up @@ -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",
});
```
Expand Down
36 changes: 10 additions & 26 deletions packages/junior-evals/src/fixture/gateway.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,32 +16,24 @@ 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<boolean> {
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
* agent can check the image, for example after an edit. A 1x1 PNG does not
* 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;
});
}

Expand Down Expand Up @@ -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 (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.");
2 changes: 1 addition & 1 deletion packages/junior/src/chat/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
144 changes: 41 additions & 103 deletions packages/junior/src/chat/tools/web/image-generate.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";

Expand All @@ -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({
Expand Down Expand Up @@ -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,
Expand All @@ -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,
};
Expand Down Expand Up @@ -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<Pick<ToolHooks, "writeGeneratedArtifacts">>,
Expand Down Expand Up @@ -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) => ({
Expand All @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion packages/junior/tests/component/config/chat-config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down
Loading
Loading