ai-chatbot/tests/routes/model-classification.test.ts

import { expect, test } from "@playwright/test";
import {
  createAtlasCloudImagePreset,
  inferAtlasCloudMediaKind,
  resolveAtlasCloudDisplayName,
} from "@/lib/ai/atlascloud-models";
import { createAtlasCloudRequestBody } from "@/lib/ai/atlascloud-request";
import {
  is3DModelDescriptor,
  isImageGenerationModelDescriptor,
  isVideoModelDescriptor,
} from "@/lib/ai/model-classification";

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("preserves administrator-defined AtlasCloud display names", () => {
    expect(
      resolveAtlasCloudDisplayName(
        "视觉 Flux Kontext Pro",
        "Flux Kontext Pro",
        "black-forest-labs/flux-kontext-pro"
      )
    ).toBe("视觉 Flux Kontext Pro");
    expect(
      resolveAtlasCloudDisplayName(
        "",
        "Flux Kontext Pro",
        "black-forest-labs/flux-kontext-pro"
      )
    ).toBe("Flux Kontext Pro");
  });

  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,
      },
    ]);
  });
});