diff --git a/packages/ai/test/images.test.ts b/packages/ai/test/images.test.ts new file mode 100644 index 00000000..84aac37a --- /dev/null +++ b/packages/ai/test/images.test.ts @@ -0,0 +1,124 @@ +import { readFileSync } from "node:fs"; +import { dirname, join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { describe, expect, it } from "vitest"; +import { getImageModel } from "../src/image-models.js"; +import { generateImages, images } from "../src/images.js"; +import type { ImageContent, ImagesContext, ImagesModel, ProviderImagesOptions } from "../src/types.js"; + +const __filename = fileURLToPath(import.meta.url); +const __dirname = dirname(__filename); + +type ImagesOptionsWithExtras = ProviderImagesOptions & Record; + +async function basicImageGeneration(model: ImagesModel, options?: ImagesOptionsWithExtras) { + const context: ImagesContext = { + input: [{ type: "text", text: "Generate a simple red circle on a plain white background. No text." }], + }; + + const response = await generateImages(model, context, options); + + expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("stop"); + expect(response.errorMessage).toBeFalsy(); + expect(response.output.some((item) => item.type === "image")).toBe(true); + expect(response.timestamp).toBeGreaterThan(0); +} + +async function handleStreaming(model: ImagesModel, options?: ImagesOptionsWithExtras) { + const context: ImagesContext = { + input: [{ type: "text", text: "Generate a simple blue square on a plain white background. No text." }], + }; + + const s = images(model, context, options); + let started = false; + let imageStarted = false; + let imageCompleted = false; + + for await (const event of s) { + if (event.type === "start") { + started = true; + } else if (event.type === "image_start") { + imageStarted = true; + } else if (event.type === "image_end") { + imageCompleted = true; + expect(event.image.type).toBe("image"); + expect(event.image.data.length).toBeGreaterThan(0); + } + } + + const response = await s.result(); + expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("stop"); + expect(started).toBe(true); + expect(imageStarted).toBe(true); + expect(imageCompleted).toBe(true); + expect(response.output.some((item) => item.type === "image")).toBe(true); +} + +async function handleTextAndImageOutput( + model: ImagesModel, + options?: ImagesOptionsWithExtras, +) { + if (!model.output.includes("text")) { + console.log(`Skipping text+image output test - model ${model.id} doesn't support text output`); + return; + } + + const context: ImagesContext = { + input: [{ type: "text", text: "Generate a red circle and include a brief description of the image." }], + }; + + const response = await generateImages(model, context, options); + + expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("stop"); + expect(response.output.some((item) => item.type === "image")).toBe(true); + expect(response.output.some((item) => item.type === "text" && item.text.trim().length > 0)).toBe(true); +} + +async function handleImageInput(model: ImagesModel, options?: ImagesOptionsWithExtras) { + if (!model.input.includes("image")) { + console.log(`Skipping image input test - model ${model.id} doesn't support image input`); + return; + } + + const imagePath = join(__dirname, "data", "red-circle.png"); + const imageBuffer = readFileSync(imagePath); + const imageContent: ImageContent = { + type: "image", + data: imageBuffer.toString("base64"), + mimeType: "image/png", + }; + + const context: ImagesContext = { + input: [{ type: "text", text: "Create a variation of this image with a blue background." }, imageContent], + }; + + const response = await generateImages(model, context, options); + + expect(response.stopReason, `Error: ${response.errorMessage}`).toBe("stop"); + expect(response.output.some((item) => item.type === "image")).toBe(true); +} + +describe("Images E2E Tests", () => { + describe.skipIf(!process.env.OPENROUTER_API_KEY)( + "OpenRouter Images Provider (google/gemini-2.5-flash-image)", + () => { + const model = getImageModel("openrouter", "google/gemini-2.5-flash-image"); + + it("should generate a basic image", { retry: 3 }, async () => { + await basicImageGeneration(model); + }); + + it("should handle image streaming", { retry: 3 }, async () => { + await handleStreaming(model); + }); + + it("should handle text plus image output", { retry: 3 }, async () => { + await handleTextAndImageOutput(model); + }); + + it("should handle image input", { retry: 3 }, async () => { + await handleImageInput(model); + }); + }, + ); +});