ai-chatbot/scripts/verify-ruxa-image-models.ts

import { config as loadEnv } from "dotenv";
import { and, eq, sql } from "drizzle-orm";

import { db } from "../lib/db";
import { aiModel, aiProvider } from "../lib/db/schema";
import { RUXA_NANO_BANANA_2_RUNTIME_MODEL_ID } from "../lib/ai/ruxa-image-models";

loadEnv({ path: ".env.local" });

type ModelExpectation = {
  modelId: string;
  displayName: string;
  expectedRuntimeModelId?: string;
  expectedTaskMode?: "size-image" | "resolution-based";
};

const EXPECTED_MODELS: ModelExpectation[] = [
  {
    modelId: "google/nano-banana-2",
    displayName: "nano-banana-2「专线」",
    expectedRuntimeModelId: RUXA_NANO_BANANA_2_RUNTIME_MODEL_ID,
    expectedTaskMode: "resolution-based",
  },
  { modelId: "google/nano-banana-pro", displayName: "nano-banana-pro「专线」", expectedTaskMode: "resolution-based" },
];

type Row = {
  providerId: string;
  providerName: string;
  providerType: string;
  modelId: string;
  displayName: string;
  supportsImageGen: boolean | null;
  supportsVision: boolean | null;
  isEnabled: boolean | null;
  isHidden: boolean | null;
  healthStatus: string | null;
  imageGenConfig: unknown;
  defaultParams: unknown;
};

function getConfigFlag(
  config: Record<string, unknown> | null,
  key: string
): boolean | null {
  const raw = config?.[key];
  return typeof raw === "boolean" ? raw : null;
}

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

async function main() {
  const rows = await db
    .select({
      providerId: aiProvider.id,
      providerName: aiProvider.name,
      providerType: aiProvider.type,
      modelId: aiModel.modelId,
      displayName: aiModel.displayName,
      supportsImageGen: aiModel.supportsImageGen,
      supportsVision: aiModel.supportsVision,
      isEnabled: aiModel.isEnabled,
      isHidden: aiModel.isHidden,
      defaultParams: aiModel.defaultParams,
      healthStatus: aiModel.healthStatus,
      imageGenConfig: aiModel.imageGenConfig,
    })
    .from(aiModel)
    .innerJoin(aiProvider, eq(aiModel.providerId, aiProvider.id))
    .where(
      and(
        eq(aiProvider.type, "ruxa-image"),
        sql`LOWER(${aiProvider.name}) LIKE ${"%ruxa%"}`,
        sql`LOWER(${aiModel.modelId}) IN (${"google/nano-banana-2"}, ${"google/nano-banana-pro"})`
      )
    );

  const found = new Map<string, Row[]>();
  for (const row of rows) {
    const bucket = found.get(row.providerId) ?? [];
    bucket.push(row);
    found.set(row.providerId, bucket);
  }

  const providers = [...found.entries()].map(([providerId, models]) => ({
    providerId,
    providerName: models[0]?.providerName ?? "",
    providerType: models[0]?.providerType ?? "",
    modelCount: models.length,
    models,
  }));

  const failures: string[] = [];
  const expectByModelId = new Map(
    EXPECTED_MODELS.map((entry) => [entry.modelId, entry])
  );
  const expectedSet = new Set(expectByModelId.keys());

  for (const row of rows) {
    if (!expectedSet.has(row.modelId)) {
      failures.push(`Unexpected model present in RUXA image provider: ${row.modelId}`);
      continue;
    }

    const expected = expectByModelId.get(row.modelId);
    if (!expected) {
      failures.push(`Unexpected model present in RUXA image provider: ${row.modelId}`);
      continue;
    }

    if (row.displayName !== expected.displayName) {
      failures.push(
        `${row.providerName} ${row.modelId}: displayName mismatch (expected "${expected.displayName}", got "${row.displayName}")`
      );
    }

    const config = getImageGenConfig(row.imageGenConfig);

    if (row.supportsImageGen !== true) {
      failures.push(`${row.providerName} ${row.modelId}: supportsImageGen is not true`);
    }
    if (row.supportsVision !== true) {
      failures.push(`${row.providerName} ${row.modelId}: supportsVision is not true`);
    }
    if (row.isEnabled !== true) {
      failures.push(`${row.providerName} ${row.modelId}: isEnabled is not true`);
    }
    if (row.isHidden === true) {
      failures.push(`${row.providerName} ${row.modelId}: isHidden should be false`);
    }
    if (row.healthStatus !== "healthy") {
      failures.push(
        `${row.providerName} ${row.modelId}: healthStatus is ${row.healthStatus ?? "null"}, expected healthy`
      );
    }

    if (getConfigFlag(config, "supportsSize") !== true) {
      failures.push(`${row.providerName} ${row.modelId}: imageGenConfig.supportsSize should be true`);
    }
    if (getConfigFlag(config, "supports1K") !== true) {
      failures.push(
        `${row.providerName} ${row.modelId}: imageGenConfig.supports1K should be true`
      );
    }
    if (getConfigFlag(config, "supports2K") !== true) {
      failures.push(
        `${row.providerName} ${row.modelId}: imageGenConfig.supports2K should be true`
      );
    }
    if (getConfigFlag(config, "supports4K") !== true) {
      failures.push(
        `${row.providerName} ${row.modelId}: imageGenConfig.supports4K should be true`
      );
    }

    const defaultParams = getImageGenConfig(row.defaultParams) ?? {};
    const runtimeModelFromDb = getRuntimeModelId(row, defaultParams);
    const expectedRuntimeModelId = expected.expectedRuntimeModelId;
    if (expectedRuntimeModelId && runtimeModelFromDb !== expectedRuntimeModelId) {
      failures.push(
        `${row.providerName} ${row.modelId}: default runtime model should be ${expectedRuntimeModelId}, got ${runtimeModelFromDb || "null"}`
      );
    }

    const runtimeMode = config?.ruxaTaskInputMode;
    if (runtimeMode !== expected.expectedTaskMode) {
      failures.push(
        `${row.providerName} ${row.modelId}: imageGenConfig.ruxaTaskInputMode should be ${expected.expectedTaskMode}, got ${runtimeMode}`
      );
    }
  }

  for (const provider of providers) {
    const providerModels = provider.models.map((item) => item.modelId);
    for (const expected of EXPECTED_MODELS) {
      if (!providerModels.includes(expected.modelId)) {
        failures.push(
          `${provider.providerName}: missing required model ${expected.modelId} (${expected.displayName})`
        );
      }
    }
  }

  const providerSummary = providers
    .sort((a, b) => a.providerName.localeCompare(b.providerName))
    .map((provider) => {
      return {
        providerName: provider.providerName,
        providerType: provider.providerType,
        providerId: provider.providerId,
        hasAllTargets: EXPECTED_MODELS.every((expected) =>
          provider.models.some((row) => row.modelId === expected.modelId)
        ),
        modelRows: provider.models
          .sort((a, b) => a.modelId.localeCompare(b.modelId))
          .map((row) => {
            const config = getImageGenConfig(row.imageGenConfig);
            return {
              modelId: row.modelId,
              displayName: row.displayName,
              isEnabled: row.isEnabled,
              healthStatus: row.healthStatus,
              imageSize4K: getConfigFlag(config, "supports4K"),
            };
          }),
      };
    });

  console.log(
    JSON.stringify(
      {
        expected: EXPECTED_MODELS,
        providerSummary,
      },
      null,
      2
    )
  );

  if (failures.length > 0) {
    console.error("RUXA image model check failed:");
    failures.forEach((item) => console.error(`- ${item}`));
    process.exit(1);
  }

  console.log("RUXA image model check passed.");
}

function getRuntimeModelId(row: Row, defaultParams: Record<string, unknown>) {
  const endpoint = getConfigValue(defaultParams, "endpoint");
  if (endpoint) return endpoint;
  return row.modelId;
}

function getConfigValue(config: Record<string, unknown>, key: string): string | null {
  const raw = config?.[key];
  return typeof raw === "string" && raw.trim() ? raw.trim() : null;
}

main().catch((error) => {
  console.error("RUXA image model verification failed to execute:", error);
  process.exit(1);
});