diff --git a/README.md b/README.md index eba1d41..b0d7874 100644 --- a/README.md +++ b/README.md @@ -322,7 +322,43 @@ Authentication is handled via the `AI_GATEWAY_API_KEY` environment variable. Model slugs that start with `custom-` bypass the AI SDK gateway and are routed to `packages/cli/src/models/customModel.ts`. This lets you integrate any model backend — a local server, a custom API, or a model behind a proprietary SDK. -To add a custom model, edit `models/customModel.ts` and implement the `Model` interface: +#### Built-in: `custom-http:` + +Benchmark any model served behind the OpenAI chat-completions schema without writing code: a local vLLM, llama.cpp or Ollama server, an inference provider outside the AI Gateway, or a guard/wrapper service in front of a model. The slug carries the endpoint URL: + +```bash +# Open-weight model on a local vLLM server +CUSTOM_HTTP_MODEL=Qwen/Qwen3-8B \ + yarn kora run custom-http:http://localhost:8000/v1/chat/completions + +# Hosted provider that needs an API key +CUSTOM_HTTP_MODEL=openai/gpt-oss-20b CUSTOM_HTTP_API_KEY=$GROQ_API_KEY \ + yarn kora run custom-http:https://api.groq.com/openai/v1/chat/completions +``` + +Each target turn is sent as a standard request, and the answer is read from `choices[0].message.content`: + +```json +{ + "model": "", + "messages": [ ... ] +} +``` + +KORA sends no sampling parameters to targets, so the server's defaults apply. Set temperature and output length on the server (e.g. vLLM `--generation-config`, an Ollama `Modelfile`), and make sure its context window fits a full conversation. + +| Environment variable | Description | +| --------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------- | +| `CUSTOM_HTTP_MODEL` | Value of the request's `model` field, also recorded in the run stamp. Required by servers that host several models. | +| `CUSTOM_HTTP_API_KEY` | Sent as `Authorization: Bearer `. Leave unset for unauthenticated endpoints. Keep keys here rather than in the URL, which is logged and stamped. | + +Rate limits (429), timeouts (408), server errors (5xx) and dropped connections are retried with the same backoff as gateway models; other 4xx responses (bad key, unknown model, context overflow) fail immediately. Reasoning that a model returns inline, as a `…` block or just its closing tag when the chat template opens the block, is stripped so the judges grade only the answer, as with gateway targets. A reply that is still inside `` when it ends has no answer and fails the turn. + +`custom-http` is target-only: judges and the user model stay gateway models, since they need structured output. + +#### Custom JS implementation + +For anything beyond a plain HTTP passthrough, edit `models/customModel.ts` and implement the `Model` interface: ```ts export async function createCustomModel( @@ -858,6 +894,7 @@ packages/ gatewayModel.ts AI SDK gateway model implementation modelConfig.ts Model registry loader customModel.ts Custom model hook (edit to add your own) + customHttpModel.ts custom-http: OpenAI-compatible target retry.ts Retry with exponential backoff cli.ts CLI entry point ``` diff --git a/packages/benchmark/src/stamp/runStamp.ts b/packages/benchmark/src/stamp/runStamp.ts index 660679e..956e754 100644 --- a/packages/benchmark/src/stamp/runStamp.ts +++ b/packages/benchmark/src/stamp/runStamp.ts @@ -30,6 +30,9 @@ const VProfileRef = v.object({ const VRunnerTarget = v.object({ kind: v.picklist(["web-runner", "native-runner", "custom"]), slug: v.string(), + /** Model id a `custom-http:` target was asked for; the slug only + * names the endpoint, which may serve several models. */ + model: v.optional(v.string()), }); const VTargetRef = v.union([ModelSpec.io, VRunnerTarget]); diff --git a/packages/cli/src/models/__tests__/customHttpModel.test.ts b/packages/cli/src/models/__tests__/customHttpModel.test.ts new file mode 100644 index 0000000..6107a61 --- /dev/null +++ b/packages/cli/src/models/__tests__/customHttpModel.test.ts @@ -0,0 +1,236 @@ +import {Scenario} from "@korabench/benchmark"; +import {afterEach, beforeEach, describe, expect, it, vi} from "vitest"; +import {createCustomHttpModel} from "../customHttpModel.js"; +import {createCustomModel} from "../customModel.js"; + +const URL = "http://localhost:8000/v1/chat/completions"; +const SLUG = `custom-http:${URL}`; +// Scenario isn't read by custom-http; a stub keeps the tests focused. +const stubScenario = {} as unknown as Scenario; +const noRetry = {maxRetries: 0, onRetry: () => {}}; +const hi = {messages: [{role: "user" as const, content: "hi"}]}; + +function chatResponse(content: unknown, status = 200): Response { + return new Response(JSON.stringify({choices: [{message: {content}}]}), { + status, + headers: {"content-type": "application/json"}, + }); +} + +describe("custom-http", () => { + const originalFetch = global.fetch; + let fetchMock: ReturnType; + + beforeEach(() => { + fetchMock = vi.fn(); + global.fetch = fetchMock as unknown as typeof fetch; + }); + + afterEach(() => { + global.fetch = originalFetch; + vi.unstubAllEnvs(); + vi.useRealTimers(); + }); + + it("routes custom-http slugs through createCustomModel", async () => { + vi.stubEnv("CUSTOM_HTTP_MODEL", "qwen3-8b"); + vi.stubEnv("CUSTOM_HTTP_API_KEY", "sk-test"); + fetchMock.mockResolvedValue(chatResponse("hello")); + + const model = await createCustomModel(SLUG, stubScenario); + await expect(model.getTextResponse(hi)).resolves.toBe("hello"); + + const [calledUrl, init] = fetchMock.mock.calls[0]!; + expect(calledUrl).toBe(URL); + expect(JSON.parse(init.body).model).toBe("qwen3-8b"); + expect(init.headers).toMatchObject({authorization: "Bearer sk-test"}); + }); + + it("throws when the slug carries no URL", () => { + expect(() => createCustomHttpModel({modelSlug: "custom-http:"})).toThrow( + /missing a URL/ + ); + }); + + it("POSTs the OpenAI chat body shape", async () => { + fetchMock.mockResolvedValue(chatResponse("hello")); + const model = createCustomHttpModel({modelSlug: SLUG, retry: noRetry}); + + await model.getTextResponse({ + messages: [ + {role: "system", content: "be helpful"}, + {role: "user", content: "hi"}, + ], + temperature: 0.5, + maxTokens: 100, + }); + + const [, init] = fetchMock.mock.calls[0]!; + expect(init.method).toBe("POST"); + expect(JSON.parse(init.body)).toEqual({ + model: SLUG, + messages: [ + {role: "system", content: "be helpful"}, + {role: "user", content: "hi"}, + ], + temperature: 0.5, + max_tokens: 100, + }); + expect(init.headers).not.toHaveProperty("authorization"); + }); + + it("omits unset sampling parameters", async () => { + fetchMock.mockResolvedValue(chatResponse("ok")); + const model = createCustomHttpModel({modelSlug: SLUG, retry: noRetry}); + + await model.getTextResponse(hi); + + const body = JSON.parse(fetchMock.mock.calls[0]![1].body); + expect(body).not.toHaveProperty("temperature"); + expect(body).not.toHaveProperty("max_tokens"); + }); + + it.each([ + ["a full block", "\nThe child asks...\n\n\nHi!"], + // Templates that open the block in the prompt (DeepSeek-R1 distills, QwQ). + ["only the closing tag", "The child asks...\n\n\nHi!"], + ])("strips reasoning given %s", async (_, content) => { + fetchMock.mockResolvedValue(chatResponse(content)); + const model = createCustomHttpModel({modelSlug: SLUG, retry: noRetry}); + + await expect(model.getTextResponse(hi)).resolves.toBe("Hi!"); + }); + + it("passes plain answers through untouched", async () => { + fetchMock.mockResolvedValue(chatResponse(" Hi!\n")); + const model = createCustomHttpModel({modelSlug: SLUG, retry: noRetry}); + + await expect(model.getTextResponse(hi)).resolves.toBe(" Hi!\n"); + }); + + it("fails an unterminated block instead of grading it", async () => { + fetchMock.mockResolvedValue(chatResponse("\nThe child asks, so")); + const model = createCustomHttpModel({ + modelSlug: SLUG, + retry: {maxRetries: 3, onRetry: () => {}}, + }); + + await expect(model.getTextResponse(hi)).rejects.toThrow( + /unterminated block/ + ); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it.each([408, 429, 500, 503])("retries on %i", async status => { + vi.useFakeTimers(); + fetchMock + .mockResolvedValueOnce(new Response("busy", {status})) + .mockResolvedValueOnce(chatResponse("ok")); + const model = createCustomHttpModel({ + modelSlug: SLUG, + retry: {maxRetries: 1, onRetry: () => {}}, + }); + + const pending = model.getTextResponse(hi); + await vi.runAllTimersAsync(); + + await expect(pending).resolves.toBe("ok"); + expect(fetchMock).toHaveBeenCalledTimes(2); + }); + + // OpenAI-compatible servers tag these "invalid_request_error", which the + // message heuristic in retry.ts would otherwise treat as transient. + it.each([ + [400, "maximum context length exceeded"], + [401, "Invalid API Key"], + [404, "The model `qwen3` does not exist"], + ])("does not retry a %i", async (status, message) => { + fetchMock.mockResolvedValue( + new Response( + JSON.stringify({error: {message, type: "invalid_request_error"}}), + {status} + ) + ); + const model = createCustomHttpModel({ + modelSlug: SLUG, + retry: {maxRetries: 3, onRetry: () => {}}, + }); + + await expect(model.getTextResponse(hi)).rejects.toThrow( + `custom-http ${status}: ` + ); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("hints at CUSTOM_HTTP_MODEL when the slug is rejected as a model", async () => { + fetchMock.mockResolvedValue( + new Response('{"error":{"message":"invalid model name"}}', {status: 400}) + ); + const unset = createCustomHttpModel({modelSlug: SLUG, retry: noRetry}); + const set = createCustomHttpModel({ + modelSlug: SLUG, + model: "qwen3-8b", + retry: noRetry, + }); + + await expect(unset.getTextResponse(hi)).rejects.toThrow( + /set CUSTOM_HTTP_MODEL/ + ); + await expect(set.getTextResponse(hi)).rejects.not.toThrow( + /CUSTOM_HTTP_MODEL/ + ); + }); + + it("retries network failures and names the cause", async () => { + vi.useFakeTimers(); + fetchMock.mockRejectedValue( + new TypeError("fetch failed", { + cause: Object.assign(new Error("read ECONNRESET"), { + code: "ECONNRESET", + }), + }) + ); + const model = createCustomHttpModel({ + modelSlug: SLUG, + retry: {maxRetries: 2, onRetry: () => {}}, + }); + + const assertion = expect(model.getTextResponse(hi)).rejects.toThrow( + "fetch failed (ECONNRESET)" + ); + await vi.runAllTimersAsync(); + + await assertion; + expect(fetchMock).toHaveBeenCalledTimes(3); + }); + + it("throws when the response lacks choices[0].message.content", async () => { + fetchMock.mockResolvedValue(chatResponse(null)); + const model = createCustomHttpModel({ + modelSlug: SLUG, + retry: {maxRetries: 3, onRetry: () => {}}, + }); + + await expect(model.getTextResponse(hi)).rejects.toThrow( + /missing choices\[0\]\.message\.content/ + ); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("rejects structured requests", async () => { + const model = createCustomHttpModel({modelSlug: SLUG}); + + await expect( + model.getStructuredResponse({...hi, outputType: {} as never}) + ).rejects.toThrow(/do not support structured output/); + }); +}); + +describe("createCustomModel — unknown custom slug", () => { + it("points at custom-http in the error", async () => { + const model = await createCustomModel("custom-unknown", stubScenario); + await expect(model.getTextResponse(hi)).rejects.toThrow( + /custom-http:/ + ); + }); +}); diff --git a/packages/cli/src/models/customHttpModel.ts b/packages/cli/src/models/customHttpModel.ts new file mode 100644 index 0000000..8ff7bef --- /dev/null +++ b/packages/cli/src/models/customHttpModel.ts @@ -0,0 +1,177 @@ +import {ModelRequest, TypedModelRequest} from "@korabench/core"; +import { + createLogRetryHandler, + isRetryableError, + RetryOptions, + withRetry, +} from "../retry.js"; +import {Model} from "./model.js"; + +const CUSTOM_HTTP_PREFIX = "custom-http:"; +const THINK_CLOSE = ""; + +interface CustomHttpModelConfig { + modelSlug: string; + /** Sent as the request's `model` field. Defaults to the full slug, which + * suits wrapper services that ignore the field; OpenAI-compatible servers + * that host several models (vLLM, Ollama, Groq, Together, ...) need the + * real model id here. */ + model?: string; + /** Optional bearer token sent as `Authorization: Bearer `. When + * omitted, no auth header is sent (suits a local, unauthenticated server). */ + apiKey?: string; + retry?: RetryOptions; +} + +/** A failure whose cause is known, so retrying is decided by that cause rather + * than by keywords in the message: OpenAI-compatible servers put + * "invalid_request_error" in every 4xx body, which the message heuristic would + * treat as transient. */ +class CustomHttpError extends Error { + constructor( + message: string, + readonly retryable: boolean + ) { + super(message); + this.name = "CustomHttpError"; + } +} + +export function isCustomHttpSlug(slug: string): boolean { + return slug.startsWith(CUSTOM_HTTP_PREFIX); +} + +function isRetryableStatus(status: number): boolean { + return status === 408 || status === 429 || status >= 500; +} + +function shouldRetry(error: unknown): boolean { + return error instanceof CustomHttpError + ? error.retryable + : isRetryableError(error); +} + +// Reasoning models served without a reasoning parser (e.g. Qwen3 or +// DeepSeek-R1 on a plain vLLM/llama.cpp endpoint) put their chain of thought +// inline in `content`. Gateway targets never expose it to the judge, so drop +// everything up to the closing tag. Some chat templates open the block +// in the prompt, so the reply carries only the closing tag. A block that never +// closes (reasoning cut off by max tokens) has no answer to grade. +function stripReasoning(content: string): string { + const close = content.indexOf(THINK_CLOSE); + if (close !== -1) { + return content.slice(close + THINK_CLOSE.length).trim(); + } + if (content.trimStart().startsWith("")) { + throw new CustomHttpError( + `custom-http response is an unterminated block, with no ` + + `answer after the reasoning (did it hit the server's max tokens?)`, + false + ); + } + return content; +} + +function describeFetchError(error: unknown): string { + if (!(error instanceof Error)) return String(error); + // Node's fetch reports every network failure as "fetch failed" and keeps + // the actual reason (ECONNRESET, ECONNREFUSED, ...) on `cause`. + const cause = error.cause as {code?: unknown; message?: unknown} | undefined; + const detail = cause?.code ?? cause?.message; + return detail ? `${error.message} (${String(detail)})` : error.message; +} + +/** + * Target model that POSTs each conversation to any endpoint implementing the + * OpenAI chat-completions schema. The slug carries the URL: + * + * custom-http:http://localhost:8000/v1/chat/completions + * custom-http:https://api.groq.com/openai/v1/chat/completions + * + * This lets the benchmark target self-hosted open-weight models, inference + * providers outside the AI Gateway, or guard/wrapper services without + * editing customModel.ts. It is a target-only model: structured responses + * (judges, scenario generation) stay on gateway models. + */ +export function createCustomHttpModel(config: CustomHttpModelConfig): Model { + const url = config.modelSlug.slice(CUSTOM_HTTP_PREFIX.length); + if (!url) { + throw new Error( + `custom-http slug is missing a URL. Use "custom-http:", e.g. ` + + `"custom-http:http://localhost:8000/v1/chat/completions".` + ); + } + + const headers: Record = {"content-type": "application/json"}; + if (config.apiKey) headers["authorization"] = `Bearer ${config.apiKey}`; + + const retryOptions: RetryOptions = { + ...config.retry, + onRetry: config.retry?.onRetry ?? createLogRetryHandler(config.modelSlug), + shouldRetry: config.retry?.shouldRetry ?? shouldRetry, + }; + + async function send(body: string): Promise { + try { + return await fetch(url, {method: "POST", headers, body}); + } catch (error) { + throw new CustomHttpError( + `custom-http request to ${url} failed: ${describeFetchError(error)}`, + true + ); + } + } + + async function postChat(request: ModelRequest): Promise { + const r = await send( + JSON.stringify({ + model: config.model ?? config.modelSlug, + messages: request.messages, + ...(request.temperature !== undefined && { + temperature: request.temperature, + }), + ...(request.maxTokens !== undefined && { + max_tokens: request.maxTokens, + }), + }) + ); + if (!r.ok) { + const text = await r.text(); + // Servers hosting several models reject the slug as a model id. + const hint = + config.model === undefined && (r.status === 400 || r.status === 404) + ? ` (the request's model was the slug; set CUSTOM_HTTP_MODEL to ` + + `the id the server expects)` + : ""; + throw new CustomHttpError( + `custom-http ${r.status}: ${text.slice(0, 300)}${hint}`, + isRetryableStatus(r.status) + ); + } + const data = (await r.json()) as { + choices?: Array<{message?: {content?: string | null}}>; + }; + const content = data.choices?.[0]?.message?.content; + if (typeof content !== "string") { + throw new CustomHttpError( + `custom-http response missing choices[0].message.content: ` + + JSON.stringify(data).slice(0, 300), + false + ); + } + return stripReasoning(content); + } + + return { + getTextResponse(request: ModelRequest): Promise { + return withRetry(() => postChat(request), retryOptions); + }, + + async getStructuredResponse(_request: TypedModelRequest): Promise { + throw new Error( + `custom-http targets do not support structured output. Use a ` + + `gateway model for judges. Slug: ${config.modelSlug}` + ); + }, + }; +} diff --git a/packages/cli/src/models/customModel.ts b/packages/cli/src/models/customModel.ts index 4dbc327..3414501 100644 --- a/packages/cli/src/models/customModel.ts +++ b/packages/cli/src/models/customModel.ts @@ -1,4 +1,5 @@ import {Scenario} from "@korabench/benchmark"; +import {createCustomHttpModel, isCustomHttpSlug} from "./customHttpModel.js"; import {Model} from "./model.js"; import { createNativeRunnerModel, @@ -24,12 +25,20 @@ export async function createCustomModel( const apiKey = process.env.WEB_RUNNER_API_KEY; return createWebRunnerModel({modelSlug, webRunnerUrl, apiKey}); } + if (isCustomHttpSlug(modelSlug)) { + return createCustomHttpModel({ + modelSlug, + model: process.env.CUSTOM_HTTP_MODEL, + apiKey: process.env.CUSTOM_HTTP_API_KEY, + }); + } return { async getTextResponse() { throw new Error( `Custom model "${modelSlug}" is not implemented. ` + - `Provide an implementation in customModel.ts.` + `Provide an implementation in customModel.ts, or use ` + + `"custom-http:" for an OpenAI-compatible endpoint.` ); }, diff --git a/packages/cli/src/retry.ts b/packages/cli/src/retry.ts index f5db572..a1d9b0f 100644 --- a/packages/cli/src/retry.ts +++ b/packages/cli/src/retry.ts @@ -15,9 +15,15 @@ export interface RetryOptions { jitterFactor?: number; /** Called before each retry with attempt info */ onRetry?: (attempt: number, error: Error, delayMs: number) => void; + /** Decides whether an error is transient. Default: `isRetryableError`, which + * matches on the error message; callers that know the failure's cause (e.g. + * an HTTP status) can decide exactly. */ + shouldRetry?: (error: unknown) => boolean; } -const defaultOptions: Required> = { +type RetryTuning = Required>; + +const defaultOptions: RetryTuning = { maxRetries: 5, initialDelayMs: 1000, maxDelayMs: 60000, @@ -28,7 +34,7 @@ const defaultOptions: Required> = { /** * Determines if an error is retryable based on its characteristics. */ -function isRetryableError(error: unknown): boolean { +export function isRetryableError(error: unknown): boolean { if (!(error instanceof Error)) { return false; } @@ -138,7 +144,7 @@ function getRetryAfterMs(error: unknown): number | undefined { */ function calculateDelay( attempt: number, - options: Required>, + options: RetryTuning, retryAfterMs?: number ): number { // Respect retry-after header if present @@ -186,7 +192,7 @@ export async function withRetry( lastError = error instanceof Error ? error : new Error(String(error)); // Don't retry if this is not a retryable error - if (!isRetryableError(error)) { + if (!(options?.shouldRetry ?? isRetryableError)(error)) { throw lastError; } diff --git a/packages/cli/src/stamp/__tests__/buildRunStamp.test.ts b/packages/cli/src/stamp/__tests__/buildRunStamp.test.ts index f11e0b9..18af013 100644 --- a/packages/cli/src/stamp/__tests__/buildRunStamp.test.ts +++ b/packages/cli/src/stamp/__tests__/buildRunStamp.test.ts @@ -3,7 +3,7 @@ import {mkdtempSync, writeFileSync} from "node:fs"; import {tmpdir} from "node:os"; import * as path from "node:path"; import * as v from "valibot"; -import {describe, expect, it} from "vitest"; +import {afterEach, describe, expect, it, vi} from "vitest"; import {makeRoles} from "../../profiles/__tests__/fixtures.js"; import {EffectiveProfile} from "../../profiles/effectiveProfile.js"; import {Profile} from "../../profiles/profile.js"; @@ -40,6 +40,43 @@ describe("resolveTargetRef", () => { ])("maps %s to a %s reference", (slug, kind) => { expect(resolveTargetRef(modelsJsonPath, slug)).toEqual({kind, slug}); }); + + describe("custom-http", () => { + const slug = "custom-http:http://localhost:8000/v1/chat/completions"; + + afterEach(() => { + vi.unstubAllEnvs(); + }); + + it("records the model the endpoint was asked for", async () => { + vi.stubEnv("CUSTOM_HTTP_MODEL", "qwen3-8b"); + expect(resolveTargetRef(modelsJsonPath, slug)).toEqual({ + kind: "custom", + slug, + model: "qwen3-8b", + }); + + // The field survives the persisted schema, not just the builder. + const stamp = await buildRunStamp({ + effective, + modelsJsonPath, + target: slug, + }); + expect(v.parse(RunStamp.io, stamp).models.target).toEqual({ + kind: "custom", + slug, + model: "qwen3-8b", + }); + }); + + it("omits the model when CUSTOM_HTTP_MODEL is unset", () => { + vi.stubEnv("CUSTOM_HTTP_MODEL", ""); + expect(resolveTargetRef(modelsJsonPath, slug)).toEqual({ + kind: "custom", + slug, + }); + }); + }); }); describe("buildRunStamp", () => { diff --git a/packages/cli/src/stamp/buildRunStamp.ts b/packages/cli/src/stamp/buildRunStamp.ts index 91354cb..40a4ec7 100644 --- a/packages/cli/src/stamp/buildRunStamp.ts +++ b/packages/cli/src/stamp/buildRunStamp.ts @@ -5,6 +5,7 @@ import { RunStamp, TargetRef, } from "@korabench/benchmark"; +import {isCustomHttpSlug} from "../models/customHttpModel.js"; import {resolveModelConfig} from "../models/modelConfig.js"; import {isNativeRunnerSlug} from "../models/nativeRunnerModel.js"; import {isWebRunnerSlug} from "../models/webRunnerModel.js"; @@ -32,6 +33,10 @@ export function resolveTargetRef( modelsJsonPath: string, slug: string ): TargetRef { + if (isCustomHttpSlug(slug)) { + const model = process.env.CUSTOM_HTTP_MODEL; + return {kind: "custom", slug, ...(model ? {model} : {})}; + } if (slug.startsWith("custom-")) return {kind: "custom", slug}; if (isNativeRunnerSlug(slug)) return {kind: "native-runner", slug}; if (isWebRunnerSlug(slug)) return {kind: "web-runner", slug};