584 lines
18 KiB
TypeScript
584 lines
18 KiB
TypeScript
export type Protocol = "openai" | "openai_responses" | "claude";
|
||
|
||
const PROBE_TIMEOUT_MS = 45_000;
|
||
const PROBE_MAX_TOKENS = 256;
|
||
/** 存入 D1 前在 Worker 侧截断的流式输出上限 */
|
||
const STREAM_OUTPUT_CAP = 4096;
|
||
/** 非 2xx 时读取响应体截断长度（scheduler 还会再截断入库） */
|
||
const HTTP_ERROR_BODY_MAX = 6000;
|
||
|
||
function tryFormatJsonBody(raw: string): string {
|
||
const t = raw.trim();
|
||
if (!t.startsWith("{") && !t.startsWith("[")) return raw;
|
||
try {
|
||
return JSON.stringify(JSON.parse(t), null, 2);
|
||
} catch {
|
||
return raw;
|
||
}
|
||
}
|
||
|
||
async function readHttpErrorBody(res: Response): Promise<string> {
|
||
try {
|
||
const raw = (await res.text()).trim();
|
||
if (!raw) return "";
|
||
const formatted = tryFormatJsonBody(raw);
|
||
if (formatted.length <= HTTP_ERROR_BODY_MAX) return formatted;
|
||
return formatted.slice(0, HTTP_ERROR_BODY_MAX) + "…";
|
||
} catch {
|
||
return "";
|
||
}
|
||
}
|
||
|
||
function buildHttpErrorMessage(res: Response, body: string): string {
|
||
const line = `HTTP ${res.status}${res.statusText ? ` ${res.statusText}` : ""}`;
|
||
if (!body) return `${line}\n（响应体为空；常见原因：网关未返回 JSON 错误详情）`;
|
||
return `${line}\n\n${body}`;
|
||
}
|
||
|
||
function normalizeBase(url: string): string {
|
||
return url.replace(/\/+$/, "");
|
||
}
|
||
|
||
function appendCap(base: string, add: string, max: number): string {
|
||
if (base.length >= max) return base;
|
||
const room = max - base.length;
|
||
return base + (add.length <= room ? add : add.slice(0, room));
|
||
}
|
||
|
||
/** 经典 Chat Completions（/v1/chat/completions），兼容绝大多数 OpenAI 兼容网关 */
|
||
function openAiChatUrl(base: string): string {
|
||
const b = normalizeBase(base);
|
||
if (b.endsWith("/chat/completions")) return b;
|
||
if (b.endsWith("/v1")) return `${b}/chat/completions`;
|
||
return `${b}/chat/completions`;
|
||
}
|
||
|
||
/** OpenAI Responses API（/v1/responses） */
|
||
function openAiResponsesUrl(base: string): string {
|
||
const b = normalizeBase(base);
|
||
if (b.endsWith("/responses")) return b;
|
||
if (b.endsWith("/v1")) return `${b}/responses`;
|
||
return `${b}/v1/responses`;
|
||
}
|
||
|
||
function claudeUrl(base: string): string {
|
||
const b = normalizeBase(base);
|
||
if (b.endsWith("/messages")) return b;
|
||
if (b.endsWith("/v1")) return `${b}/messages`;
|
||
return `${b}/v1/messages`;
|
||
}
|
||
|
||
/**
|
||
* Chat Completions 流式里 `choices[0].delta` 的正文抽取。
|
||
* 对齐 OpenAI 常见形态：delta.content 字符串；多段/多模态下为 part 数组；
|
||
* 部分国产/推理模型会在 delta.reasoning_content 等字段里先出字。
|
||
*/
|
||
function streamDeltaTextFromOpenAiChatChoice(delta: unknown): string {
|
||
if (delta == null || typeof delta !== "object") return "";
|
||
const d = delta as Record<string, unknown>;
|
||
const c = d.content;
|
||
if (typeof c === "string" && c !== "") return c;
|
||
if (Array.isArray(c)) {
|
||
let s = "";
|
||
for (const part of c) {
|
||
if (part && typeof part === "object") {
|
||
const p = part as { type?: string; text?: string };
|
||
if (typeof p.text === "string") s += p.text;
|
||
}
|
||
}
|
||
return s;
|
||
}
|
||
for (const k of ["reasoning_content", "reasoning"] as const) {
|
||
const v = d[k];
|
||
if (typeof v === "string" && v !== "") return v;
|
||
}
|
||
return "";
|
||
}
|
||
|
||
async function readStreamOpenAIChat(
|
||
res: Response,
|
||
started: number
|
||
): Promise<{ firstTokenMs: number | null; httpStatus: number; outputText: string }> {
|
||
if (!res.ok || !res.body) {
|
||
return { firstTokenMs: null, httpStatus: res.status, outputText: "" };
|
||
}
|
||
const reader = res.body.getReader();
|
||
const dec = new TextDecoder();
|
||
let buf = "";
|
||
let firstTokenMs: number | null = null;
|
||
let outputText = "";
|
||
try {
|
||
for (;;) {
|
||
const { done, value } = await reader.read();
|
||
if (done) break;
|
||
buf += dec.decode(value, { stream: true });
|
||
const lines = buf.split("\n");
|
||
buf = lines.pop() ?? "";
|
||
for (const line of lines) {
|
||
const t = line.trim();
|
||
if (!t.startsWith("data:")) continue;
|
||
const payload = t.slice(5).trim();
|
||
if (payload === "[DONE]") {
|
||
return { firstTokenMs, httpStatus: res.status, outputText };
|
||
}
|
||
try {
|
||
const obj = JSON.parse(payload) as {
|
||
choices?: Array<{ delta?: unknown }>;
|
||
};
|
||
const delta = obj.choices?.[0]?.delta;
|
||
const chunk = streamDeltaTextFromOpenAiChatChoice(delta);
|
||
if (chunk !== "") {
|
||
if (firstTokenMs == null) firstTokenMs = Math.max(0, Date.now() - started);
|
||
outputText = appendCap(outputText, chunk, STREAM_OUTPUT_CAP);
|
||
}
|
||
} catch {
|
||
/* ignore bad json line */
|
||
}
|
||
}
|
||
}
|
||
} finally {
|
||
reader.releaseLock();
|
||
}
|
||
return { firstTokenMs, httpStatus: res.status, outputText };
|
||
}
|
||
|
||
async function readStreamOpenAIResponses(
|
||
res: Response,
|
||
started: number
|
||
): Promise<{ firstTokenMs: number | null; httpStatus: number; outputText: string }> {
|
||
if (!res.ok || !res.body) {
|
||
return { firstTokenMs: null, httpStatus: res.status, outputText: "" };
|
||
}
|
||
const reader = res.body.getReader();
|
||
const dec = new TextDecoder();
|
||
let buf = "";
|
||
let firstTokenMs: number | null = null;
|
||
let outputText = "";
|
||
try {
|
||
for (;;) {
|
||
const { done, value } = await reader.read();
|
||
if (done) break;
|
||
buf += dec.decode(value, { stream: true });
|
||
const lines = buf.split("\n");
|
||
buf = lines.pop() ?? "";
|
||
for (const line of lines) {
|
||
const t = line.trim();
|
||
if (!t.startsWith("data:")) continue;
|
||
const payload = t.slice(5).trim();
|
||
if (!payload || payload === "[DONE]") continue;
|
||
try {
|
||
const obj = JSON.parse(payload) as { type?: string; delta?: string };
|
||
if (obj.type === "response.output_text.delta" && obj.delta != null && obj.delta !== "") {
|
||
if (firstTokenMs == null) firstTokenMs = Math.max(0, Date.now() - started);
|
||
outputText = appendCap(outputText, obj.delta, STREAM_OUTPUT_CAP);
|
||
}
|
||
} catch {
|
||
/* ignore */
|
||
}
|
||
}
|
||
}
|
||
} finally {
|
||
reader.releaseLock();
|
||
}
|
||
return { firstTokenMs, httpStatus: res.status, outputText };
|
||
}
|
||
|
||
function openAiChatMessageContentToText(content: unknown): string {
|
||
if (typeof content === "string") return content;
|
||
if (!Array.isArray(content)) return "";
|
||
let s = "";
|
||
for (const part of content) {
|
||
if (part && typeof part === "object") {
|
||
const p = part as { type?: string; text?: string };
|
||
if (typeof p.text === "string") s += p.text;
|
||
}
|
||
}
|
||
return s;
|
||
}
|
||
|
||
function parseOpenAiChatNonStream(
|
||
json: unknown,
|
||
started: number
|
||
): { firstTokenMs: number | null; outputText: string } {
|
||
const obj = json as { choices?: Array<{ message?: { content?: unknown } }> };
|
||
const raw = openAiChatMessageContentToText(obj.choices?.[0]?.message?.content);
|
||
const outputText = raw.length <= STREAM_OUTPUT_CAP ? raw : raw.slice(0, STREAM_OUTPUT_CAP);
|
||
const firstTokenMs = outputText !== "" ? Math.max(0, Date.now() - started) : null;
|
||
return { firstTokenMs, outputText };
|
||
}
|
||
|
||
function extractOpenAiResponsesText(data: unknown): string {
|
||
if (data == null || typeof data !== "object") return "";
|
||
const d = data as Record<string, unknown>;
|
||
if (typeof d.output_text === "string") return d.output_text;
|
||
|
||
const out = d.output;
|
||
if (!Array.isArray(out)) return "";
|
||
|
||
const parts: string[] = [];
|
||
const walk = (node: unknown): void => {
|
||
if (node == null) return;
|
||
if (typeof node === "string") {
|
||
parts.push(node);
|
||
return;
|
||
}
|
||
if (Array.isArray(node)) {
|
||
for (const x of node) walk(x);
|
||
return;
|
||
}
|
||
if (typeof node !== "object") return;
|
||
const o = node as Record<string, unknown>;
|
||
if (o.type === "output_text" && typeof o.text === "string") {
|
||
parts.push(o.text);
|
||
return;
|
||
}
|
||
if (typeof o.text === "string" && typeof o.type === "string" && o.type.includes("text")) {
|
||
parts.push(o.text);
|
||
return;
|
||
}
|
||
if (o.content != null) walk(o.content);
|
||
if (o.output != null) walk(o.output);
|
||
};
|
||
for (const item of out) walk(item);
|
||
return parts.join("");
|
||
}
|
||
|
||
function parseOpenAiResponsesNonStream(
|
||
json: unknown,
|
||
started: number
|
||
): { firstTokenMs: number | null; outputText: string } {
|
||
const raw = extractOpenAiResponsesText(json);
|
||
const outputText = raw.length <= STREAM_OUTPUT_CAP ? raw : raw.slice(0, STREAM_OUTPUT_CAP);
|
||
const firstTokenMs = outputText !== "" ? Math.max(0, Date.now() - started) : null;
|
||
return { firstTokenMs, outputText };
|
||
}
|
||
|
||
function extractClaudeNonStreamText(json: unknown): string {
|
||
if (json == null || typeof json !== "object") return "";
|
||
const d = json as { content?: unknown };
|
||
const content = d.content;
|
||
if (!Array.isArray(content)) return "";
|
||
let s = "";
|
||
for (const block of content) {
|
||
if (block && typeof block === "object") {
|
||
const b = block as { type?: string; text?: string };
|
||
if (b.type === "text" && typeof b.text === "string") s += b.text;
|
||
}
|
||
}
|
||
return s;
|
||
}
|
||
|
||
function parseClaudeNonStream(
|
||
json: unknown,
|
||
started: number
|
||
): { firstTokenMs: number | null; outputText: string } {
|
||
const raw = extractClaudeNonStreamText(json);
|
||
const outputText = raw.length <= STREAM_OUTPUT_CAP ? raw : raw.slice(0, STREAM_OUTPUT_CAP);
|
||
const firstTokenMs = outputText !== "" ? Math.max(0, Date.now() - started) : null;
|
||
return { firstTokenMs, outputText };
|
||
}
|
||
|
||
/** 首字可来自正文或 thinking；输出只累计 assistant 可见正文 delta.text */
|
||
function claudeStreamDeltaParts(
|
||
data: string,
|
||
currentEvent: string
|
||
): { outputChunk: string | null; marksFirstToken: boolean } {
|
||
try {
|
||
const obj = JSON.parse(data) as {
|
||
type?: string;
|
||
delta?: { type?: string; text?: string; thinking?: string };
|
||
};
|
||
const match = currentEvent === "content_block_delta" || obj.type === "content_block_delta";
|
||
if (!match) return { outputChunk: null, marksFirstToken: false };
|
||
const d = obj.delta;
|
||
if (!d) return { outputChunk: null, marksFirstToken: false };
|
||
const hasThinking = typeof d.thinking === "string" && d.thinking !== "";
|
||
const hasText = typeof d.text === "string" && d.text !== "";
|
||
return {
|
||
outputChunk: hasText ? (d.text as string) : null,
|
||
marksFirstToken: hasThinking || hasText,
|
||
};
|
||
} catch {
|
||
return { outputChunk: null, marksFirstToken: false };
|
||
}
|
||
}
|
||
|
||
async function readStreamClaude(
|
||
res: Response,
|
||
started: number
|
||
): Promise<{ firstTokenMs: number | null; httpStatus: number; outputText: string }> {
|
||
if (!res.ok || !res.body) {
|
||
return { firstTokenMs: null, httpStatus: res.status, outputText: "" };
|
||
}
|
||
const reader = res.body.getReader();
|
||
const dec = new TextDecoder();
|
||
let buf = "";
|
||
let currentEvent = "";
|
||
let firstTokenMs: number | null = null;
|
||
let outputText = "";
|
||
try {
|
||
for (;;) {
|
||
const { done, value } = await reader.read();
|
||
if (done) break;
|
||
buf += dec.decode(value, { stream: true });
|
||
let idx: number;
|
||
while ((idx = buf.indexOf("\n")) >= 0) {
|
||
let line = buf.slice(0, idx);
|
||
buf = buf.slice(idx + 1);
|
||
if (line.endsWith("\r")) line = line.slice(0, -1);
|
||
if (line === "") {
|
||
currentEvent = "";
|
||
continue;
|
||
}
|
||
if (line.startsWith("event:")) {
|
||
currentEvent = line.slice(6).trim();
|
||
continue;
|
||
}
|
||
if (!line.startsWith("data:")) continue;
|
||
const data = line.slice(5).trim();
|
||
const parts = claudeStreamDeltaParts(data, currentEvent);
|
||
if (parts.marksFirstToken && firstTokenMs == null) {
|
||
firstTokenMs = Math.max(0, Date.now() - started);
|
||
}
|
||
if (parts.outputChunk) {
|
||
outputText = appendCap(outputText, parts.outputChunk, STREAM_OUTPUT_CAP);
|
||
}
|
||
}
|
||
}
|
||
} finally {
|
||
reader.releaseLock();
|
||
}
|
||
return { firstTokenMs, httpStatus: res.status, outputText };
|
||
}
|
||
|
||
export type ProbeResult = {
|
||
ok: boolean;
|
||
firstTokenMs: number | null;
|
||
httpStatus: number | null;
|
||
errorMessage: string | null;
|
||
requestMessage: string;
|
||
responseText: string | null;
|
||
};
|
||
|
||
export async function runProbe(params: {
|
||
apiBaseUrl: string;
|
||
apiKey: string;
|
||
model: string;
|
||
protocol: Protocol;
|
||
userMessage: string;
|
||
/** 默认 true（流式）；false 时使用非流式 JSON 响应 */
|
||
stream?: boolean;
|
||
}): Promise<ProbeResult> {
|
||
const userMessage = params.userMessage.trim() || "ping";
|
||
const useStream = params.stream !== false;
|
||
const ac = new AbortController();
|
||
const t = setTimeout(() => ac.abort(), PROBE_TIMEOUT_MS);
|
||
const started = Date.now();
|
||
try {
|
||
if (params.protocol === "openai") {
|
||
const url = openAiChatUrl(params.apiBaseUrl);
|
||
const res = await fetch(url, {
|
||
method: "POST",
|
||
signal: ac.signal,
|
||
headers: {
|
||
"Content-Type": "application/json",
|
||
Authorization: `Bearer ${params.apiKey}`,
|
||
},
|
||
body: JSON.stringify({
|
||
model: params.model,
|
||
messages: [{ role: "user", content: userMessage }],
|
||
max_tokens: PROBE_MAX_TOKENS,
|
||
stream: useStream,
|
||
}),
|
||
});
|
||
if (!res.ok) {
|
||
const errBody = await readHttpErrorBody(res);
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: res.status,
|
||
errorMessage: buildHttpErrorMessage(res, errBody),
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
}
|
||
if (useStream) {
|
||
const { firstTokenMs, httpStatus, outputText } = await readStreamOpenAIChat(res, started);
|
||
const ok = firstTokenMs != null;
|
||
return {
|
||
ok,
|
||
firstTokenMs,
|
||
httpStatus,
|
||
errorMessage: ok ? null : "no_stream_token",
|
||
requestMessage: userMessage,
|
||
responseText: outputText === "" ? null : outputText,
|
||
};
|
||
}
|
||
let json: unknown;
|
||
try {
|
||
json = await res.json();
|
||
} catch {
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: res.status,
|
||
errorMessage: "invalid_json_body",
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
}
|
||
const { firstTokenMs, outputText } = parseOpenAiChatNonStream(json, started);
|
||
const ok = firstTokenMs != null;
|
||
return {
|
||
ok,
|
||
firstTokenMs,
|
||
httpStatus: res.status,
|
||
errorMessage: ok ? null : "no_response_content",
|
||
requestMessage: userMessage,
|
||
responseText: outputText === "" ? null : outputText,
|
||
};
|
||
}
|
||
if (params.protocol === "openai_responses") {
|
||
const url = openAiResponsesUrl(params.apiBaseUrl);
|
||
const res = await fetch(url, {
|
||
method: "POST",
|
||
signal: ac.signal,
|
||
headers: {
|
||
"Content-Type": "application/json",
|
||
Authorization: `Bearer ${params.apiKey}`,
|
||
},
|
||
body: JSON.stringify({
|
||
model: params.model,
|
||
input: userMessage,
|
||
stream: useStream,
|
||
}),
|
||
});
|
||
if (!res.ok) {
|
||
const errBody = await readHttpErrorBody(res);
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: res.status,
|
||
errorMessage: buildHttpErrorMessage(res, errBody),
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
}
|
||
if (useStream) {
|
||
const { firstTokenMs, httpStatus, outputText } = await readStreamOpenAIResponses(res, started);
|
||
const ok = firstTokenMs != null;
|
||
return {
|
||
ok,
|
||
firstTokenMs,
|
||
httpStatus,
|
||
errorMessage: ok ? null : "no_stream_token",
|
||
requestMessage: userMessage,
|
||
responseText: outputText === "" ? null : outputText,
|
||
};
|
||
}
|
||
let json: unknown;
|
||
try {
|
||
json = await res.json();
|
||
} catch {
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: res.status,
|
||
errorMessage: "invalid_json_body",
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
}
|
||
const { firstTokenMs, outputText } = parseOpenAiResponsesNonStream(json, started);
|
||
const ok = firstTokenMs != null;
|
||
return {
|
||
ok,
|
||
firstTokenMs,
|
||
httpStatus: res.status,
|
||
errorMessage: ok ? null : "no_response_content",
|
||
requestMessage: userMessage,
|
||
responseText: outputText === "" ? null : outputText,
|
||
};
|
||
}
|
||
const url = claudeUrl(params.apiBaseUrl);
|
||
const res = await fetch(url, {
|
||
method: "POST",
|
||
signal: ac.signal,
|
||
headers: {
|
||
"Content-Type": "application/json",
|
||
"x-api-key": params.apiKey,
|
||
"anthropic-version": "2023-06-01",
|
||
},
|
||
body: JSON.stringify({
|
||
model: params.model,
|
||
max_tokens: PROBE_MAX_TOKENS,
|
||
messages: [{ role: "user", content: userMessage }],
|
||
stream: useStream,
|
||
}),
|
||
});
|
||
if (!res.ok) {
|
||
const errBody = await readHttpErrorBody(res);
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: res.status,
|
||
errorMessage: buildHttpErrorMessage(res, errBody),
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
}
|
||
if (useStream) {
|
||
const { firstTokenMs, httpStatus, outputText } = await readStreamClaude(res, started);
|
||
const ok = firstTokenMs != null;
|
||
return {
|
||
ok,
|
||
firstTokenMs,
|
||
httpStatus,
|
||
errorMessage: ok ? null : "no_stream_token",
|
||
requestMessage: userMessage,
|
||
responseText: outputText === "" ? null : outputText,
|
||
};
|
||
}
|
||
let json: unknown;
|
||
try {
|
||
json = await res.json();
|
||
} catch {
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: res.status,
|
||
errorMessage: "invalid_json_body",
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
}
|
||
const { firstTokenMs, outputText } = parseClaudeNonStream(json, started);
|
||
const ok = firstTokenMs != null;
|
||
return {
|
||
ok,
|
||
firstTokenMs,
|
||
httpStatus: res.status,
|
||
errorMessage: ok ? null : "no_response_content",
|
||
requestMessage: userMessage,
|
||
responseText: outputText === "" ? null : outputText,
|
||
};
|
||
} catch (e) {
|
||
const msg = e instanceof Error ? (e.name === "AbortError" ? "timeout" : e.message) : "unknown_error";
|
||
return {
|
||
ok: false,
|
||
firstTokenMs: null,
|
||
httpStatus: null,
|
||
errorMessage: truncateErr(msg),
|
||
requestMessage: userMessage,
|
||
responseText: null,
|
||
};
|
||
} finally {
|
||
clearTimeout(t);
|
||
}
|
||
}
|
||
|
||
function truncateErr(s: string, max = 120): string {
|
||
if (s.length <= max) return s;
|
||
return s.slice(0, max);
|
||
}
|