can1357
3f32971944
OpenRouter previously omitted `max_tokens` entirely (except for specific models) to prevent unintended provider filtering when a model's catalog default `maxTokens` exceeded an upstream's actual capacity. This could lead to incorrect routing if a model had a high catalog cap but individual providers under OpenRouter did not.
456 lines
16 KiB
TypeScript
456 lines
16 KiB
TypeScript
import type { ResolvedThinkingLevel } from "@oh-my-pi/pi-agent-core";
|
|
import type {
|
|
Api,
|
|
ApiKeyResolver,
|
|
AssistantMessage,
|
|
AssistantMessageEvent,
|
|
AssistantMessageEventStream,
|
|
Context,
|
|
Effort,
|
|
Model,
|
|
SimpleStreamOptions,
|
|
} from "@oh-my-pi/pi-ai";
|
|
import { streamSimple } from "@oh-my-pi/pi-ai";
|
|
import type { CanonicalModelVariant } from "@oh-my-pi/pi-catalog/identity";
|
|
import { replaceTabs, truncateToWidth } from "@oh-my-pi/pi-tui";
|
|
import { formatDuration, getProjectDir } from "@oh-my-pi/pi-utils";
|
|
import chalk from "chalk";
|
|
import type { ApiKeyResolverModel } from "../config/api-key-resolver";
|
|
import { type CanonicalModelQueryOptions, ModelRegistry } from "../config/model-registry";
|
|
import { formatModelSelectorValue, formatModelString, getModelMatchPreferences, resolveCliModel } from "../config/model-resolver";
|
|
import { Settings } from "../config/settings";
|
|
import benchPrompt from "../prompts/bench.md" with { type: "text" };
|
|
import { discoverAuthStorage } from "../sdk";
|
|
import { resolveThinkingLevelForModel, shouldDisableReasoning, toReasoningEffort } from "../thinking";
|
|
|
|
const DEFAULT_RUNS = 1;
|
|
const DEFAULT_MAX_TOKENS = 512;
|
|
const ERROR_WIDTH = 110;
|
|
const BENCH_PROMPT = benchPrompt.trim();
|
|
|
|
export interface BenchCommandArgs {
|
|
models: string[];
|
|
flags: {
|
|
runs?: number;
|
|
maxTokens?: number;
|
|
prompt?: string;
|
|
json?: boolean;
|
|
};
|
|
}
|
|
|
|
export interface BenchModelRegistry {
|
|
getAll(): Model<Api>[];
|
|
getApiKey(model: Model<Api>, sessionId?: string): Promise<string | undefined>;
|
|
resolver(model: ApiKeyResolverModel, sessionId?: string): ApiKeyResolver;
|
|
resolveCanonicalModel?(canonicalId: string, options?: CanonicalModelQueryOptions): Model<Api> | undefined;
|
|
getCanonicalVariants?(canonicalId: string, options?: CanonicalModelQueryOptions): CanonicalModelVariant[];
|
|
getCanonicalId?(model: Model<Api>): string | undefined;
|
|
}
|
|
|
|
export interface BenchRuntime {
|
|
modelRegistry: BenchModelRegistry;
|
|
settings?: Settings;
|
|
close?: () => void;
|
|
}
|
|
|
|
export interface BenchRunSuccess {
|
|
ok: true;
|
|
ttftMs: number;
|
|
durationMs: number;
|
|
outputTokens: number;
|
|
/** Generation throughput measured over the post-first-token window. */
|
|
tokensPerSecond: number;
|
|
}
|
|
|
|
export interface BenchRunFailure {
|
|
ok: false;
|
|
error: string;
|
|
}
|
|
|
|
export type BenchRunResult = BenchRunSuccess | BenchRunFailure;
|
|
|
|
export interface BenchAverages {
|
|
ttftMs: number;
|
|
durationMs: number;
|
|
outputTokens: number;
|
|
tokensPerSecond: number;
|
|
}
|
|
|
|
export interface BenchModelReport {
|
|
/** Selector as the user typed it (e.g. "opus" or "gemini-3.5:low"). */
|
|
selector: string;
|
|
/** Resolved `provider/id`. */
|
|
model: string;
|
|
/** Explicit thinking level from a `:level` selector suffix; undefined = provider default. */
|
|
thinking?: ResolvedThinkingLevel;
|
|
results: BenchRunResult[];
|
|
/** Averages over successful runs; null when every run failed. */
|
|
average: BenchAverages | null;
|
|
}
|
|
|
|
export interface BenchSummary {
|
|
runs: number;
|
|
maxTokens: number;
|
|
models: BenchModelReport[];
|
|
failures: number;
|
|
}
|
|
|
|
type BenchStreamSimple = (
|
|
model: Model<Api>,
|
|
context: Context,
|
|
options?: SimpleStreamOptions,
|
|
) => AssistantMessageEventStream;
|
|
|
|
export interface BenchDependencies {
|
|
createRuntime?: () => Promise<BenchRuntime>;
|
|
randomSessionId?: () => string;
|
|
writeStdout?: (text: string) => void;
|
|
writeStderr?: (text: string) => void;
|
|
setExitCode?: (code: number) => void;
|
|
streamSimple?: BenchStreamSimple;
|
|
now?: () => number;
|
|
stdoutIsTTY?: boolean;
|
|
}
|
|
|
|
function getErrorMessage(error: unknown): string {
|
|
if (error instanceof Error && error.message) return error.message;
|
|
return String(error);
|
|
}
|
|
|
|
function normalizePositiveInteger(name: string, value: number | undefined, fallback: number): number {
|
|
if (value === undefined) return fallback;
|
|
if (!Number.isInteger(value) || value <= 0) {
|
|
throw new Error(`Expected --${name} to be a positive integer, got ${value}`);
|
|
}
|
|
return value;
|
|
}
|
|
|
|
function isFirstTokenEvent(event: AssistantMessageEvent): boolean {
|
|
switch (event.type) {
|
|
case "text_delta":
|
|
case "thinking_delta":
|
|
case "toolcall_delta":
|
|
return event.delta.length > 0;
|
|
case "text_end":
|
|
case "thinking_end":
|
|
return event.content.length > 0;
|
|
default:
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Tokens/s over the generation window (duration minus TTFT) so queue/prefill
|
|
* latency does not dilute throughput. Falls back to total duration when the
|
|
* response arrived as a single chunk (TTFT ~ duration).
|
|
*/
|
|
export function computeTokensPerSecond(
|
|
outputTokens: number,
|
|
durationMs: number,
|
|
ttftMs: number,
|
|
deltaChunkCount: number,
|
|
): number {
|
|
const decodeMs = durationMs - ttftMs;
|
|
// Fall back to total duration when the response arrived as a single chunk/non-streaming.
|
|
const windowMs = decodeMs > 0 && deltaChunkCount >= 2 ? decodeMs : durationMs;
|
|
return windowMs > 0 ? (outputTokens * 1000) / windowMs : 0;
|
|
}
|
|
|
|
interface BenchRequestOptions {
|
|
apiKey: ApiKeyResolver;
|
|
sessionId: string;
|
|
prompt: string;
|
|
maxTokens: number;
|
|
/** Explicit effort from a `:level` selector suffix; absent = provider default. */
|
|
reasoning?: Effort;
|
|
/** Only set for an explicit `:off` suffix — some endpoints reject disablement. */
|
|
disableReasoning?: boolean;
|
|
}
|
|
|
|
async function runBenchRequest(
|
|
model: Model<Api>,
|
|
options: BenchRequestOptions,
|
|
streamFn: BenchStreamSimple,
|
|
now: () => number,
|
|
): Promise<BenchRunResult> {
|
|
const startedAt = now();
|
|
let firstTokenAt: number | undefined;
|
|
try {
|
|
const context: Context = {
|
|
// Codex's Responses endpoint 400s with "Instructions are required" when no
|
|
// system prompt is present — same guard as eval's completion bridge.
|
|
systemPrompt: ["You are a helpful assistant."],
|
|
messages: [{ role: "user", content: options.prompt, timestamp: Date.now(), attribution: "user" }],
|
|
};
|
|
const stream = streamFn(model, context, {
|
|
apiKey: options.apiKey,
|
|
sessionId: options.sessionId,
|
|
maxTokens:
|
|
model.maxTokens !== null && Number.isFinite(model.maxTokens) && model.maxTokens > 0
|
|
? Math.min(options.maxTokens, model.maxTokens)
|
|
: options.maxTokens,
|
|
reasoning: options.reasoning,
|
|
disableReasoning: options.disableReasoning,
|
|
// pi-ai opts every OpenRouter request into response caching (1h TTL).
|
|
// Bench sends a byte-identical request each run, so within the TTL
|
|
// OpenRouter replays the cached generation with zeroed usage — the run
|
|
// shows "tokens 0, TPS 0.0" at line speed. Opt back out so every run
|
|
// measures a fresh generation.
|
|
headers: model.provider === "openrouter" ? { "X-OpenRouter-Cache": "false" } : undefined,
|
|
});
|
|
let message: AssistantMessage | undefined;
|
|
let deltaChunkCount = 0;
|
|
for await (const event of stream) {
|
|
if (firstTokenAt === undefined && isFirstTokenEvent(event)) {
|
|
firstTokenAt = now();
|
|
}
|
|
if (
|
|
(event.type === "text_delta" || event.type === "thinking_delta" || event.type === "toolcall_delta") &&
|
|
event.delta.length > 0
|
|
) {
|
|
deltaChunkCount++;
|
|
}
|
|
if (event.type === "error") {
|
|
return { ok: false, error: event.error.errorMessage ?? "request failed" };
|
|
}
|
|
if (event.type === "done") {
|
|
message = event.message;
|
|
}
|
|
}
|
|
message ??= await stream.result();
|
|
if (message.stopReason === "error" || message.errorMessage) {
|
|
return { ok: false, error: message.errorMessage ?? "request failed" };
|
|
}
|
|
const rawDuration = message.duration ?? now() - startedAt;
|
|
const durationMs = Number.isFinite(rawDuration) && rawDuration > 0 ? rawDuration : 0;
|
|
const rawTtft = message.ttft ?? (firstTokenAt === undefined ? durationMs : firstTokenAt - startedAt);
|
|
const ttftMs = Number.isFinite(rawTtft) && rawTtft > 0 ? rawTtft : 0;
|
|
const outputTokens = Number.isFinite(message.usage.output) && message.usage.output > 0 ? message.usage.output : 0;
|
|
return {
|
|
ok: true,
|
|
ttftMs,
|
|
durationMs,
|
|
outputTokens,
|
|
tokensPerSecond: computeTokensPerSecond(outputTokens, durationMs, ttftMs, deltaChunkCount),
|
|
};
|
|
} catch (error) {
|
|
return { ok: false, error: getErrorMessage(error) };
|
|
}
|
|
}
|
|
|
|
function buildModelReport(
|
|
selector: string,
|
|
model: Model<Api>,
|
|
thinking: ResolvedThinkingLevel | undefined,
|
|
results: BenchRunResult[],
|
|
): BenchModelReport {
|
|
const successes = results.filter((result): result is BenchRunSuccess => result.ok);
|
|
const average =
|
|
successes.length === 0
|
|
? null
|
|
: {
|
|
ttftMs: successes.reduce((sum, r) => sum + r.ttftMs, 0) / successes.length,
|
|
durationMs: successes.reduce((sum, r) => sum + r.durationMs, 0) / successes.length,
|
|
outputTokens: successes.reduce((sum, r) => sum + r.outputTokens, 0) / successes.length,
|
|
tokensPerSecond: successes.reduce((sum, r) => sum + r.tokensPerSecond, 0) / successes.length,
|
|
};
|
|
return { selector, model: formatModelString(model), thinking, results, average };
|
|
}
|
|
|
|
function formatBenchModelLabel(report: BenchModelReport): string {
|
|
return formatModelSelectorValue(report.model, report.thinking);
|
|
}
|
|
|
|
function formatMs(ms: number): string {
|
|
return formatDuration(Math.max(0, Math.round(ms)));
|
|
}
|
|
|
|
function formatRunLine(result: BenchRunResult, index: number, total: number): string {
|
|
const prefix = chalk.dim(`run ${index + 1}/${total}`);
|
|
if (result.ok) {
|
|
return ` ${chalk.green("✓")} ${prefix} ${chalk.dim("TTFT")} ${formatMs(result.ttftMs)} ${chalk.dim("TPS")} ${result.tokensPerSecond.toFixed(1)}/s ${chalk.dim("tokens")} ${result.outputTokens} ${chalk.dim("total")} ${formatMs(result.durationMs)}`;
|
|
}
|
|
return ` ${chalk.red("✗")} ${prefix} ${chalk.red(truncateToWidth(replaceTabs(result.error).replace(/\r?\n/g, " "), ERROR_WIDTH))}`;
|
|
}
|
|
|
|
export function formatBenchTable(summary: BenchSummary): string {
|
|
const ranked = [...summary.models].sort((a, b) => {
|
|
if (a.average === null && b.average === null) return 0;
|
|
if (a.average === null) return 1;
|
|
if (b.average === null) return -1;
|
|
return b.average.tokensPerSecond - a.average.tokensPerSecond;
|
|
});
|
|
const rows = ranked.map(report => ({
|
|
model: formatBenchModelLabel(report),
|
|
ttft: report.average ? formatMs(report.average.ttftMs) : "-",
|
|
tps: report.average ? `${report.average.tokensPerSecond.toFixed(1)}/s` : "-",
|
|
tokens: report.average ? String(Math.round(report.average.outputTokens)) : "-",
|
|
total: report.average ? formatMs(report.average.durationMs) : "-",
|
|
failed: report.results.filter(result => !result.ok).length,
|
|
}));
|
|
const headers = { model: "model", ttft: "TTFT", tps: "TPS", tokens: "tokens", total: "total" } as const;
|
|
const width = (key: keyof typeof headers): number =>
|
|
Math.max(headers[key].length, ...rows.map(row => row[key].length));
|
|
const lines = [
|
|
[
|
|
headers.model.padEnd(width("model")),
|
|
headers.ttft.padEnd(width("ttft")),
|
|
headers.tps.padEnd(width("tps")),
|
|
headers.tokens.padEnd(width("tokens")),
|
|
headers.total.padEnd(width("total")),
|
|
]
|
|
.join(" ")
|
|
.trimEnd(),
|
|
];
|
|
for (const row of rows) {
|
|
const failedSuffix = row.failed > 0 ? ` ${chalk.red(`(${row.failed} failed)`)}` : "";
|
|
lines.push(
|
|
[
|
|
row.model.padEnd(width("model")),
|
|
row.ttft.padEnd(width("ttft")),
|
|
row.tps.padEnd(width("tps")),
|
|
row.tokens.padEnd(width("tokens")),
|
|
row.total.padEnd(width("total")),
|
|
]
|
|
.join(" ")
|
|
.trimEnd() + failedSuffix,
|
|
);
|
|
}
|
|
return `${lines.map((line, index) => (index === 0 ? chalk.dim(line) : line)).join("\n")}\n`;
|
|
}
|
|
|
|
async function createDefaultRuntime(): Promise<BenchRuntime> {
|
|
const authStorage = await discoverAuthStorage();
|
|
try {
|
|
const settings = await Settings.init({ cwd: getProjectDir() });
|
|
const modelRegistry = new ModelRegistry(authStorage);
|
|
return {
|
|
modelRegistry,
|
|
settings,
|
|
close: () => authStorage.close(),
|
|
};
|
|
} catch (error) {
|
|
authStorage.close();
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
interface BenchTarget {
|
|
selector: string;
|
|
model: Model<Api>;
|
|
thinking: ResolvedThinkingLevel | undefined;
|
|
}
|
|
|
|
function resolveBenchModels(
|
|
selectors: string[],
|
|
modelRegistry: BenchModelRegistry,
|
|
settings: Settings | undefined,
|
|
writeStderr: (text: string) => void,
|
|
): BenchTarget[] {
|
|
const preferences = getModelMatchPreferences(settings);
|
|
const resolved: BenchTarget[] = [];
|
|
const errors: string[] = [];
|
|
for (const selector of selectors) {
|
|
const result = resolveCliModel({ cliModel: selector, modelRegistry, preferences });
|
|
if (result.error) {
|
|
errors.push(`${selector}: ${result.error}`);
|
|
continue;
|
|
}
|
|
if (!result.model) {
|
|
errors.push(`${selector}: model not found`);
|
|
continue;
|
|
}
|
|
if (result.warning) writeStderr(`${chalk.yellow(`Warning: ${result.warning}`)}\n`);
|
|
resolved.push({
|
|
selector,
|
|
model: result.model,
|
|
thinking: resolveThinkingLevelForModel(result.model, result.thinkingLevel),
|
|
});
|
|
}
|
|
if (errors.length > 0) {
|
|
throw new Error(`Could not resolve ${errors.length === 1 ? "model" : "models"}:\n${errors.join("\n")}`);
|
|
}
|
|
return resolved;
|
|
}
|
|
|
|
export async function runBenchCommand(command: BenchCommandArgs, deps: BenchDependencies = {}): Promise<BenchSummary> {
|
|
const runs = normalizePositiveInteger("runs", command.flags.runs, DEFAULT_RUNS);
|
|
const maxTokens = normalizePositiveInteger("max-tokens", command.flags.maxTokens, DEFAULT_MAX_TOKENS);
|
|
const prompt = command.flags.prompt?.trim() || BENCH_PROMPT;
|
|
const json = command.flags.json === true;
|
|
const randomSessionId = deps.randomSessionId ?? (() => Bun.randomUUIDv7());
|
|
const writeStdout = deps.writeStdout ?? ((text: string) => process.stdout.write(text));
|
|
const writeStderr = deps.writeStderr ?? ((text: string) => process.stderr.write(text));
|
|
const setExitCode =
|
|
deps.setExitCode ??
|
|
((code: number) => {
|
|
process.exitCode = code;
|
|
});
|
|
const streamFn = deps.streamSimple ?? streamSimple;
|
|
const now = deps.now ?? (() => performance.now());
|
|
const interactive = deps.stdoutIsTTY ?? process.stdout.isTTY === true;
|
|
if (command.models.length === 0) {
|
|
throw new Error("Pass at least one model selector, e.g. `omp bench opus gpt-5.2`");
|
|
}
|
|
|
|
const runtime = await (deps.createRuntime ?? createDefaultRuntime)();
|
|
try {
|
|
const targets = resolveBenchModels(command.models, runtime.modelRegistry, runtime.settings, writeStderr);
|
|
const reports: BenchModelReport[] = [];
|
|
for (const { selector, model, thinking } of targets) {
|
|
if (!json) {
|
|
const resolvedModel = formatModelSelectorValue(formatModelString(model), thinking);
|
|
const resolvedNote = selector === resolvedModel ? "" : chalk.dim(` (${selector})`);
|
|
writeStdout(`${chalk.bold(resolvedModel)}${resolvedNote}\n`);
|
|
}
|
|
const results: BenchRunResult[] = [];
|
|
for (let index = 0; index < runs; index++) {
|
|
const sessionId = randomSessionId();
|
|
const initialKey = await runtime.modelRegistry.getApiKey(model, sessionId);
|
|
if (!initialKey) {
|
|
const failure: BenchRunFailure = {
|
|
ok: false,
|
|
error: `No credentials for provider "${model.provider}". Run \`omp\` and use /login, or set the provider API key.`,
|
|
};
|
|
results.push(failure);
|
|
if (!json) writeStdout(`${formatRunLine(failure, index, runs)}\n`);
|
|
break; // remaining runs would fail identically
|
|
}
|
|
if (!json && interactive) {
|
|
writeStdout(chalk.dim(` … run ${index + 1}/${runs} streaming`));
|
|
}
|
|
const result = await runBenchRequest(
|
|
model,
|
|
{
|
|
apiKey: runtime.modelRegistry.resolver(model, sessionId),
|
|
sessionId,
|
|
prompt,
|
|
maxTokens,
|
|
reasoning: toReasoningEffort(thinking),
|
|
disableReasoning: shouldDisableReasoning(thinking) ? true : undefined,
|
|
},
|
|
streamFn,
|
|
now,
|
|
);
|
|
results.push(result);
|
|
if (!json) {
|
|
if (interactive) writeStdout("\r\x1b[2K");
|
|
writeStdout(`${formatRunLine(result, index, runs)}\n`);
|
|
}
|
|
}
|
|
reports.push(buildModelReport(selector, model, thinking, results));
|
|
}
|
|
const failures = reports.reduce((sum, report) => sum + report.results.filter(result => !result.ok).length, 0);
|
|
const summary: BenchSummary = { runs, maxTokens, models: reports, failures };
|
|
if (json) {
|
|
writeStdout(`${JSON.stringify(summary, null, 2)}\n`);
|
|
} else if (reports.length > 1 || runs > 1) {
|
|
writeStdout(`\n${formatBenchTable(summary)}`);
|
|
}
|
|
if (failures > 0) setExitCode(1);
|
|
return summary;
|
|
} finally {
|
|
runtime.close?.();
|
|
}
|
|
}
|