ai-chatbot/scripts/test-vip98k-image-models.ts

import { createHash } from "node:crypto";
import { mkdir, writeFile } from "node:fs/promises";
import { dirname, resolve } from "node:path";
import dotenv from "dotenv";
import sharp from "sharp";

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

const BASE_URL = String(process.env.VIP98K_BASE_URL || "https://api.vip98k.com")
  .trim()
  .replace(/\/+$/, "")
  .replace(/\/v1$/i, "");
const DEFAULT_REFERENCE_URLS = [
  "https://dummyimage.com/512x512/2563eb/ffffff.png&text=BLUE+A",
  "https://dummyimage.com/512x512/dc2626/ffffff.png&text=RED+B",
];
const REFERENCE_URLS = (() => {
  const configured = String(process.env.VIP98K_IMAGE_TEST_REFERENCES || "").trim();
  if (!configured) return DEFAULT_REFERENCE_URLS;
  const urls = configured
    .split(/\r?\n/)
    .map((url) => url.trim())
    .filter((url) => /^https?:\/\//i.test(url));
  if (urls.length < 2) {
    throw new Error("VIP98K_IMAGE_TEST_REFERENCES must contain at least two URLs");
  }
  return urls;
})();
const USES_CUSTOM_REFERENCES = Boolean(
  String(process.env.VIP98K_IMAGE_TEST_REFERENCES || "").trim(),
);
const MODEL_FILTER = new Set(
  String(process.env.VIP98K_IMAGE_TEST_MODELS || "")
    .split(",")
    .map((modelId) => modelId.trim())
    .filter(Boolean),
);
const TEST_PROMPT =
  String(process.env.VIP98K_IMAGE_TEST_PROMPT || "").trim() ||
  "Use BOTH references. Place the blue A card on the left and the red B card on the right. " +
    "Preserve both labels and colors in a landscape 16:9 composition at 1024x576.";
const TERMINAL_STATUSES = new Set(["succeeded", "success", "completed", "done"]);
const FAILURE_STATUSES = new Set(["failed", "error", "timeout", "canceled", "cancelled"]);

type JsonRecord = Record<string, any>;
type Invocation = {
  method?: string;
  path?: string;
  content_type?: string;
  request_example?: JsonRecord;
  result_path?: string;
};

function asRecord(value: unknown): JsonRecord {
  return value && typeof value === "object" && !Array.isArray(value)
    ? (value as JsonRecord)
    : {};
}

async function readJson(response: Response) {
  const text = await response.text();
  let payload: unknown = text;
  try {
    payload = text ? JSON.parse(text) : null;
  } catch {
    // Keep the raw provider body for diagnostics.
  }
  if (!response.ok) {
    const record = asRecord(payload);
    const message = record.error?.message || record.error || record.message || text;
    throw new Error(`${String(message || "VIP98K request failed")} (HTTP ${response.status})`);
  }
  return payload;
}

function requestKeys(invocation: Invocation) {
  return Object.keys(asRecord(invocation.request_example)).map((key) => key.toLowerCase());
}

function supportsMultipleImages(invocation: Invocation) {
  return requestKeys(invocation).some((key) =>
    ["file_urls", "images", "image_urls", "reference_images"].includes(key),
  );
}

function extractStatus(payload: unknown) {
  const root = asRecord(payload);
  const data = asRecord(root.data || root);
  return String(data.status || data.state || root.status || "").toLowerCase();
}

function extractTaskId(payload: unknown) {
  const root = asRecord(payload);
  const data = asRecord(root.data || root);
  return String(data.taskId || data.task_id || data.id || root.id || "");
}

async function fetchCatalog(apiKey: string) {
  const [catalogPayload, authorizedPayload] = await Promise.all([
    readJson(
      await fetch(`${BASE_URL}/api/v2/catalog/models`, {
        signal: AbortSignal.timeout(60_000),
      }),
    ),
    readJson(
      await fetch(`${BASE_URL}/v1/models`, {
        headers: { Authorization: `Bearer ${apiKey}` },
        signal: AbortSignal.timeout(60_000),
      }),
    ),
  ]);
  const catalog = asRecord(catalogPayload);
  const authorized = asRecord(authorizedPayload);
  const authorizedIds = new Set(
    (Array.isArray(authorized.data) ? authorized.data : []).map((model) =>
      String(asRecord(model).id || ""),
    ),
  );
  const sections = Array.isArray(asRecord(catalog.data).sections)
    ? asRecord(catalog.data).sections
    : [];
  return sections
    .flatMap((section: unknown) => {
      const record = asRecord(section);
      return Array.isArray(record.models) ? record.models : [];
    })
    .map(asRecord)
    .filter(
      (model: JsonRecord) => {
        const capabilities = (Array.isArray(model.capabilities) ? model.capabilities : []).map(
          (capability: unknown) => String(capability).toLowerCase(),
        );
        return (
          authorizedIds.has(String(model.id || "")) &&
          String(model.type || model.category || "").toLowerCase() === "image" &&
          capabilities.some((capability: string) =>
            ["text_to_image", "text-to-image", "openai.images.generations"].includes(
              capability,
            ),
          )
        );
      },
    );
}

async function resolveOutputUrl(apiKey: string, taskId: string) {
  const payload = await readJson(
    await fetch(
      `${BASE_URL}/api/external/task/result?taskId=${encodeURIComponent(taskId)}`,
      {
        headers: { Authorization: `Bearer ${apiKey}` },
        signal: AbortSignal.timeout(60_000),
      },
    ),
  );
  const data = asRecord(asRecord(payload).data);
  const outputs = asRecord(data.outputs);
  const generatedImages = Array.isArray(outputs.generated_image)
    ? outputs.generated_image
    : [];
  const first = asRecord(generatedImages[0]);
  const url = String(first.image_url || first.file_download_url || "");
  if (!url.startsWith("http")) throw new Error("VIP98K delivery result has no image URL");
  return url;
}

async function pollTask(apiKey: string, taskId: string, resultPath: string) {
  const startedAt = Date.now();
  while (Date.now() - startedAt < 15 * 60_000) {
    await new Promise((resolvePromise) => setTimeout(resolvePromise, 5000));
    const payload = await readJson(
      await fetch(`${BASE_URL}${resultPath}`, {
        headers: { Authorization: `Bearer ${apiKey}` },
        signal: AbortSignal.timeout(60_000),
      }),
    );
    const status = extractStatus(payload);
    if (FAILURE_STATUSES.has(status)) throw new Error(`VIP98K task failed: ${status}`);
    if (TERMINAL_STATUSES.has(status)) {
      return {
        durationMs: Date.now() - startedAt,
        outputUrl: await resolveOutputUrl(apiKey, taskId),
      };
    }
  }
  throw new Error("VIP98K task did not finish within 15 minutes");
}

async function analyzeImage(url: string, modelId: string) {
  const response = await fetch(url, { signal: AbortSignal.timeout(120_000) });
  if (!response.ok) throw new Error(`Generated image download failed (HTTP ${response.status})`);
  const buffer = Buffer.from(await response.arrayBuffer());
  const image = sharp(buffer).removeAlpha();
  const metadata = await image.metadata();
  const outputPath = resolve(
    ".local",
    `vip98k-${modelId.replace(/[^a-z0-9._-]+/gi, "-")}-latest.${
      metadata.format === "jpeg" ? "jpg" : metadata.format || "png"
    }`,
  );
  await mkdir(dirname(outputPath), { recursive: true });
  await writeFile(outputPath, buffer);
  const { data, info } = await image.resize(256, 256, { fit: "fill" }).raw().toBuffer({
    resolveWithObject: true,
  });
  let bluePixels = 0;
  let redPixels = 0;
  for (let offset = 0; offset < data.length; offset += info.channels) {
    const red = data[offset] || 0;
    const green = data[offset + 1] || 0;
    const blue = data[offset + 2] || 0;
    if (blue > 100 && blue > red * 1.25 && blue > green * 1.1) bluePixels += 1;
    if (red > 100 && red > blue * 1.25 && red > green * 1.15) redPixels += 1;
  }
  const pixels = info.width * info.height;
  return {
    width: metadata.width || 0,
    height: metadata.height || 0,
    format: metadata.format || "unknown",
    blueShare: Number((bluePixels / pixels).toFixed(4)),
    redShare: Number((redPixels / pixels).toFixed(4)),
    bothReferencesDetected: bluePixels / pixels > 0.03 && redPixels / pixels > 0.03,
    outputPath,
    sha256: createHash("sha256").update(buffer).digest("hex"),
  };
}

async function testModel(apiKey: string, model: JsonRecord) {
  const modelId = String(model.id || "");
  const invocations = (Array.isArray(model.invocations)
    ? model.invocations
    : [model.invocation]
  ).map(asRecord) as Invocation[];
  const invocation = invocations.find(supportsMultipleImages);
  if (!invocation) {
    return {
      modelId,
      status: "unsupported",
      eligibleForMultiImageEditing: false,
      reason: "The published V2 contract exposes text-to-image only",
    };
  }

  const keys = requestKeys(invocation);
  const nativeAspectRatio = keys.some((key) => ["ratio", "aspect_ratio"].includes(key));
  const nativeSize = keys.some((key) =>
    ["size", "image_size", "resolution", "width", "height"].includes(key),
  );
  const body = { ...asRecord(invocation.request_example) };
  body.model = modelId;
  body.file_urls = REFERENCE_URLS;
  body.prompt = TEST_PROMPT;
  const submitted = await readJson(
    await fetch(`${BASE_URL}${invocation.path}`, {
      method: String(invocation.method || "POST").toUpperCase(),
      headers: {
        Authorization: `Bearer ${apiKey}`,
        "Content-Type": invocation.content_type || "application/json",
      },
      body: JSON.stringify(body),
      signal: AbortSignal.timeout(120_000),
    }),
  );
  const taskId = extractTaskId(submitted);
  if (!taskId) throw new Error(`${modelId} returned no task ID`);
  const resultPath = String(invocation.result_path || "").replace(
    "{task_id}",
    encodeURIComponent(taskId),
  );
  if (!resultPath) throw new Error(`${modelId} publishes no result path`);
  const completed = await pollTask(apiKey, taskId, resultPath);
  const image = await analyzeImage(completed.outputUrl, modelId);
  const passed = USES_CUSTOM_REFERENCES || image.bothReferencesDetected;
  return {
    modelId,
    status: passed ? "passed" : "failed",
    eligibleForMultiImageEditing: passed,
    referenceVerification: USES_CUSTOM_REFERENCES ? "manual" : "color-marker",
    publishedRequestKeys: keys.sort(),
    requested: { referenceImageCount: 2, aspectRatio: "16:9", size: "1024x576" },
    nativeParameterSupport: { aspectRatio: nativeAspectRatio, size: nativeSize },
    actual: image,
    durationMs: completed.durationMs,
  };
}

async function main() {
  const apiKey = String(process.env.VIP98K_API_KEY || "").trim();
  if (!apiKey) throw new Error("VIP98K_API_KEY is required");
  const authorizedModels = await fetchCatalog(apiKey);
  const models =
    MODEL_FILTER.size > 0
      ? authorizedModels.filter((model: JsonRecord) =>
          MODEL_FILTER.has(String(model.id || "")),
        )
      : authorizedModels;
  if (models.length === 0) throw new Error("No authorized VIP98K image models matched");
  const results: unknown[] = [];
  for (const model of models) {
    try {
      results.push(await testModel(apiKey, model));
    } catch (error) {
      results.push({
        modelId: String(model.id || ""),
        status: "failed",
        eligibleForMultiImageEditing: false,
        error: error instanceof Error ? error.message.slice(0, 1000) : String(error),
      });
    }
  }
  const report = {
    generatedAt: new Date().toISOString(),
    baseUrl: BASE_URL,
    authorizedImageModelCount: authorizedModels.length,
    testedImageModelCount: models.length,
    results,
  };
  const outputPath = resolve(
    process.env.VIP98K_IMAGE_TEST_REPORT || ".local/vip98k-image-model-tests.json",
  );
  await mkdir(dirname(outputPath), { recursive: true });
  await writeFile(outputPath, `${JSON.stringify(report, null, 2)}\n`, "utf8");
  console.log(JSON.stringify({ ...report, reportPath: outputPath }, null, 2));
  if (!results.some((result) => asRecord(result).status === "passed")) process.exitCode = 1;
}

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