Merge remote-tracking branch 'upstream/main' into feat/profiles-and-alias
This commit is contained in:
Ogrodev
@@ -0,0 +1,823 @@
|
||||
import type {
|
||||
Api,
|
||||
AssistantMessage,
|
||||
AssistantMessageEvent,
|
||||
AssistantMessageEventStream,
|
||||
Context,
|
||||
Model,
|
||||
OAuthAccess,
|
||||
OAuthAccessResolution,
|
||||
SimpleStreamOptions,
|
||||
} from "@oh-my-pi/pi-ai";
|
||||
import { streamSimple } from "@oh-my-pi/pi-ai";
|
||||
import { replaceTabs, truncateToWidth } from "@oh-my-pi/pi-tui";
|
||||
import { formatDuration, getProjectDir } from "@oh-my-pi/pi-utils";
|
||||
import chalk from "chalk";
|
||||
import type { CanonicalModelVariant } from "../config/model-equivalence";
|
||||
import { type CanonicalModelQueryOptions, ModelRegistry } from "../config/model-registry";
|
||||
import {
|
||||
formatModelString,
|
||||
type ModelMatchPreferences,
|
||||
resolveAllowedModels,
|
||||
resolveCliModel,
|
||||
resolveModelRoleValue,
|
||||
} from "../config/model-resolver";
|
||||
import { Settings } from "../config/settings";
|
||||
import dryBalanceBenchPrompt from "../prompts/dry-balance-bench.md" with { type: "text" };
|
||||
import { discoverAuthStorage } from "../sdk";
|
||||
|
||||
const DEFAULT_SAMPLE_COUNT = 100;
|
||||
const DEFAULT_CONCURRENCY = 32;
|
||||
const BENCH_MAX_TOKENS = 512;
|
||||
const BENCH_RENDER_INTERVAL_MS = 80;
|
||||
const BENCH_ACCOUNT_WIDTH = 60;
|
||||
const BENCH_ERROR_WIDTH = 110;
|
||||
const BENCH_SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"] as const;
|
||||
const DRY_BALANCE_BENCH_PROMPT = dryBalanceBenchPrompt.trim();
|
||||
|
||||
export interface DryBalanceCommandArgs {
|
||||
model?: string;
|
||||
flags: {
|
||||
model?: string;
|
||||
count?: number;
|
||||
concurrency?: number;
|
||||
json?: boolean;
|
||||
bench?: boolean;
|
||||
};
|
||||
}
|
||||
|
||||
export interface DryBalanceAuthOptions {
|
||||
baseUrl?: string;
|
||||
modelId?: string;
|
||||
signal?: AbortSignal;
|
||||
}
|
||||
|
||||
export interface DryBalanceAuthStorage {
|
||||
getOAuthAccess(
|
||||
provider: string,
|
||||
sessionId?: string,
|
||||
options?: DryBalanceAuthOptions,
|
||||
): Promise<OAuthAccess | undefined>;
|
||||
getOAuthAccesses?(provider: string, options?: DryBalanceAuthOptions): Promise<OAuthAccessResolution[]>;
|
||||
}
|
||||
|
||||
export interface DryBalanceModelRegistry {
|
||||
authStorage: DryBalanceAuthStorage;
|
||||
getAll(): Model<Api>[];
|
||||
getAvailable(): Model<Api>[];
|
||||
getApiKey(model: Model<Api>, sessionId?: string): Promise<string | undefined>;
|
||||
getCanonicalVariants(canonicalId: string, options?: CanonicalModelQueryOptions): CanonicalModelVariant[];
|
||||
resolveCanonicalModel?(canonicalId: string, options?: CanonicalModelQueryOptions): Model<Api> | undefined;
|
||||
getCanonicalId?(model: Model<Api>): string | undefined;
|
||||
}
|
||||
|
||||
export interface DryBalanceRuntime {
|
||||
modelRegistry: DryBalanceModelRegistry;
|
||||
settings?: Settings;
|
||||
close?: () => void;
|
||||
}
|
||||
|
||||
export interface DryBalanceAccountStat {
|
||||
account: string;
|
||||
count: number;
|
||||
percent: number;
|
||||
}
|
||||
|
||||
export interface DryBalanceFailureStat {
|
||||
reason: string;
|
||||
count: number;
|
||||
percent: number;
|
||||
}
|
||||
|
||||
export interface DryBalanceBenchSuccessResult {
|
||||
ok: true;
|
||||
account: string;
|
||||
ttftMs: number;
|
||||
durationMs: number;
|
||||
outputTokens: number;
|
||||
tokensPerSecond: number;
|
||||
}
|
||||
|
||||
export interface DryBalanceBenchFailureResult {
|
||||
ok: false;
|
||||
account?: string;
|
||||
error: string;
|
||||
}
|
||||
|
||||
export type DryBalanceBenchResult = DryBalanceBenchSuccessResult | DryBalanceBenchFailureResult;
|
||||
|
||||
export interface DryBalanceBenchSummary {
|
||||
total: number;
|
||||
success: {
|
||||
total: number;
|
||||
averageTtftMs: number | null;
|
||||
averageTokensPerSecond: number | null;
|
||||
};
|
||||
failure: {
|
||||
total: number;
|
||||
reasons: DryBalanceFailureStat[];
|
||||
};
|
||||
results: DryBalanceBenchResult[];
|
||||
}
|
||||
|
||||
export interface DryBalanceSummary {
|
||||
model: string;
|
||||
provider: string;
|
||||
samples: number;
|
||||
concurrency: number;
|
||||
success: {
|
||||
total: number;
|
||||
accounts: DryBalanceAccountStat[];
|
||||
};
|
||||
failure: {
|
||||
total: number;
|
||||
reasons: DryBalanceFailureStat[];
|
||||
};
|
||||
bench?: DryBalanceBenchSummary;
|
||||
}
|
||||
|
||||
type DryBalanceStreamSimple = (
|
||||
model: Model<Api>,
|
||||
context: Context,
|
||||
options?: SimpleStreamOptions,
|
||||
) => AssistantMessageEventStream;
|
||||
|
||||
export interface DryBalanceDependencies {
|
||||
createRuntime?: () => Promise<DryBalanceRuntime>;
|
||||
randomSessionId?: () => string;
|
||||
writeStdout?: (text: string) => void;
|
||||
writeStderr?: (text: string) => void;
|
||||
setExitCode?: (code: number) => void;
|
||||
streamSimple?: DryBalanceStreamSimple;
|
||||
now?: () => number;
|
||||
stdoutIsTTY?: boolean;
|
||||
stderrIsTTY?: boolean;
|
||||
}
|
||||
|
||||
type DryBalanceAttemptResult =
|
||||
| {
|
||||
ok: true;
|
||||
account: string;
|
||||
}
|
||||
| {
|
||||
ok: false;
|
||||
reason: string;
|
||||
};
|
||||
|
||||
type DryBalanceBenchProgressStatus =
|
||||
| { state: "waiting" }
|
||||
| { state: "running"; account: string }
|
||||
| { state: "success"; result: DryBalanceBenchSuccessResult }
|
||||
| { state: "failure"; result: DryBalanceBenchFailureResult };
|
||||
|
||||
interface DryBalanceBenchProgressSink {
|
||||
markRunning(index: number, account: string): void;
|
||||
complete(index: number, result: DryBalanceBenchResult): void;
|
||||
close(): void;
|
||||
}
|
||||
|
||||
type DryBalanceBenchTarget =
|
||||
| {
|
||||
ok: true;
|
||||
account: string;
|
||||
accessToken: string;
|
||||
}
|
||||
| {
|
||||
ok: false;
|
||||
account: string;
|
||||
error: string;
|
||||
};
|
||||
|
||||
function normalizePositiveInteger(name: string, value: number | undefined, fallback: number): number {
|
||||
const resolved = value ?? fallback;
|
||||
if (!Number.isInteger(resolved) || resolved <= 0) {
|
||||
throw new Error(`--${name} must be a positive integer`);
|
||||
}
|
||||
return resolved;
|
||||
}
|
||||
|
||||
function getErrorMessage(error: unknown): string {
|
||||
if (error instanceof Error && error.message) return error.message;
|
||||
const message = String(error);
|
||||
return message ? message : "Unknown error";
|
||||
}
|
||||
|
||||
function extractAccount(access: {
|
||||
email?: string;
|
||||
accountId?: string;
|
||||
projectId?: string;
|
||||
enterpriseUrl?: string;
|
||||
}): string {
|
||||
return access.email ?? access.accountId ?? access.projectId ?? access.enterpriseUrl ?? "(unknown oauth account)";
|
||||
}
|
||||
|
||||
function getBenchTargetKey(access: {
|
||||
credentialId?: number;
|
||||
email?: string;
|
||||
accountId?: string;
|
||||
projectId?: string;
|
||||
enterpriseUrl?: string;
|
||||
accessToken?: string;
|
||||
}): string {
|
||||
return (
|
||||
access.email ??
|
||||
access.accountId ??
|
||||
access.projectId ??
|
||||
access.enterpriseUrl ??
|
||||
(access.credentialId === undefined ? access.accessToken : `credential:${access.credentialId}`) ??
|
||||
"(unknown oauth account)"
|
||||
);
|
||||
}
|
||||
|
||||
function sanitizeBenchText(text: string, width: number): string {
|
||||
return truncateToWidth(replaceTabs(text).replace(/\r?\n/g, " "), width);
|
||||
}
|
||||
|
||||
function formatBenchIndex(index: number, total: number): string {
|
||||
return `#${String(index + 1).padStart(String(total).length, "0")}`;
|
||||
}
|
||||
|
||||
function formatBenchAccount(account: string | undefined): string {
|
||||
return account ? sanitizeBenchText(account, BENCH_ACCOUNT_WIDTH) : chalk.dim("(no account)");
|
||||
}
|
||||
|
||||
function formatBenchDuration(ms: number): string {
|
||||
return formatDuration(Math.max(0, Math.round(ms)));
|
||||
}
|
||||
|
||||
function formatBenchTps(tokensPerSecond: number): string {
|
||||
return `${tokensPerSecond.toFixed(1)}/s`;
|
||||
}
|
||||
|
||||
function isBenchSuccess(result: DryBalanceBenchResult): result is DryBalanceBenchSuccessResult {
|
||||
return result.ok;
|
||||
}
|
||||
|
||||
function isBenchFirstTokenEvent(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;
|
||||
}
|
||||
}
|
||||
|
||||
function resolveBenchMaxTokens(model: Model<Api>): number {
|
||||
return Number.isFinite(model.maxTokens) && model.maxTokens > 0
|
||||
? Math.min(BENCH_MAX_TOKENS, model.maxTokens)
|
||||
: BENCH_MAX_TOKENS;
|
||||
}
|
||||
|
||||
function normalizeBenchMs(value: number): number {
|
||||
return Number.isFinite(value) && value > 0 ? value : 0;
|
||||
}
|
||||
|
||||
function renderBenchResultLine(index: number, total: number, result: DryBalanceBenchResult): string {
|
||||
const prefix = formatBenchIndex(index, total);
|
||||
if (result.ok) {
|
||||
return `${chalk.green("✓")} ${prefix} ${formatBenchAccount(result.account)} ${chalk.dim("TTFT")} ${formatBenchDuration(
|
||||
result.ttftMs,
|
||||
)} ${chalk.dim("TPS")} ${formatBenchTps(result.tokensPerSecond)}`;
|
||||
}
|
||||
return `${chalk.red("✗")} ${prefix} ${formatBenchAccount(result.account)} ${chalk.red(
|
||||
sanitizeBenchText(result.error, BENCH_ERROR_WIDTH),
|
||||
)}`;
|
||||
}
|
||||
|
||||
function renderBenchStatusLine(
|
||||
status: DryBalanceBenchProgressStatus,
|
||||
index: number,
|
||||
total: number,
|
||||
frame: number,
|
||||
): string {
|
||||
const prefix = formatBenchIndex(index, total);
|
||||
switch (status.state) {
|
||||
case "waiting":
|
||||
return `${chalk.dim("○")} ${prefix} ${chalk.dim("waiting")}`;
|
||||
case "running": {
|
||||
const spinner = BENCH_SPINNER_FRAMES[frame % BENCH_SPINNER_FRAMES.length] ?? "*";
|
||||
return `${chalk.yellow(spinner)} ${prefix} ${formatBenchAccount(status.account)} ${chalk.dim("sending request")}`;
|
||||
}
|
||||
case "success":
|
||||
return renderBenchResultLine(index, total, status.result);
|
||||
case "failure":
|
||||
return renderBenchResultLine(index, total, status.result);
|
||||
}
|
||||
}
|
||||
|
||||
function createBenchProgressSink(
|
||||
total: number,
|
||||
write: (text: string) => void,
|
||||
interactive: boolean,
|
||||
): DryBalanceBenchProgressSink {
|
||||
const statuses: DryBalanceBenchProgressStatus[] = Array.from({ length: total }, () => ({ state: "waiting" }));
|
||||
if (!interactive) {
|
||||
return {
|
||||
markRunning(index, account) {
|
||||
statuses[index] = { state: "running", account };
|
||||
write(`${renderBenchStatusLine(statuses[index], index, total, 0)}\n`);
|
||||
},
|
||||
complete(index, result) {
|
||||
statuses[index] = result.ok ? { state: "success", result } : { state: "failure", result };
|
||||
write(`${renderBenchResultLine(index, total, result)}\n`);
|
||||
},
|
||||
close() {},
|
||||
};
|
||||
}
|
||||
|
||||
let frame = 0;
|
||||
let lineCount = 0;
|
||||
let timer: NodeJS.Timeout | undefined;
|
||||
const render = (): void => {
|
||||
const lines = [
|
||||
chalk.bold("bench requests"),
|
||||
...statuses.map((status, index) => renderBenchStatusLine(status, index, total, frame)),
|
||||
];
|
||||
if (lineCount > 0) write(`\x1b[${lineCount}A`);
|
||||
write(`${lines.map(line => `\x1b[2K${line}`).join("\n")}\n`);
|
||||
lineCount = lines.length;
|
||||
};
|
||||
render();
|
||||
timer = setInterval(() => {
|
||||
frame += 1;
|
||||
render();
|
||||
}, BENCH_RENDER_INTERVAL_MS);
|
||||
timer.unref?.();
|
||||
return {
|
||||
markRunning(index, account) {
|
||||
statuses[index] = { state: "running", account };
|
||||
render();
|
||||
},
|
||||
complete(index, result) {
|
||||
statuses[index] = result.ok ? { state: "success", result } : { state: "failure", result };
|
||||
render();
|
||||
},
|
||||
close() {
|
||||
if (timer) {
|
||||
clearInterval(timer);
|
||||
timer = undefined;
|
||||
}
|
||||
render();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function runBenchRequest(
|
||||
model: Model<Api>,
|
||||
sessionId: string,
|
||||
account: string,
|
||||
accessToken: string,
|
||||
streamFn: DryBalanceStreamSimple,
|
||||
now: () => number,
|
||||
): Promise<DryBalanceBenchResult> {
|
||||
const startedAt = now();
|
||||
let firstTokenAt: number | undefined;
|
||||
try {
|
||||
const context: Context = {
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: DRY_BALANCE_BENCH_PROMPT,
|
||||
timestamp: Date.now(),
|
||||
attribution: "user",
|
||||
},
|
||||
],
|
||||
};
|
||||
const stream = streamFn(model, context, {
|
||||
apiKey: accessToken,
|
||||
sessionId,
|
||||
maxTokens: resolveBenchMaxTokens(model),
|
||||
temperature: 0.2,
|
||||
disableReasoning: true,
|
||||
hideThinkingSummary: true,
|
||||
});
|
||||
let message: AssistantMessage | undefined;
|
||||
for await (const event of stream) {
|
||||
if (firstTokenAt === undefined && isBenchFirstTokenEvent(event)) {
|
||||
firstTokenAt = now();
|
||||
}
|
||||
if (event.type === "error") {
|
||||
return { ok: false, account, 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, account, error: message.errorMessage ?? "request failed" };
|
||||
}
|
||||
const durationMs = normalizeBenchMs(message.duration ?? now() - startedAt);
|
||||
const ttftMs = normalizeBenchMs(
|
||||
message.ttft ?? (firstTokenAt === undefined ? durationMs : firstTokenAt - startedAt),
|
||||
);
|
||||
const outputTokens = Number.isFinite(message.usage.output) && message.usage.output > 0 ? message.usage.output : 0;
|
||||
const tokensPerSecond = durationMs > 0 ? (outputTokens * 1000) / durationMs : 0;
|
||||
return {
|
||||
ok: true,
|
||||
account,
|
||||
ttftMs,
|
||||
durationMs,
|
||||
outputTokens,
|
||||
tokensPerSecond,
|
||||
};
|
||||
} catch (error) {
|
||||
return { ok: false, account, error: getErrorMessage(error) };
|
||||
}
|
||||
}
|
||||
|
||||
async function resolveBenchTargets(
|
||||
model: Model<Api>,
|
||||
authStorage: DryBalanceAuthStorage,
|
||||
): Promise<DryBalanceBenchTarget[]> {
|
||||
const resolved = authStorage.getOAuthAccesses
|
||||
? await authStorage.getOAuthAccesses(model.provider, {
|
||||
baseUrl: model.baseUrl,
|
||||
modelId: model.id,
|
||||
})
|
||||
: await authStorage
|
||||
.getOAuthAccess(model.provider, undefined, {
|
||||
baseUrl: model.baseUrl,
|
||||
modelId: model.id,
|
||||
})
|
||||
.then(access => (access ? [{ ok: true as const, ...access }] : []));
|
||||
const targets: DryBalanceBenchTarget[] = [];
|
||||
const seen = new Set<string>();
|
||||
for (const entry of resolved) {
|
||||
const key = getBenchTargetKey(entry);
|
||||
if (seen.has(key)) continue;
|
||||
seen.add(key);
|
||||
const account = extractAccount(entry);
|
||||
if (entry.ok) {
|
||||
targets.push({ ok: true, account, accessToken: entry.accessToken });
|
||||
} else {
|
||||
targets.push({ ok: false, account, error: entry.error });
|
||||
}
|
||||
}
|
||||
return targets;
|
||||
}
|
||||
|
||||
async function runBenchTargets(
|
||||
model: Model<Api>,
|
||||
targets: DryBalanceBenchTarget[],
|
||||
randomSessionId: () => string,
|
||||
progress: DryBalanceBenchProgressSink | undefined,
|
||||
streamFn: DryBalanceStreamSimple,
|
||||
now: () => number,
|
||||
): Promise<DryBalanceBenchResult[]> {
|
||||
return Promise.all(
|
||||
targets.map(async (target, index) => {
|
||||
if (!target.ok) {
|
||||
const result: DryBalanceBenchFailureResult = {
|
||||
ok: false,
|
||||
account: target.account,
|
||||
error: target.error,
|
||||
};
|
||||
progress?.complete(index, result);
|
||||
return result;
|
||||
}
|
||||
progress?.markRunning(index, target.account);
|
||||
const result = await runBenchRequest(
|
||||
model,
|
||||
randomSessionId(),
|
||||
target.account,
|
||||
target.accessToken,
|
||||
streamFn,
|
||||
now,
|
||||
);
|
||||
progress?.complete(index, result);
|
||||
return result;
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
async function createDefaultRuntime(): Promise<DryBalanceRuntime> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
|
||||
async function resolveDryBalanceModel(
|
||||
modelSelector: string | undefined,
|
||||
modelRegistry: DryBalanceModelRegistry,
|
||||
settings: Settings | undefined,
|
||||
randomSessionId: () => string,
|
||||
): Promise<{ model: Model<Api>; warning?: string }> {
|
||||
const preferences: ModelMatchPreferences = {
|
||||
usageOrder: settings?.getStorage()?.getModelUsageOrder(),
|
||||
};
|
||||
if (modelSelector) {
|
||||
const resolved = resolveCliModel({
|
||||
cliModel: modelSelector,
|
||||
modelRegistry,
|
||||
preferences,
|
||||
});
|
||||
if (resolved.error) throw new Error(resolved.error);
|
||||
if (!resolved.model) throw new Error(`Model "${modelSelector}" not found`);
|
||||
return { model: resolved.model, warning: resolved.warning };
|
||||
}
|
||||
|
||||
const allowedModels = await resolveAllowedModels(modelRegistry, settings, preferences);
|
||||
if (allowedModels.length === 0) {
|
||||
throw new Error(
|
||||
"No models available. Use --model to select a model or configure enabledModels/default model settings.",
|
||||
);
|
||||
}
|
||||
|
||||
const defaultRoleSpec = resolveModelRoleValue(settings?.getModelRole("default"), allowedModels, {
|
||||
settings,
|
||||
matchPreferences: preferences,
|
||||
modelRegistry,
|
||||
});
|
||||
if (defaultRoleSpec.model) {
|
||||
return { model: defaultRoleSpec.model, warning: defaultRoleSpec.warning };
|
||||
}
|
||||
|
||||
for (const candidate of allowedModels) {
|
||||
const apiKey = await modelRegistry.getApiKey(candidate, randomSessionId());
|
||||
if (apiKey) return { model: candidate };
|
||||
}
|
||||
|
||||
return {
|
||||
model: allowedModels[0],
|
||||
warning:
|
||||
"No allowed model had usable credentials during default resolution; dry-balance will report OAuth failures for the first allowed model.",
|
||||
};
|
||||
}
|
||||
|
||||
async function runOneAttempt(
|
||||
model: Model<Api>,
|
||||
modelRegistry: DryBalanceModelRegistry,
|
||||
sessionId: string,
|
||||
): Promise<DryBalanceAttemptResult> {
|
||||
try {
|
||||
// AuthStorage.getOAuthAccess shares the OAuth credential ranking, refresh,
|
||||
// usage-limit, broker, and session-sticky path used by getApiKey(), while
|
||||
// returning the selected account metadata instead of bearer bytes.
|
||||
const access = await modelRegistry.authStorage.getOAuthAccess(model.provider, sessionId, {
|
||||
baseUrl: model.baseUrl,
|
||||
modelId: model.id,
|
||||
});
|
||||
if (!access) return { ok: false, reason: "no OAuth access resolved" };
|
||||
return { ok: true, account: extractAccount(access) };
|
||||
} catch (error) {
|
||||
return { ok: false, reason: getErrorMessage(error) };
|
||||
}
|
||||
}
|
||||
|
||||
async function mapConcurrent<T, R>(
|
||||
items: T[],
|
||||
concurrency: number,
|
||||
fn: (item: T, index: number) => Promise<R>,
|
||||
): Promise<R[]> {
|
||||
const results = new Array<R>(items.length);
|
||||
let nextIndex = 0;
|
||||
const workerCount = Math.min(concurrency, items.length);
|
||||
await Promise.all(
|
||||
Array.from({ length: workerCount }, async () => {
|
||||
while (true) {
|
||||
const index = nextIndex;
|
||||
nextIndex += 1;
|
||||
if (index >= items.length) return;
|
||||
results[index] = await fn(items[index], index);
|
||||
}
|
||||
}),
|
||||
);
|
||||
return results;
|
||||
}
|
||||
|
||||
function sortedStats(
|
||||
map: Map<string, number>,
|
||||
samples: number,
|
||||
): Array<{ label: string; count: number; percent: number }> {
|
||||
return [...map.entries()]
|
||||
.map(([label, count]) => ({ label, count, percent: (count / samples) * 100 }))
|
||||
.sort((left, right) => right.count - left.count || left.label.localeCompare(right.label));
|
||||
}
|
||||
|
||||
function summarizeBenchResults(results: DryBalanceBenchResult[]): DryBalanceBenchSummary | undefined {
|
||||
if (results.length === 0) return undefined;
|
||||
const successes = results.filter(isBenchSuccess);
|
||||
const failureReasons = new Map<string, number>();
|
||||
for (const result of results) {
|
||||
if (!result.ok) {
|
||||
failureReasons.set(result.error, (failureReasons.get(result.error) ?? 0) + 1);
|
||||
}
|
||||
}
|
||||
const average = (values: number[]): number | null =>
|
||||
values.length === 0 ? null : values.reduce((sum, value) => sum + value, 0) / values.length;
|
||||
return {
|
||||
total: results.length,
|
||||
success: {
|
||||
total: successes.length,
|
||||
averageTtftMs: average(successes.map(result => result.ttftMs)),
|
||||
averageTokensPerSecond: average(successes.map(result => result.tokensPerSecond)),
|
||||
},
|
||||
failure: {
|
||||
total: results.length - successes.length,
|
||||
reasons: sortedStats(failureReasons, results.length).map(stat => ({
|
||||
reason: stat.label,
|
||||
count: stat.count,
|
||||
percent: stat.percent,
|
||||
})),
|
||||
},
|
||||
results,
|
||||
};
|
||||
}
|
||||
|
||||
function summarizeResults(
|
||||
model: Model<Api>,
|
||||
samples: number,
|
||||
concurrency: number,
|
||||
results: DryBalanceAttemptResult[],
|
||||
): DryBalanceSummary {
|
||||
const accounts = new Map<string, number>();
|
||||
const reasons = new Map<string, number>();
|
||||
for (const result of results) {
|
||||
if (result.ok) {
|
||||
accounts.set(result.account, (accounts.get(result.account) ?? 0) + 1);
|
||||
} else {
|
||||
reasons.set(result.reason, (reasons.get(result.reason) ?? 0) + 1);
|
||||
}
|
||||
}
|
||||
const accountStats: DryBalanceAccountStat[] = sortedStats(accounts, samples).map(stat => ({
|
||||
account: stat.label,
|
||||
count: stat.count,
|
||||
percent: stat.percent,
|
||||
}));
|
||||
const failureStats: DryBalanceFailureStat[] = sortedStats(reasons, samples).map(stat => ({
|
||||
reason: stat.label,
|
||||
count: stat.count,
|
||||
percent: stat.percent,
|
||||
}));
|
||||
const summary: DryBalanceSummary = {
|
||||
model: formatModelString(model),
|
||||
provider: model.provider,
|
||||
samples,
|
||||
concurrency,
|
||||
success: {
|
||||
total: results.filter(result => result.ok).length,
|
||||
accounts: accountStats,
|
||||
},
|
||||
failure: {
|
||||
total: results.filter(result => !result.ok).length,
|
||||
reasons: failureStats,
|
||||
},
|
||||
};
|
||||
return summary;
|
||||
}
|
||||
|
||||
function formatRows(rows: Array<{ count: number; percent: number; label: string }>): string[] {
|
||||
if (rows.length === 0) return [` ${chalk.dim("(none)")}`];
|
||||
const maxCountWidth = Math.max(...rows.map(row => row.count.toString().length));
|
||||
return rows.map(row => {
|
||||
const count = row.count.toString().padStart(maxCountWidth);
|
||||
const percent = `${row.percent.toFixed(1)}%`.padStart(6);
|
||||
return ` ${count} ${percent} ${row.label}`;
|
||||
});
|
||||
}
|
||||
|
||||
export function formatDryBalanceText(summary: DryBalanceSummary): string {
|
||||
const accountRows = summary.success.accounts.map(row => ({
|
||||
count: row.count,
|
||||
percent: row.percent,
|
||||
label: row.account,
|
||||
}));
|
||||
const failureRows = summary.failure.reasons.map(row => ({
|
||||
count: row.count,
|
||||
percent: row.percent,
|
||||
label: row.reason,
|
||||
}));
|
||||
const lines = [
|
||||
chalk.bold("dry-balance"),
|
||||
`model: ${summary.model}`,
|
||||
`provider: ${summary.provider}`,
|
||||
`samples: ${summary.samples}`,
|
||||
`concurrency: ${summary.concurrency}`,
|
||||
"",
|
||||
`${chalk.green("success")} ${summary.success.total}`,
|
||||
...formatRows(accountRows),
|
||||
"",
|
||||
`${summary.failure.total > 0 ? chalk.red("failure") : chalk.dim("failure")} ${summary.failure.total}`,
|
||||
...formatRows(failureRows),
|
||||
];
|
||||
if (summary.bench) {
|
||||
const avgTtft =
|
||||
summary.bench.success.averageTtftMs === null ? "-" : formatBenchDuration(summary.bench.success.averageTtftMs);
|
||||
const avgTps =
|
||||
summary.bench.success.averageTokensPerSecond === null
|
||||
? "-"
|
||||
: formatBenchTps(summary.bench.success.averageTokensPerSecond);
|
||||
const benchFailureRows = summary.bench.failure.reasons.map(row => ({
|
||||
count: row.count,
|
||||
percent: row.percent,
|
||||
label: row.reason,
|
||||
}));
|
||||
lines.push(
|
||||
"",
|
||||
chalk.bold("bench"),
|
||||
`requests: ${summary.bench.total}`,
|
||||
`${chalk.green("success")} ${summary.bench.success.total}`,
|
||||
`avg TTFT: ${avgTtft}`,
|
||||
`avg TPS: ${avgTps}`,
|
||||
"",
|
||||
`${summary.bench.failure.total > 0 ? chalk.red("failure") : chalk.dim("failure")} ${summary.bench.failure.total}`,
|
||||
...formatRows(benchFailureRows),
|
||||
);
|
||||
}
|
||||
return `${lines.join("\n")}\n`;
|
||||
}
|
||||
|
||||
export async function runDryBalanceCommand(
|
||||
command: DryBalanceCommandArgs,
|
||||
deps: DryBalanceDependencies = {},
|
||||
): Promise<DryBalanceSummary> {
|
||||
const isBench = command.flags.bench === true;
|
||||
const samples = isBench ? 0 : normalizePositiveInteger("count", command.flags.count, DEFAULT_SAMPLE_COUNT);
|
||||
const concurrency = isBench
|
||||
? 0
|
||||
: Math.min(samples, normalizePositiveInteger("concurrency", command.flags.concurrency, DEFAULT_CONCURRENCY));
|
||||
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 runtime = await (deps.createRuntime ?? createDefaultRuntime)();
|
||||
let progress: DryBalanceBenchProgressSink | undefined;
|
||||
let progressClosed = false;
|
||||
const closeProgress = (): void => {
|
||||
if (progressClosed) return;
|
||||
progressClosed = true;
|
||||
progress?.close();
|
||||
};
|
||||
try {
|
||||
const modelSelector = command.flags.model ?? command.model;
|
||||
const { model, warning } = await resolveDryBalanceModel(
|
||||
modelSelector,
|
||||
runtime.modelRegistry,
|
||||
runtime.settings,
|
||||
randomSessionId,
|
||||
);
|
||||
if (warning) writeStderr(`${chalk.yellow(`Warning: ${warning}`)}\n`);
|
||||
let results: DryBalanceAttemptResult[];
|
||||
let benchResults: DryBalanceBenchResult[] | undefined;
|
||||
let summarySamples = samples;
|
||||
let summaryConcurrency = concurrency;
|
||||
if (isBench) {
|
||||
const targets = await resolveBenchTargets(model, runtime.modelRegistry.authStorage);
|
||||
if (targets.length === 0) throw new Error(`No OAuth accounts resolved for provider ${model.provider}`);
|
||||
summarySamples = targets.length;
|
||||
summaryConcurrency = targets.length;
|
||||
const progressWrite = command.flags.json ? writeStderr : writeStdout;
|
||||
const progressInteractive = command.flags.json
|
||||
? (deps.stderrIsTTY ?? process.stderr.isTTY === true)
|
||||
: (deps.stdoutIsTTY ?? process.stdout.isTTY === true);
|
||||
progress = createBenchProgressSink(targets.length, progressWrite, progressInteractive);
|
||||
benchResults = await runBenchTargets(model, targets, randomSessionId, progress, streamFn, now);
|
||||
results = targets.map(target =>
|
||||
target.ok ? { ok: true, account: target.account } : { ok: false, reason: target.error },
|
||||
);
|
||||
} else {
|
||||
const sessionIds = Array.from({ length: samples }, () => randomSessionId());
|
||||
results = await mapConcurrent(sessionIds, concurrency, sessionId =>
|
||||
runOneAttempt(model, runtime.modelRegistry, sessionId),
|
||||
);
|
||||
}
|
||||
closeProgress();
|
||||
const summary = summarizeResults(model, summarySamples, summaryConcurrency, results);
|
||||
if (benchResults) {
|
||||
const benchSummary = summarizeBenchResults(benchResults);
|
||||
if (benchSummary) summary.bench = benchSummary;
|
||||
}
|
||||
if (command.flags.json) {
|
||||
writeStdout(`${JSON.stringify(summary, null, 2)}\n`);
|
||||
} else {
|
||||
writeStdout(formatDryBalanceText(summary));
|
||||
}
|
||||
if (summary.failure.total > 0 || (summary.bench?.failure.total ?? 0) > 0) setExitCode(1);
|
||||
return summary;
|
||||
} finally {
|
||||
closeProgress();
|
||||
runtime.close?.();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user