ai-chatbot/scripts/configure-image-provider-allowlist.ts

import dotenv from "dotenv";
import postgres from "postgres";
import {
  fetchVip98kModelPresets,
  type Vip98kModelPreset,
} from "../lib/ai/vip98k-models";
import {
  defaultVip98kRuntimeConfig,
  VIP98K_RUNTIME_CONFIG_KEY,
} from "../lib/ai/vip98k-runtime-config";

dotenv.config({ path: ".env.local" });
dotenv.config();

const ALLOWED_IMAGE_PROVIDER_TYPES = ["alibaba-image", "vip98k"] as const;

function hasFlag(name: string) {
  return process.argv.includes(name);
}

function toJsonValue(value: unknown): postgres.JSONValue {
  return JSON.parse(JSON.stringify(value)) as postgres.JSONValue;
}

function isQualifiedVip98kImagePreset(preset: Vip98kModelPreset) {
  return preset.supportsImageGen && !preset.isHidden;
}

async function main() {
  const apply = hasFlag("--apply");
  const databaseUrl = String(process.env.POSTGRES_URL || "").trim();
  if (!databaseUrl) throw new Error("POSTGRES_URL is required");

  const sql = postgres(databaseUrl, { prepare: false, max: 1 });
  try {
    const providers = await sql<
      Array<{
        id: string;
        type: string;
        name: string;
        apiKey: string | null;
        baseUrl: string | null;
        isEnabled: boolean;
      }>
    >`select id, type, name, "apiKey", "baseUrl", "isEnabled" from "AiProvider" order by "createdAt"`;
    const vip98kProviders = providers.filter((provider) => provider.type === "vip98k");
    const alibabaProviders = providers.filter(
      (provider) => provider.type === "alibaba-image",
    );
    if (vip98kProviders.length !== 1) {
      throw new Error(`Expected one VIP98K provider, found ${vip98kProviders.length}`);
    }
    if (alibabaProviders.length === 0) {
      throw new Error("Alibaba official image provider is missing");
    }
    const vip98kProvider = vip98kProviders[0];
    if (!vip98kProvider) throw new Error("VIP98K provider is missing");
    const vip98kApiKey = String(
      process.env.VIP98K_API_KEY || vip98kProvider.apiKey || "",
    ).trim();
    if (!vip98kApiKey) throw new Error("VIP98K API key is not configured");
    const vip98kBaseUrl = String(
      process.env.VIP98K_BASE_URL ||
        vip98kProvider.baseUrl ||
        "https://api.vip98k.com",
    ).trim();
    const presets = await fetchVip98kModelPresets({
      baseURL: vip98kBaseUrl,
      apiKey: vip98kApiKey,
    });
    const qualifiedVip98kPresets = presets.filter(isQualifiedVip98kImagePreset);
    if (qualifiedVip98kPresets.length === 0) {
      throw new Error("VIP98K catalog contains no qualified image model");
    }
    const existingVip98kModels = await sql<
      Array<{
        id: string;
        modelId: string;
      }>
    >`
      select id, "modelId"
      from "AiModel"
      where "providerId" = ${vip98kProvider.id}::uuid
    `;
    const existingVip98kByModelId = new Map(
      existingVip98kModels.map((model) => [model.modelId, model]),
    );

    const disallowedImageModels = await sql<
      Array<{ id: string; providerType: string; modelId: string }>
    >`
      select m.id, p.type as "providerType", m."modelId"
      from "AiModel" m join "AiProvider" p on p.id = m."providerId"
      where m."supportsImageGen"
        and p.type != all(${ALLOWED_IMAGE_PROVIDER_TYPES}::text[])
      order by p.type, m."modelId"
    `;
    const staleVip98kImageModels = await sql<Array<{ id: string; modelId: string }>>`
      select m.id, m."modelId"
      from "AiModel" m
      where m."providerId" = ${vip98kProvider.id}::uuid
        and m."supportsImageGen"
        and m."modelId" != all(${qualifiedVip98kPresets.map((preset) => preset.modelId)}::text[])
      order by m."modelId"
    `;
    const imageOnlyProviders = await sql<Array<{ id: string; type: string; name: string }>>`
      select p.id, p.type, p.name
      from "AiProvider" p
      where p.type != all(${ALLOWED_IMAGE_PROVIDER_TYPES}::text[])
        and (
          p.type like '%-image'
          or p.type in ('bananaimage', 'gemini', 'huanke-gemini')
          or (
            exists (select 1 from "AiModel" m where m."providerId" = p.id and m."supportsImageGen")
            and not exists (
              select 1 from "AiModel" m
              where m."providerId" = p.id and not m."supportsImageGen" and m."isEnabled"
            )
          )
        )
      order by p.type, p.name
    `;

    console.log(
      JSON.stringify(
        {
          mode: apply ? "apply" : "dry-run",
          allowedImageProviderTypes: ALLOWED_IMAGE_PROVIDER_TYPES,
          qualifiedVip98kModels: qualifiedVip98kPresets.map((preset) => preset.modelId),
          protectedVip98kModels: qualifiedVip98kPresets.map(
            (preset) => preset.modelId,
          ),
          disallowedImageModelsToDisable: disallowedImageModels.length,
          staleVip98kImageModelsToDisable: staleVip98kImageModels.length,
          imageOnlyProvidersToDisable: imageOnlyProviders.map((provider) => ({
            type: provider.type,
            name: provider.name,
          })),
        },
        null,
        2,
      ),
    );
    if (!apply) return;

    await sql.begin(async (tx) => {
      await tx`
        insert into "SystemSettings" (key, value, "updatedAt")
        values (
          ${VIP98K_RUNTIME_CONFIG_KEY},
          ${tx.json(toJsonValue(defaultVip98kRuntimeConfig))},
          now()
        )
        on conflict (key) do nothing
      `;
      await tx`
        update "AiProvider"
        set "baseUrl" = 'https://api.vip98k.com', "apiKey" = ${vip98kApiKey},
            "isEnabled" = true, "updatedAt" = now()
        where id = ${vip98kProvider.id}::uuid
      `;
      await tx`
        update "AiProvider" set "isEnabled" = true, "updatedAt" = now()
        where type = 'alibaba-image'
      `;
      await tx`
        update "AiModel" m set
          "isEnabled" = true, "excludeFromHealthBan" = true,
          "healthStatus" = 'healthy', "healthConsecutiveFailures" = 0,
          "healthLastEventType" = null, "healthLastError" = null,
          "healthAutoDisabledAt" = null, "healthAutoDisabledReason" = null,
          "healthRecoverySuccessStreak" = 0,
          "healthObservationStartedAt" = null,
          "healthUpdatedAt" = now(), "updatedAt" = now()
        from "AiProvider" p
        where p.id = m."providerId" and p.type = 'alibaba-image'
          and m."supportsImageGen" and not m."isHidden"
      `;

      for (const preset of qualifiedVip98kPresets) {
        const capabilities = tx.json(toJsonValue(preset.capabilities));
        const defaultParams = tx.json(toJsonValue(preset.defaultParams));
        const imageGenConfig = preset.imageGenConfig
          ? tx.json(toJsonValue(preset.imageGenConfig))
          : null;
        const existing = existingVip98kByModelId.get(preset.modelId);
        if (existing) {
          await tx`
            update "AiModel" set
              "displayName" = ${preset.displayName}, description = ${preset.description},
              "supportsVision" = ${preset.supportsVision}, "supportsImageGen" = true,
              "isEnabled" = true, "isHidden" = false,
              "excludeFromHealthBan" = true, "healthStatus" = 'healthy',
              "healthConsecutiveFailures" = 0, "healthLastEventType" = null,
              "healthLastError" = null, "healthAutoDisabledAt" = null,
              "healthAutoDisabledReason" = null, "healthRecoverySuccessStreak" = 0,
              "healthObservationStartedAt" = null, "healthUpdatedAt" = now(),
              "sortOrder" = ${preset.sortOrder},
              capabilities = ${capabilities}, "defaultParams" = ${defaultParams},
              "imageGenConfig" = ${imageGenConfig}, "fallbackModelId1" = null,
              "fallbackModelId2" = null, "updatedAt" = now()
            where id = ${existing.id}::uuid
          `;
        } else {
          await tx`
            insert into "AiModel" (
              "providerId", "modelId", "displayName", description, "supportsVision",
              "supportsImageGen", "isEnabled", "isHidden", "sortOrder", capabilities,
              "defaultParams", "imageGenConfig", "isDefault", "isDefaultTextEditModel",
              "excludeFromHealthBan", "healthStatus"
            ) values (
              ${vip98kProvider.id}::uuid, ${preset.modelId}, ${preset.displayName},
              ${preset.description}, ${preset.supportsVision}, true, true, false,
              ${preset.sortOrder}, ${capabilities}, ${defaultParams}, ${imageGenConfig},
              false, false, true, 'healthy'
            )
          `;
        }
      }

      await tx`
        update "AiModel" m set
          "isEnabled" = false, "isHidden" = true, "isDefault" = false,
          "isDefaultTextEditModel" = false, "updatedAt" = now()
        from "AiProvider" p
        where p.id = m."providerId" and m."supportsImageGen"
          and p.type != all(${ALLOWED_IMAGE_PROVIDER_TYPES}::text[])
      `;
      await tx`
        update "AiModel" set
          "isEnabled" = false, "isHidden" = true, "isDefault" = false,
          "isDefaultTextEditModel" = false, "updatedAt" = now()
        where "providerId" = ${vip98kProvider.id}::uuid and "supportsImageGen"
          and "modelId" != all(${qualifiedVip98kPresets.map((preset) => preset.modelId)}::text[])
      `;
      if (imageOnlyProviders.length > 0) {
        await tx`
          update "AiProvider" set "isEnabled" = false, "updatedAt" = now()
          where id = any(${imageOnlyProviders.map((provider) => provider.id)}::uuid[])
        `;
      }
    });

    const invalidModels = await sql<
      Array<{ providerType: string; modelId: string }>
    >`
      select p.type as "providerType", m."modelId"
      from "AiModel" m join "AiProvider" p on p.id = m."providerId"
      where m."supportsImageGen" and m."isEnabled"
        and p.type != all(${ALLOWED_IMAGE_PROVIDER_TYPES}::text[])
    `;
    const activeVip98kModels = await sql<Array<{ modelId: string }>>`
      select m."modelId"
      from "AiModel" m join "AiProvider" p on p.id = m."providerId"
      where p.type = 'vip98k' and p."isEnabled" and m."supportsImageGen"
        and m."isEnabled" and not m."isHidden"
      order by m."modelId"
    `;
    if (invalidModels.length > 0) {
      throw new Error("Disallowed image models remain enabled after apply");
    }
    if (activeVip98kModels.length !== qualifiedVip98kPresets.length) {
      throw new Error("VIP98K active image model count does not match qualified catalog");
    }
    console.log(
      JSON.stringify(
        {
          success: true,
          activeVip98kModels: activeVip98kModels.map((model) => model.modelId),
          disabledImageModelCount: disallowedImageModels.length + staleVip98kImageModels.length,
        },
        null,
        2,
      ),
    );
  } finally {
    await sql.end();
  }
}

main().catch((error) => {
  console.error(error instanceof Error ? error.message : String(error));
  process.exit(1);
});