ai-chatbot/lib/ai/model-classification.ts

import { inferKreaModelKind } from "@/lib/ai/krea-models";

const STATIC_VIDEO_PROVIDER_TYPES = new Set([
  "bltcy-video",
  "doubao-video",
  "volcano-video",
  "qianwen-video",
  "fal-video",
  "ruxa-video",
]);

const STATIC_IMAGE_PROVIDER_TYPES = new Set([
  "gemini",
  "huanke-gemini",
  "alibaba-image",
  "bfl-image",
  "ruxa-image",
  "fal-image",
  "snapgen-image",
  "grsai-gpt-image",
  "openai-image",
  "bananaimage",
  "visionary",
  "atlascloud",
  "stepfun-image",
  "midjourney",
  "volcano-vision",
]);

const RUNWAY_IMAGE_MODEL_IDS = new Set([
  "gen4_image",
  "gen4_image_turbo",
  "gemini_2.5_flash",
]);

const RUNWAY_VIDEO_MODEL_PREFIXES = [
  "gen4.5",
  "gen4_turbo",
  "gen4_aleph",
  "gen3a_turbo",
  "veo3.1_fast",
  "veo3.1",
  "veo3",
  "act_two",
  "upscale_v1",
];

const VIDEO_MODE_KEYS = new Set([
  "text-to-video",
  "first-frame",
  "first-last-frame",
  "reference-images",
]);

const TEXT_TO_IMAGE_MODE_KEYS = new Set([
  "text-to-image",
  "text2image",
  "txt2img",
  "image-generation",
  "image_generate",
  "generate-image",
]);

const IMAGE_EDIT_ONLY_KIND_KEYS = new Set([
  "image-edit",
  "image-to-image",
  "img2img",
  "reference-image",
  "reference-images",
]);

const TEXT_TO_IMAGE_KIND_KEYS = new Set([
  "text-to-image",
  "text2image",
  "txt2img",
  "image",
  "image-generation",
  "image_generate",
]);

function normalizeText(value?: string | null) {
  return String(value || "").trim().toLowerCase();
}

function toRecord(value: unknown) {
  if (!value || typeof value !== "object" || Array.isArray(value)) {
    return null;
  }
  return value as Record<string, unknown>;
}

function getStringArrayField(source: unknown, key: string) {
  const record = toRecord(source);
  const value = record?.[key];
  if (!Array.isArray(value)) return [] as string[];
  return value.map((item) => normalizeText(String(item || ""))).filter(Boolean);
}

function getStringField(source: unknown, key: string) {
  const record = toRecord(source);
  return normalizeText(String(record?.[key] || ""));
}

function hasVideoModes(capabilities: unknown) {
  if (!capabilities || typeof capabilities !== "object") return false;
  const supportedModes = (capabilities as { supportedModes?: unknown }).supportedModes;
  if (!Array.isArray(supportedModes)) return false;
  return supportedModes.some((mode) => VIDEO_MODE_KEYS.has(String(mode || "").trim()));
}

export function inferRunwayModelKind(modelId?: string | null): "image" | "video" | null {
  const normalized = normalizeText(modelId);
  if (!normalized) return null;
  if (RUNWAY_IMAGE_MODEL_IDS.has(normalized)) return "image";
  if (RUNWAY_VIDEO_MODEL_PREFIXES.some((prefix) => normalized === prefix || normalized.startsWith(`${prefix}_`))) {
    return "video";
  }
  return null;
}

export function isStaticVideoProviderType(providerType?: string | null): boolean {
  return STATIC_VIDEO_PROVIDER_TYPES.has(normalizeText(providerType));
}

export function isVideoModelDescriptor(params: {
  providerType?: string | null;
  modelId?: string | null;
  capabilities?: unknown;
  supportsImageGen?: boolean | null;
}): boolean {
  const providerType = normalizeText(params.providerType);

  if (isStaticVideoProviderType(providerType)) {
    return true;
  }

  if (hasVideoModes(params.capabilities)) {
    return true;
  }

  if (providerType === "runway") {
    return inferRunwayModelKind(params.modelId) === "video";
  }

  if (providerType === "krea") {
    return inferKreaModelKind(params.modelId) === "video";
  }

  return false;
}

export function isImageGenerationModelDescriptor(params: {
  providerType?: string | null;
  modelId?: string | null;
  capabilities?: unknown;
  supportsImageGen?: boolean | null;
}): boolean {
  const providerType = normalizeText(params.providerType);

  if (providerType === "runway") {
    const runwayModelKind = inferRunwayModelKind(params.modelId);
    if (runwayModelKind === "image") return true;
    if (runwayModelKind === "video") return false;
  }

  if (providerType === "krea") {
    const kreaModelKind = inferKreaModelKind(params.modelId);
    if (kreaModelKind === "image" || kreaModelKind === "image-enhance") return true;
    if (kreaModelKind === "video") return false;
  }

  if (isVideoModelDescriptor(params)) {
    return false;
  }

  if (params.supportsImageGen === true) {
    return true;
  }

  return STATIC_IMAGE_PROVIDER_TYPES.has(providerType);
}

export function modelSupportsTextToImageDescriptor(params: {
  providerType?: string | null;
  modelId?: string | null;
  capabilities?: unknown;
  defaultParams?: unknown;
  supportsImageGen?: boolean | null;
}): boolean {
  if (
    !isImageGenerationModelDescriptor({
      providerType: params.providerType,
      modelId: params.modelId,
      capabilities: params.capabilities,
      supportsImageGen: params.supportsImageGen,
    })
  ) {
    return false;
  }

  const providerType = normalizeText(params.providerType);
  const modelId = normalizeText(params.modelId);
  const capabilityModes = getStringArrayField(params.capabilities, "supportedModes");
  const defaultModes = getStringArrayField(params.defaultParams, "supportedModes");
  const modes = new Set([...capabilityModes, ...defaultModes]);

  if ([...modes].some((mode) => TEXT_TO_IMAGE_MODE_KEYS.has(mode))) {
    return true;
  }

  const kind =
    getStringField(params.capabilities, "kind") ||
    getStringField(params.defaultParams, "kind");
  if (TEXT_TO_IMAGE_KIND_KEYS.has(kind)) {
    return true;
  }

  if (providerType === "stepfun-image") {
    return true;
  }

  const runtimeEndpoint =
    getStringField(params.capabilities, "endpoint") ||
    getStringField(params.defaultParams, "endpoint");
  const runtimeModelText = `${modelId} ${runtimeEndpoint}`.trim();

  if (
    IMAGE_EDIT_ONLY_KIND_KEYS.has(kind) ||
    [...modes].some((mode) => IMAGE_EDIT_ONLY_KIND_KEYS.has(mode)) ||
    runtimeModelText.endsWith("/edit") ||
    runtimeModelText.includes("image-edit")
  ) {
    return false;
  }

  if (providerType === "alibaba-image") {
    return false;
  }

  return true;
}