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

import sharp from "sharp";

const DEFAULT_BFL_BASE_URL = "https://api.bfl.ai";
const DEFAULT_BFL_MODEL_ID = "flux-2-klein-4b";
const DEFAULT_TIMEOUT_MS = 15_000;
const PROBE_IMAGE_SIZE = 64;

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

export type BflModelProbeResult = {
  supported: boolean;
  code: BflCapabilityCode;
  message: string;
  status?: number;
};

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

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

function encodeModelPath(modelId: string) {
  return modelId
    .split("/")
    .filter(Boolean)
    .map((part) => encodeURIComponent(part))
    .join("/");
}

function toErrorMessage(payload: any, fallback: string) {
  const direct = String(
    payload?.error?.message ||
      payload?.error ||
      payload?.message ||
      payload?.detail ||
      payload?.msg ||
      ""
  ).trim();
  return direct || fallback;
}

function looksLikeModelInvalid(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 hasTaskIdentity(payload: any) {
  return Boolean(
    (typeof payload?.id === "string" && payload.id.trim()) ||
      (typeof payload?.polling_url === "string" && payload.polling_url.trim()) ||
      (typeof payload?.response?.id === "string" && payload.response.id.trim()) ||
      (typeof payload?.response?.polling_url === "string" &&
        payload.response.polling_url.trim())
  );
}

let cachedProbeImageBase64: string | null = null;

async function getProbeImageBase64() {
  if (cachedProbeImageBase64) {
    return cachedProbeImageBase64;
  }

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

  cachedProbeImageBase64 = buffer.toString("base64");
  return cachedProbeImageBase64;
}

export async function probeBflModelSupport(params: BflProbeParams): Promise<BflModelProbeResult> {
  const apiKey = String(params.apiKey || "").trim();
  const modelId = String(params.modelId || DEFAULT_BFL_MODEL_ID).trim() || DEFAULT_BFL_MODEL_ID;
  const baseUrl = normalizeBaseUrl(params.baseUrl);
  const timeoutMs = Math.max(3_000, Number(params.timeoutMs || DEFAULT_TIMEOUT_MS));

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

  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), timeoutMs);

  try {
    const endpoint = `${baseUrl}/v1/${encodeModelPath(modelId)}`;
    const inputImage = await getProbeImageBase64();
    const response = await fetch(endpoint, {
      method: "POST",
      headers: {
        "x-key": apiKey,
        "Content-Type": "application/json",
        Accept: "application/json",
      },
      body: JSON.stringify({
        prompt: "probe image editing capability; preserve the input image composition",
        input_image: inputImage,
        width: PROBE_IMAGE_SIZE,
        height: PROBE_IMAGE_SIZE,
        safety_tolerance: 2,
      }),
      cache: "no-store",
      signal: controller.signal,
    });

    const payload = await response.json().catch(() => null);
    const message = toErrorMessage(payload, `BFL 探针请求失败（${response.status}）`);

    if (response.ok && hasTaskIdentity(payload)) {
      return {
        supported: true,
        code: "SUPPORTED",
        message: "模型可用。",
        status: response.status,
      };
    }

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

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

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