ai-chatbot/lib/realtime/replicate-capability.ts

import sharp from "sharp";

const DEFAULT_REPLICATE_BASE_URL = "https://api.replicate.com/v1";
const DEFAULT_REPLICATE_MODEL = "black-forest-labs/flux-2-klein-4b";
const DEFAULT_TIMEOUT_MS = 20_000;

export type ReplicateCapabilityCode =
  | "SUPPORTED"
  | "AUTH_FAILED"
  | "MODEL_INVALID"
  | "UPSTREAM_ERROR";

export type ReplicateModelProbeResult = {
  supported: boolean;
  code: ReplicateCapabilityCode;
  message: string;
  status?: number;
};

type ReplicateProbeParams = {
  apiKey: string;
  modelId?: string;
  baseUrl?: string;
  timeoutMs?: number;
};

let cachedProbeImageDataUri: string | null = null;

function trimString(value: unknown) {
  return String(value || "").trim();
}

function normalizeBaseUrl(raw: string | undefined) {
  const next = trimString(raw || DEFAULT_REPLICATE_BASE_URL);
  return (next || DEFAULT_REPLICATE_BASE_URL).replace(/\/+$/, "");
}

function toErrorMessage(payload: any, fallback: string) {
  const direct = trimString(
    payload?.detail ||
      payload?.error ||
      payload?.message ||
      payload?.title ||
      payload?.input?.message ||
      payload?.non_field_errors?.[0]
  );
  return direct || fallback;
}

function isModelInvalid(message: string) {
  const normalized = message.toLowerCase();
  return (
    normalized.includes("model") &&
    (normalized.includes("not found") ||
      normalized.includes("invalid") ||
      normalized.includes("unknown") ||
      normalized.includes("does not exist"))
  );
}

function splitModelIdentifier(modelId: string) {
  const [owner, ...rest] = modelId.split("/").filter(Boolean);
  const name = rest.join("/");
  if (!owner || !name) {
    throw new Error(`无效的 Replicate 模型标识：${modelId}`);
  }
  return { owner, name };
}

function shouldRetryWithAltField(message: string) {
  const normalized = message.toLowerCase();
  return (
    normalized.includes("images") ||
    normalized.includes("image") ||
    normalized.includes("input_image") ||
    normalized.includes("schema") ||
    normalized.includes("validation")
  );
}

async function getProbeImageDataUri() {
  if (cachedProbeImageDataUri) return cachedProbeImageDataUri;

  const buffer = await sharp({
    create: {
      width: 64,
      height: 64,
      channels: 3,
      background: { r: 255, g: 255, b: 255 },
    },
  })
    .png()
    .toBuffer();

  cachedProbeImageDataUri = `data:image/png;base64,${buffer.toString("base64")}`;
  return cachedProbeImageDataUri;
}

async function createPrediction(params: {
  apiKey: string;
  baseUrl: string;
  modelId: string;
  input: Record<string, unknown>;
  waitSeconds?: number;
  signal?: AbortSignal;
}) {
  const { owner, name } = splitModelIdentifier(params.modelId);
  const endpoint = `${params.baseUrl}/models/${encodeURIComponent(owner)}/${encodeURIComponent(name)}/predictions`;
  const response = await fetch(endpoint, {
    method: "POST",
    headers: {
      Authorization: `Bearer ${params.apiKey}`,
      "Content-Type": "application/json",
      Prefer: `wait=${Math.max(1, Number(params.waitSeconds || 1))}`,
    },
    body: JSON.stringify({ input: params.input }),
    cache: "no-store",
    signal: params.signal,
  });

  const payload = await response.json().catch(() => null);
  return { response, payload };
}

export async function probeReplicateModelSupport(
  params: ReplicateProbeParams
): Promise<ReplicateModelProbeResult> {
  const apiKey = trimString(params.apiKey);
  const modelId = trimString(params.modelId || DEFAULT_REPLICATE_MODEL) || DEFAULT_REPLICATE_MODEL;
  const baseUrl = normalizeBaseUrl(params.baseUrl);

  if (!apiKey) {
    return {
      supported: false,
      code: "AUTH_FAILED",
      message: "缺少 Replicate API Token。",
    };
  }

  const controller = new AbortController();
  const timer = setTimeout(
    () => controller.abort(),
    Math.max(3_000, Number(params.timeoutMs || DEFAULT_TIMEOUT_MS))
  );

  try {
    const image = await getProbeImageDataUri();
    const candidateInputs = [
      { prompt: "white square photo", images: [image] },
      { prompt: "white square photo", image },
      { prompt: "white square photo", input_image: image },
    ];

    for (let index = 0; index < candidateInputs.length; index += 1) {
      const { response, payload } = await createPrediction({
        apiKey,
        baseUrl,
        modelId,
        input: candidateInputs[index],
        waitSeconds: 1,
        signal: controller.signal,
      });

      const message = toErrorMessage(payload, `Replicate 探针请求失败（${response.status}）`);

      if (
        response.ok &&
        (trimString(payload?.id) ||
          trimString(payload?.status) ||
          trimString(payload?.urls?.get))
      ) {
        return {
          supported: true,
          code: "SUPPORTED",
          message: "模型可用。",
          status: response.status,
        };
      }

      if (response.status === 401 || response.status === 403) {
        return {
          supported: false,
          code: "AUTH_FAILED",
          message: "Replicate API Token 无效或无权限，请检查后重试。",
          status: response.status,
        };
      }

      if (response.status === 404 || isModelInvalid(message)) {
        return {
          supported: false,
          code: "MODEL_INVALID",
          message: `模型 ${modelId} 不存在或不可用，请更换模型后重试。`,
          status: response.status,
        };
      }

      if (index < candidateInputs.length - 1 && shouldRetryWithAltField(message)) {
        continue;
      }

      return {
        supported: false,
        code: "UPSTREAM_ERROR",
        message,
        status: response.status,
      };
    }

    return {
      supported: false,
      code: "UPSTREAM_ERROR",
      message: "Replicate 探针请求失败。",
    };
  } catch (error) {
    if (error instanceof Error && error.name === "AbortError") {
      return {
        supported: false,
        code: "UPSTREAM_ERROR",
        message: "Replicate 探针请求超时，请稍后重试。",
      };
    }

    return {
      supported: false,
      code: "UPSTREAM_ERROR",
      message: error instanceof Error ? error.message : "Replicate 探针请求失败。",
    };
  } finally {
    clearTimeout(timer);
  }
}