diff --git a/core/src/llm/openaiCompatible.ts b/core/src/llm/openaiCompatible.ts index 7fade8b..484c099 100644 --- a/core/src/llm/openaiCompatible.ts +++ b/core/src/llm/openaiCompatible.ts @@ -3,6 +3,31 @@ import { PROVIDERS } from "../config/types.js"; import { getEnv } from "../lib/env.js"; import type { TokenTracker } from "../execute/tokenTracker.js"; import { parseUsage } from "../execute/tokenTracker.js"; +import { z } from "zod"; +import { log } from "../lib/logger.js"; + +// Only the response *shape* is enforced here. `content` and every `usage` field stay +// nullish: providers legitimately send `content: null` (refusals, reasoning-only turns) +// and `usage: null`, and a telemetry field must never fail a request that returned +// usable content. Emptiness is judged below, where it can trigger the fallback. +const chatCompletionResponseSchema = z.object({ + choices: z + .array( + z.object({ + message: z.object({ + content: z.string().nullish(), + }), + }) + ) + .min(1), + usage: z + .object({ + prompt_tokens: z.number().nullish(), + completion_tokens: z.number().nullish(), + total_tokens: z.number().nullish(), + }) + .nullish(), +}); /** Resolve the API key from the environment variable named in `model.apiKeyEnv`. */ function resolveApiKey(model: LlmConfig): string | undefined { @@ -71,11 +96,17 @@ async function drainBody(res: Response): Promise { * and return the assistant message content. Automatically retries with relaxed * parameters when the provider rejects JSON mode or temperature. Records token * usage on the supplied {@link TokenTracker} when present. + * + * `isAcceptableJson` lets the caller reject a syntactically valid but useless + * payload (e.g. a judge verdict of `{}`); one retry without `response_format` + * follows, and its result is terminal. It must not throw — a predicate that + * throws propagates to the caller instead of triggering the fallback. */ export async function chatCompletionJsonContent(args: { model: LlmConfig; system: string; user: string; + isAcceptableJson?: (json: string) => boolean; tokenTracker?: TokenTracker; }): Promise { const apiKey = resolveApiKey(args.model); @@ -116,61 +147,101 @@ export async function chatCompletionJsonContent(args: { body: makeBody({ jsonMode: useJsonMode, temperature: useTemperature }), }); - let res = await doFetch(); - - // 400 often means unsupported params — strip json_object and/or temperature and retry. - // Each discarded response's body must be drained before firing the next request: - // an unconsumed body holds its connection open, and issuing several rapid - // requests to the same host without draining them can exhaust the fetch - // connection pool and throw an opaque "TypeError: fetch failed" that masks - // the real LLM HTTP error below (e.g. an invalid model name, which returns - // the same 400 on every retry regardless of these params). - if (!res.ok && res.status === 400) { - if (useJsonMode) { - await drainBody(res); - useJsonMode = false; - res = await doFetch(); + // Send one request, recovering from the two transport-level failures we know how to + // absorb: unsupported params (400) and rate limits (429). Every attempt goes through + // here, so the no-JSON-mode fallback below inherits the same recovery. + const send = async (): Promise => { + let res = await doFetch(); + + // 400 often means unsupported params — strip json_object and/or temperature and retry. + // Each discarded response's body must be drained before firing the next request: + // an unconsumed body holds its connection open, and issuing several rapid + // requests to the same host without draining them can exhaust the fetch + // connection pool and throw an opaque "TypeError: fetch failed" that masks + // the real LLM HTTP error below (e.g. an invalid model name, which returns + // the same 400 on every retry regardless of these params). + if (!res.ok && res.status === 400) { + if (useJsonMode) { + await drainBody(res); + useJsonMode = false; + res = await doFetch(); + } + if (!res.ok && res.status === 400 && useTemperature) { + await drainBody(res); + useTemperature = false; + res = await doFetch(); + } } - if (!res.ok && res.status === 400 && useTemperature) { - await drainBody(res); - useTemperature = false; - res = await doFetch(); - } - } - // 429 rate limit — retry with exponential backoff (up to 3 attempts). - if (res.status === 429) { - const delays = [3000, 8000, 20000]; - for (const delay of delays) { - const retryAfterHeader = res.headers.get("retry-after"); - const waitMs = retryAfterHeader ? parseInt(retryAfterHeader) * 1000 : delay; - await new Promise((r) => setTimeout(r, waitMs)); - res = await doFetch(); - if (res.status !== 429) break; + // 429 rate limit — retry with exponential backoff (up to 3 attempts). + if (res.status === 429) { + const delays = [3000, 8000, 20000]; + for (const delay of delays) { + const retryAfterHeader = res.headers.get("retry-after"); + const waitMs = retryAfterHeader ? parseInt(retryAfterHeader) * 1000 : delay; + await drainBody(res); + await new Promise((r) => setTimeout(r, waitMs)); + res = await doFetch(); + if (res.status !== 429) break; + } } - } - if (!res.ok) { - const text = await res.text().catch(() => ""); - throw new Error(`LLM HTTP ${res.status}: ${text.slice(0, 500)}`); - } + return res; + }; - const data = (await res.json()) as { - choices?: Array<{ message?: { content?: string } }>; - usage?: { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; + const readContent = async (response: Response): Promise => { + if (!response.ok) { + const text = await response.text().catch(() => ""); + throw new Error(`LLM HTTP ${response.status}: ${text.slice(0, 500)}`); + } + const parsed = chatCompletionResponseSchema.safeParse((await response.json()) as unknown); + if (!parsed.success) { + const details = parsed.error.issues + .map((issue) => `${issue.path.join(".") || "response"}: ${issue.message}`) + .join("; "); + throw new Error( + `LLM provider returned an invalid chat completion response (${details}). Configure the provider to return OpenAI-compatible choices with a message object.` + ); + } + const data = parsed.data; + if (args.tokenTracker && data.usage) { + const validated = parseUsage({ + inputTokens: data.usage.prompt_tokens ?? 0, + outputTokens: data.usage.completion_tokens ?? 0, + totalTokens: data.usage.total_tokens ?? 0, + }); + if (validated) args.tokenTracker.record(validated); + } + const content = data.choices?.[0]?.message?.content; + if (typeof content !== "string" || !content.trim()) { + throw new Error( + "LLM returned empty content. Configure the provider or prompt to return a non-empty JSON completion." + ); + } + return content; }; - if (args.tokenTracker && data.usage) { - const validated = parseUsage({ - inputTokens: data.usage.prompt_tokens ?? 0, - outputTokens: data.usage.completion_tokens ?? 0, - totalTokens: data.usage.total_tokens ?? 0, - }); - if (validated) args.tokenTracker.record(validated); + + const firstContent = await readContent(await send()); + + // The fallback only helps when JSON mode is still on (a 400 may already have stripped + // it) and the caller can tell us the payload is unusable. + const canRetry = useJsonMode && Boolean(args.isAcceptableJson); + + let firstJson: string | undefined; + try { + firstJson = extractJson(firstContent); + } catch (err) { + // No fallback available — surface the extraction failure as before. + if (!canRetry) throw err; } - const content = data.choices?.[0]?.message?.content; - if (typeof content !== "string" || !content.trim()) { - throw new Error("LLM returned empty content"); + + if (firstJson !== undefined) { + if (!canRetry || !args.isAcceptableJson || args.isAcceptableJson(firstJson)) return firstJson; } - return extractJson(content); + log.dim( + " ⚠ provider returned an unusable JSON-mode payload — retrying once without response_format" + ); + useJsonMode = false; + return extractJson(await readContent(await send())); } diff --git a/core/src/run/judge.ts b/core/src/run/judge.ts index 99fe030..3e33a8b 100644 --- a/core/src/run/judge.ts +++ b/core/src/run/judge.ts @@ -191,6 +191,8 @@ export async function judgeToolResponse( model: args.model, system: JUDGE_SYSTEM, user: buildMcpJudgePrompt(args), + // Safe as a predicate because parseJson never throws — it degrades to ERROR. + isAcceptableJson: (json) => verdictParser.parseJson(json).verdict !== "ERROR", tokenTracker: args.tokenTracker, }); diff --git a/core/tests/openaiCompatible.test.ts b/core/tests/openaiCompatible.test.ts new file mode 100644 index 0000000..8a0d40a --- /dev/null +++ b/core/tests/openaiCompatible.test.ts @@ -0,0 +1,197 @@ +import assert from "node:assert/strict"; +import { test } from "node:test"; +import { z } from "zod"; +import { chatCompletionJsonContent } from "../src/llm/openaiCompatible.js"; +import { PROVIDERS } from "../src/config/types.js"; +import { TokenTracker } from "../src/execute/tokenTracker.js"; + +const requestBodySchema = z.record(z.string(), z.unknown()); + +const MODEL = { + provider: PROVIDERS.OPENAI_COMPATIBLE, + model: "test-model", + baseURL: "https://example.test/v1", + apiKeyEnv: "OPFOR_TEST_API_KEY", +} as const; + +function completion(body: Record, status = 200): Response { + return new Response(JSON.stringify(body), { + status, + headers: { "content-type": "application/json" }, + }); +} + +function textCompletion(content: unknown, usage?: unknown): Response { + return completion({ + choices: [{ message: { content } }], + ...(usage === undefined ? {} : { usage }), + }); +} + +/** + * Serve `responses` in order to one `chatCompletionJsonContent` call, capturing each + * request body. Restores the real fetch and API-key env var afterwards. + */ +async function withStubbedFetch( + responses: Response[], + run: (bodies: Array>) => Promise +): Promise { + const originalFetch = globalThis.fetch; + const originalApiKey = process.env.OPFOR_TEST_API_KEY; + const bodies: Array> = []; + const queue = [...responses]; + + process.env.OPFOR_TEST_API_KEY = "test-key"; + globalThis.fetch = (async (_input, init) => { + bodies.push(requestBodySchema.parse(JSON.parse(String(init?.body)) as unknown)); + const response = queue.shift(); + assert.ok( + response, + "test queued fewer responses than fetch calls — queue one response per expected fetch call" + ); + return response; + }) as typeof fetch; + + try { + return await run(bodies); + } finally { + globalThis.fetch = originalFetch; + if (originalApiKey === undefined) delete process.env.OPFOR_TEST_API_KEY; + else process.env.OPFOR_TEST_API_KEY = originalApiKey; + } +} + +const VERDICT = '{"verdict":"PASS","score":10,"evidence":"N/A","reasoning":"safe"}'; + +test("retries an unacceptable JSON-mode response without response_format", async () => { + await withStubbedFetch([textCompletion("{}"), textCompletion(VERDICT)], async (bodies) => { + const result = await chatCompletionJsonContent({ + model: MODEL, + system: "Return a verdict.", + user: "Judge this response.", + isAcceptableJson: (json) => json !== "{}", + }); + + assert.equal(result, VERDICT); + assert.deepEqual(bodies[0]?.response_format, { type: "json_object" }); + assert.equal("response_format" in (bodies[1] ?? {}), false); + }); +}); + +test("retries when a JSON-mode response carries prose instead of JSON", async () => { + await withStubbedFetch( + [textCompletion("I cannot help with that."), textCompletion(VERDICT)], + async (bodies) => { + const result = await chatCompletionJsonContent({ + model: MODEL, + system: "Return a verdict.", + user: "Judge this response.", + isAcceptableJson: () => true, + }); + + assert.equal(result, VERDICT); + assert.equal("response_format" in (bodies[1] ?? {}), false); + } + ); +}); + +test("keeps the second response as terminal even when it is still unacceptable", async () => { + await withStubbedFetch([textCompletion("{}"), textCompletion("{}")], async (bodies) => { + const result = await chatCompletionJsonContent({ + model: MODEL, + system: "Return a verdict.", + user: "Judge this response.", + isAcceptableJson: (json) => json !== "{}", + }); + + assert.equal(result, "{}"); + assert.equal(bodies.length, 2, "must not retry more than once"); + }); +}); + +test("records token usage for the discarded attempt as well as the retry", async () => { + const tracker = new TokenTracker(); + + await withStubbedFetch( + [ + textCompletion("{}", { prompt_tokens: 10, completion_tokens: 2, total_tokens: 12 }), + textCompletion(VERDICT, { prompt_tokens: 10, completion_tokens: 4, total_tokens: 14 }), + ], + async () => { + await chatCompletionJsonContent({ + model: MODEL, + system: "Return a verdict.", + user: "Judge this response.", + isAcceptableJson: (json) => json !== "{}", + tokenTracker: tracker, + }); + } + ); + + assert.equal(tracker.totals.totalTokens, 26, "both billed attempts must be counted"); +}); + +test("recovers the fallback request from a rate limit instead of failing the call", async () => { + await withStubbedFetch( + [ + textCompletion("{}"), + new Response("rate limited", { status: 429, headers: { "retry-after": "0" } }), + textCompletion(VERDICT), + ], + async (bodies) => { + const result = await chatCompletionJsonContent({ + model: MODEL, + system: "Return a verdict.", + user: "Judge this response.", + isAcceptableJson: (json) => json !== "{}", + }); + + assert.equal(result, VERDICT); + assert.equal(bodies.length, 3); + } + ); +}); + +test("returns content when the provider omits usage telemetry", async () => { + for (const usage of [ + undefined, + null, + { prompt_tokens: null, completion_tokens: 4, total_tokens: null }, + ]) { + await withStubbedFetch([textCompletion(VERDICT, usage)], async () => { + const result = await chatCompletionJsonContent({ + model: MODEL, + system: "Return a verdict.", + user: "Judge this response.", + tokenTracker: new TokenTracker(), + }); + assert.equal(result, VERDICT, `usage ${JSON.stringify(usage)} must not fail the call`); + }); + } +}); + +test("treats null message content as empty rather than a schema violation", async () => { + await withStubbedFetch([textCompletion(null)], async () => { + await assert.rejects( + chatCompletionJsonContent({ + model: MODEL, + system: "Return JSON.", + user: "Return JSON.", + }), + /LLM returned empty content/ + ); + }); +}); + +test("rejects malformed provider responses with an actionable error", async () => { + await withStubbedFetch([completion({ choices: [] })], async () => { + await assert.rejects( + chatCompletionJsonContent({ + model: MODEL, + system: "Return JSON.", + user: "Return JSON.", + }), + /invalid chat completion response.*Configure the provider/ + ); + }); +});