ai-chatbot/lib/style-pro/model-options.ts

import "server-only";

import type { DiscountableImageGenConfigLike } from "@/lib/ai/image-model-pricing";
import {
  getImageCreditCostFromConfig,
  getImageModelCreditDiscountLabel,
} from "@/lib/ai/image-model-pricing";
import {
  getAiModelFrontendAvailability,
  type ModelFrontendAvailabilityStatus,
} from "@/lib/ai/model-frontend-availability";
import {
  getEnabledAiModelsWithProvider,
  getFrontendDisplayAiModelsWithProvider,
} from "@/lib/db/queries";

export type StyleProModelEntry = Awaited<
  ReturnType<typeof getEnabledAiModelsWithProvider>
>[number];
export type StyleProDisplayModelEntry = Awaited<
  ReturnType<typeof getFrontendDisplayAiModelsWithProvider>
>[number];

export type PublicStyleProModelOption = {
  id: string;
  modelId: string;
  displayName: string;
  avatarUrl: string | null;
  providerName: string;
  providerType: string;
  isDefault: boolean;
  unitCost: number;
  creditDiscountLabel: string | null;
  availabilityStatus: ModelFrontendAvailabilityStatus;
  availabilityLabel: string | null;
  selectable: boolean;
  disabledReason: string | null;
};

const VIDEO_PROVIDER_TYPES = new Set([
  "grsai-video",
  "qianwen-video",
  "doubao-video",
  "ruxa-video",
  "volcano-video",
  "fal-video",
  "runway",
]);

const VIDEO_MODEL_KEYWORDS = [
  "/generate/video/",
  "video",
  "veo",
  "sora",
  "kling",
  "hailuo",
  "seedance",
  "runway",
  "gen-3",
  "gen-4",
  "gen 3",
  "gen 4",
  "wan-2.",
  "wan 2.",
  "t2v",
  "i2v",
];

const EXPLICIT_IMAGE_MODEL_KEYWORDS = [
  "/generate/image/",
  "gpt-image",
  "sora_image",
  "qwen-image",
  "seedream",
  "nano-banana",
  "banana",
  "flux",
  "z-image",
  "image-edit",
  "gen4_image",
  "gen-4 image",
];

function hasVideoGenerationConfig(value: unknown) {
  if (!value || typeof value !== "object") return false;
  const config = value as Record<string, unknown>;
  return (
    "creditCostPerVideo" in config ||
    "creditCostPerVideoByDuration" in config ||
    "creditCostPerVideoByResolution" in config ||
    "videoDurationOptions" in config ||
    "videoResolutionOptions" in config
  );
}

function hasVideoCapabilities(value: unknown) {
  if (!value || typeof value !== "object") return false;
  const serialized = JSON.stringify(value).toLowerCase();
  return (
    serialized.includes("text-to-video") ||
    serialized.includes("image-to-video")
  );
}

function isExplicitImageGenerationModel(
  entry: StyleProModelEntry | StyleProDisplayModelEntry
) {
  const searchable = [
    entry.model.modelId,
    entry.model.displayName,
    entry.model.description,
  ]
    .filter(Boolean)
    .join(" ")
    .toLowerCase();

  if (
    searchable.includes("image-to-video") ||
    searchable.includes("text-to-video") ||
    searchable.includes("i2v") ||
    searchable.includes("t2v")
  ) {
    return false;
  }

  return EXPLICIT_IMAGE_MODEL_KEYWORDS.some((keyword) =>
    searchable.includes(keyword)
  );
}

function isVideoLikeModel(
  entry: StyleProModelEntry | StyleProDisplayModelEntry
) {
  if (VIDEO_PROVIDER_TYPES.has(entry.provider.type)) return true;
  if (
    hasVideoGenerationConfig(entry.model.imageGenConfig) ||
    hasVideoCapabilities(entry.model.capabilities) ||
    hasVideoCapabilities(entry.model.defaultParams)
  ) {
    return true;
  }
  const searchable = [
    entry.provider.type,
    entry.provider.name,
    entry.model.modelId,
    entry.model.displayName,
    entry.model.description,
  ]
    .filter(Boolean)
    .join(" ")
    .toLowerCase();
  return VIDEO_MODEL_KEYWORDS.some((keyword) => searchable.includes(keyword));
}

export function isStyleProImageModel(
  entry: StyleProModelEntry | StyleProDisplayModelEntry
) {
  if (!entry.model.supportsImageGen) return false;
  if (isExplicitImageGenerationModel(entry)) return true;
  return !isVideoLikeModel(entry);
}

export async function getStyleProImageModelEntries(options?: {
  includeHidden?: boolean;
}) {
  const models = await getEnabledAiModelsWithProvider(
    false,
    options?.includeHidden === true
  );
  return models.filter(isStyleProImageModel);
}

export async function getStyleProImageModelDisplayEntries(options?: {
  includeHidden?: boolean;
}) {
  const models = await getFrontendDisplayAiModelsWithProvider(
    false,
    options?.includeHidden === true
  );
  return models.filter(isStyleProImageModel);
}

export function findStyleProImageModelByRef(
  entries: Array<StyleProModelEntry | StyleProDisplayModelEntry>,
  modelRef?: string | null
) {
  const normalized = String(modelRef || "").trim();
  if (!normalized) return null;
  return (
    entries.find(
      ({ model }) => model.id === normalized || model.modelId === normalized
    ) || null
  );
}

export function toPublicStyleProModelOption(
  entry: StyleProModelEntry | StyleProDisplayModelEntry
): PublicStyleProModelOption {
  const imageGenConfig = (entry.model.imageGenConfig ||
    null) as DiscountableImageGenConfigLike | null;
  const availability = getAiModelFrontendAvailability(entry.model);
  return {
    id: entry.model.id,
    modelId: entry.model.modelId,
    displayName: entry.model.displayName,
    avatarUrl: entry.model.avatarUrl || null,
    providerName: entry.provider.name,
    providerType: entry.provider.type,
    isDefault: entry.model.isDefault,
    unitCost: getImageCreditCostFromConfig(imageGenConfig),
    creditDiscountLabel: getImageModelCreditDiscountLabel(imageGenConfig),
    ...availability,
  };
}