Fix model test routing for image providers

2026-09-23 04:09:49 +00:00 · 2026-06-02 21:14:36 +08:00
parent cce8a50cac
commit e414975d0c
7 changed files with 309 additions and 163 deletions
@@ -0,0 +1,129 @@
+import { getApiKeys } from "@/lib/localDb";
+import { UPDATER_CONFIG } from "@/shared/constants/config";
+import { getConsistentMachineId } from "@/shared/utils/machineId";
+
+const CLI_TOKEN_SALT = "9r-cli-auth";
+
+async function getInternalHeaders() {
+  let apiKey = null;
+  try {
+    const keys = await getApiKeys();
+    apiKey = keys.find((k) => k.isActive !== false)?.key || null;
+  } catch {}
+
+  const headers = { "Content-Type": "application/json" };
+  if (apiKey) headers["Authorization"] = `Bearer ${apiKey}`;
+  headers["x-9r-cli-token"] = await getConsistentMachineId(CLI_TOKEN_SALT);
+  return headers;
+}
+
+export async function pingModelByKind(model, kind, baseUrl = `http://127.0.0.1:${process.env.PORT || UPDATER_CONFIG.appPort}`) {
+  const headers = await getInternalHeaders();
+  const start = Date.now();
+
+  if (kind === "embedding") {
+    const res = await fetch(`${baseUrl}/api/v1/embeddings`, {
+      method: "POST",
+      headers,
+      body: JSON.stringify({ model, input: "test" }),
+      signal: AbortSignal.timeout(15000),
+    });
+    const latencyMs = Date.now() - start;
+    const rawText = await res.text().catch(() => "");
+    let parsed = null;
+    try { parsed = rawText ? JSON.parse(rawText) : null; } catch {}
+
+    if (!res.ok) {
+      const detail = parsed?.error?.message || parsed?.error || rawText;
+      return { ok: false, latencyMs, error: `HTTP ${res.status}${detail ? `: ${String(detail).slice(0, 240)}` : ""}`, status: res.status };
+    }
+    const hasEmbedding = Array.isArray(parsed?.data) && parsed.data.length > 0 && Array.isArray(parsed.data[0]?.embedding);
+    if (!hasEmbedding) {
+      return { ok: false, latencyMs, status: res.status, error: "Provider returned no embedding data" };
+    }
+    return { ok: true, latencyMs, error: null, status: res.status };
+  }
+
+  if (kind === "image") {
+    const res = await fetch(`${baseUrl}/api/v1/images/generations`, {
+      method: "POST",
+      headers,
+      body: JSON.stringify({ model, prompt: "test" }),
+      signal: AbortSignal.timeout(15000),
+    });
+    const latencyMs = Date.now() - start;
+    const rawText = await res.text().catch(() => "");
+    let parsed = null;
+    try { parsed = rawText ? JSON.parse(rawText) : null; } catch {}
+
+    if (!res.ok) {
+      const detail = parsed?.error?.message || parsed?.msg || parsed?.message || parsed?.error || rawText;
+      return { ok: false, latencyMs, error: `HTTP ${res.status}${detail ? `: ${String(detail).slice(0, 240)}` : ""}`, status: res.status };
+    }
+
+    const hasImages = Array.isArray(parsed?.data) && parsed.data.length > 0;
+    if (!hasImages) {
+      return { ok: false, latencyMs, status: res.status, error: "Provider returned no image data for this model" };
+    }
+    return { ok: true, latencyMs, error: null, status: res.status };
+  }
+
+  const res = await fetch(`${baseUrl}/api/v1/chat/completions`, {
+    method: "POST",
+    headers,
+    body: JSON.stringify({
+      model,
+      max_tokens: 1,
+      stream: false,
+      messages: [{ role: "user", content: "hi" }],
+    }),
+    signal: AbortSignal.timeout(15000),
+  });
+  const latencyMs = Date.now() - start;
+
+  const rawText = await res.text().catch(() => "");
+  let parsed = null;
+  try { parsed = rawText ? JSON.parse(rawText) : null; } catch {}
+
+  if (!res.ok) {
+    const detail = parsed?.error?.message || parsed?.msg || parsed?.message || parsed?.error || rawText;
+    return { ok: false, latencyMs, error: `HTTP ${res.status}${detail ? `: ${String(detail).slice(0, 240)}` : ""}`, status: res.status };
+  }
+
+  const providerStatus = parsed?.status;
+  const providerMsg = parsed?.msg || parsed?.message;
+  const hasProviderErrorStatus = providerStatus !== undefined
+    && providerStatus !== null
+    && String(providerStatus) !== "200"
+    && String(providerStatus) !== "0";
+  if (hasProviderErrorStatus && providerMsg) {
+    return {
+      ok: false,
+      latencyMs,
+      status: res.status,
+      error: `Provider status ${providerStatus}: ${String(providerMsg).slice(0, 240)}`,
+    };
+  }
+
+  if (parsed?.error) {
+    const providerError = parsed?.error?.message || parsed?.error || "Provider returned an error";
+    return {
+      ok: false,
+      latencyMs,
+      status: res.status,
+      error: String(providerError).slice(0, 240),
+    };
+  }
+
+  const hasChoices = Array.isArray(parsed?.choices) && parsed.choices.length > 0;
+  if (!hasChoices) {
+    return {
+      ok: false,
+      latencyMs,
+      status: res.status,
+      error: "Provider returned no completion choices for this model",
+    };
+  }
+
+  return { ok: true, latencyMs, error: null, status: res.status };
+}