import { expect, test } from "@playwright/test"; import { createAtlasCloudImagePreset, inferAtlasCloudMediaKind, } from "@/lib/ai/atlascloud-models"; import { createAtlasCloudRequestBody } from "@/lib/ai/atlascloud-request"; import { is3DModelDescriptor, isImageGenerationModelDescriptor, isVideoModelDescriptor, } from "@/lib/ai/model-classification"; import { buildVip98kModelPresets } from "@/lib/ai/vip98k-models"; const atlasHunyuan3D = { model: "tencent/hunyuan3d-pro/text-to-3d", type: "Image", displayName: "Hunyuan 3D Pro Text-to-3D", categories: ["TEXT-TO-3D"], tags: ["TEXT-TO-3D", "NEW"], }; test.describe("AI model media classification", () => { test("keeps multi-reference BananaImage models classified as image models", () => { const descriptor = { providerType: "bananaimage", modelId: "gemini-3-pro-image-preview", capabilities: { supportedModes: [ "text-to-image", "image-to-image", "reference-images", ], }, supportsImageGen: true, }; expect(isVideoModelDescriptor(descriptor)).toBe(false); expect(isImageGenerationModelDescriptor(descriptor)).toBe(true); }); test("still recognizes explicit video generation modes", () => { expect( isVideoModelDescriptor({ providerType: "custom", modelId: "custom-video-model", capabilities: { supportedModes: ["text-to-video"] }, }) ).toBe(true); }); test("keeps Atlas 3D entries out of image and video categories", () => { expect(inferAtlasCloudMediaKind(atlasHunyuan3D)).toBe("3d"); const descriptor = { providerType: "atlascloud", modelId: atlasHunyuan3D.model, capabilities: { kind: "3d", mediaType: "3d", supportedModes: ["text-to-3d"], }, supportsImageGen: false, }; expect(is3DModelDescriptor(descriptor)).toBe(true); expect(isImageGenerationModelDescriptor(descriptor)).toBe(false); expect(isVideoModelDescriptor(descriptor)).toBe(false); }); test("writes 3D Atlas presets with model-level capabilities", () => { const preset = createAtlasCloudImagePreset( atlasHunyuan3D, { components: { schemas: { Input: { type: "object", properties: { model: { type: "string" }, prompt: { type: "string" }, }, required: ["model", "prompt"], }, }, }, }, 0 ); expect(preset.supportsImageGen).toBe(false); expect(preset.capabilities).toMatchObject({ kind: "3d", mediaType: "3d", supportedModes: ["text-to-3d"], }); }); test("still distinguishes Atlas image and video descriptors", () => { const imageDescriptor = { providerType: "atlascloud", modelId: "black-forest-labs/flux-2-pro/text-to-image", capabilities: { kind: "text-to-image", mediaType: "image", supportedModes: ["text-to-image"], }, supportsImageGen: true, }; const videoDescriptor = { providerType: "atlascloud", modelId: "veed/fabric-1.0/image-to-video", capabilities: { kind: "video", mediaType: "video", supportedModes: ["image-to-video"], }, supportsImageGen: false, }; expect(isImageGenerationModelDescriptor(imageDescriptor)).toBe(true); expect(isVideoModelDescriptor(imageDescriptor)).toBe(false); expect(isVideoModelDescriptor(videoDescriptor)).toBe(true); expect(isImageGenerationModelDescriptor(videoDescriptor)).toBe(false); }); test("maps Atlas reference-video image models to a valid frontend request", () => { const preset = createAtlasCloudImagePreset( { model: "google/nano-banana-2/reference-to-image", type: "Image", displayName: "Nano Banana 2 Reference to Image", categories: ["IMAGE-TO-IMAGE"], }, { paths: { "/api/v1/model/generateImage": { post: {} }, "/api/v1/model/prediction/{request_id}": { get: {} }, }, components: { schemas: { Input: { type: "object", properties: { model: { type: "string" }, prompt: { type: "string" }, resolution: { type: "string", enum: ["1k", "2k", "4k"] }, video_clips: { type: "array", minItems: 1, maxItems: 1, items: { type: "object", properties: { url: { type: "string" }, start: { type: "number", default: 0 }, ends: { type: "number", default: 0 }, fps: { type: "number", default: 1 }, }, }, }, }, required: ["model", "prompt", "video_clips"], }, }, }, } ); expect(preset.capabilities).toMatchObject({ mediaType: "image", supportsVideoReference: true, requiresVideoReference: true, maxVideoReferences: 1, }); const { body } = createAtlasCloudRequestBody({ modelId: preset.modelId, prompt: "参考视频生成商品图", referenceImages: [], referenceVideos: ["https://example.com/reference.mp4"], imageParams: { imageSize: "2K", aspectRatio: "1:1" }, modelConfig: preset.imageGenConfig, }); expect(body.video_clips).toEqual([ { url: "https://example.com/reference.mp4", start: 0, ends: 0, fps: 1, }, ]); }); test("normalizes VIP98K media and tool types before they reach the frontend", () => { const presets = buildVip98kModelPresets([ { modelId: "mislabelled-video", displayName: "Mislabelled video", type: "image", capabilities: ["text_to_video"], }, { modelId: "enhance", displayName: "Enhance", type: "image", capabilities: ["image_processing"], }, { modelId: "image-expansion", displayName: "Extend", type: "image_tool", capabilities: ["outpainting"], }, { modelId: "remove-background", displayName: "Remove background", type: "image_tool", capabilities: ["image_processing"], }, { modelId: "multimodal-model", displayName: "Multimodal model", type: "model", capabilities: ["multimodal"], inputModes: ["text", "image"], }, { modelId: "tencent/hunyuan3d-pro/text-to-3d", displayName: "Hunyuan 3D", type: "image", capabilities: ["text-to-3d"], }, ]); expect(presets.find((item) => item.modelId === "mislabelled-video")).toMatchObject({ type: "video", supportsImageGen: false, }); expect(presets.find((item) => item.modelId === "enhance")).toMatchObject({ type: "image_tool", isHidden: true, capabilities: { toolOperation: "enhance" }, }); expect(presets.find((item) => item.modelId === "image-expansion")).toMatchObject({ type: "image_tool", capabilities: { toolOperation: "extend" }, }); expect(presets.find((item) => item.modelId === "remove-background")).toMatchObject({ type: "image_tool", capabilities: { toolOperation: "remove-background" }, }); expect(presets.find((item) => item.modelId === "multimodal-model")).toMatchObject({ type: "model", supportsVision: true, supportsImageGen: false, }); expect( presets.find((item) => item.modelId === "tencent/hunyuan3d-pro/text-to-3d"), ).toMatchObject({ type: "model_3d", supportsImageGen: false, isHidden: true, capabilities: { kind: "model-3d" }, }); }); test("only shows VIP98K image models verified as released in 2026 or later", () => { const presets = buildVip98kModelPresets([ { modelId: "black-forest-labs/flux-dev", displayName: "Flux Dev", type: "image", capabilities: ["text_to_image"], }, { modelId: "google/nano-banana-2/text-to-image", displayName: "Nano Banana 2", type: "image", capabilities: ["text_to_image"], }, { modelId: "future-image-with-date", displayName: "Future image", type: "image", capabilities: ["text_to_image"], releasedAt: "2027-02-01", }, { modelId: "old-image-with-date", displayName: "Old image", type: "image", capabilities: ["text_to_image"], releasedAt: "2025-12-15", }, ]); expect(presets.find((item) => item.modelId === "black-forest-labs/flux-dev")?.isHidden).toBe(true); expect( presets.find((item) => item.modelId === "google/nano-banana-2/text-to-image")?.isHidden, ).toBe(false); expect(presets.find((item) => item.modelId === "future-image-with-date")?.isHidden).toBe(false); expect(presets.find((item) => item.modelId === "old-image-with-date")?.isHidden).toBe(true); }); });