Squashed 'third_party/pi/' content from commit d981de122

git-subtree-dir: third_party/pi
git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477

packages/evals/src/vitest-evals/reporter.ts

@@ -0,0 +1,111 @@
+import { randomUUID } from "node:crypto";
+import { appendFile, mkdir } from "node:fs/promises";
+import { join } from "node:path";
+import type { Reporter, SerializedError, TestCase, TestModule, TestRunEndReason, Vitest } from "vitest/node";
+import { isHarnessRun } from "vitest-evals/harness";
+import { PI_SESSION_SNAPSHOT_ARTIFACT, persistEvalArtifactReferences } from "./artifacts.ts";
+import { EVAL_HARNESS_ITERATION_ARTIFACT, parseEvalHarnessIterationArtifact } from "./harness-table.ts";
+import { formatHarnessComparisonReport, type HarnessObservation, summarizeHarnessComparisons } from "./summary.ts";
+
+function readFiniteNumber(value: unknown): number | undefined {
+	return typeof value === "number" && Number.isFinite(value) ? value : undefined;
+}
+
+async function appendHarnessRunReport(test: TestCase): Promise<void> {
+	const artifactDirectory = process.env.PI_EVAL_ARTIFACT_DIR?.trim();
+	if (!artifactDirectory) return;
+	const harness = test.meta().harness;
+	if (!harness || !isHarnessRun(harness.run)) return;
+
+	const run = harness.run;
+	const artifactRunId = run.artifacts?.runId;
+	const runId = typeof artifactRunId === "string" ? artifactRunId : randomUUID();
+	const metadata = Object.fromEntries(
+		Object.entries(run.artifacts ?? {}).filter(([name]) => name !== "runId" && name !== PI_SESSION_SNAPSHOT_ARTIFACT),
+	);
+	const record = {
+		schemaVersion: 1,
+		runId,
+		test: {
+			id: test.id,
+			file: test.module.relativeModuleId,
+			name: test.name,
+			fullName: test.fullName,
+			status: test.result().state,
+		},
+		harness: harness.name,
+		usage: run.usage,
+		...(run.timings ? { timings: run.timings } : {}),
+		...(run.errors.length > 0 ? { errors: run.errors } : {}),
+		artifacts: await persistEvalArtifactReferences(test.artifacts(), runId, artifactDirectory),
+		...(Object.keys(metadata).length > 0 ? { metadata } : {}),
+	};
+	await mkdir(artifactDirectory, { recursive: true, mode: 0o700 });
+	await appendFile(join(artifactDirectory, "runs.jsonl"), `${JSON.stringify(record)}\n`, {
+		encoding: "utf8",
+		flag: "a",
+		mode: 0o600,
+	});
+}
+
+function collectHarnessObservations(modules: ReadonlyArray<TestModule>): HarnessObservation[] {
+	const observations: HarnessObservation[] = [];
+	for (const module of modules) {
+		for (const test of module.children.allTests()) {
+			const harness = test.meta().harness;
+			if (!harness || !isHarnessRun(harness.run)) continue;
+			const run = harness.run;
+			const iteration = parseEvalHarnessIterationArtifact(run.artifacts?.[EVAL_HARNESS_ITERATION_ARTIFACT]);
+			if (!iteration) continue;
+			const score = readFiniteNumber(test.meta().eval?.avgScore);
+			const estimatedCostUsd = readFiniteNumber(run.usage.metadata?.estimatedCostUsd);
+			const observation = {
+				evalSet: iteration.evalSet,
+				groupKey: iteration.groupKey,
+				testName: test.name,
+				file: module.relativeModuleId,
+				harness: iteration.harness,
+				baseline: iteration.baseline,
+				candidates: iteration.candidates,
+				repetition: iteration.repetition,
+				...(run.usage.totalTokens === undefined ? {} : { totalTokens: run.usage.totalTokens }),
+				...(run.timings?.totalMs === undefined ? {} : { totalMs: run.timings.totalMs }),
+				...(estimatedCostUsd === undefined ? {} : { estimatedCostUsd }),
+			};
+			if (run.errors.length > 0) observations.push({ ...observation, outcome: "errored" });
+			else if (score !== undefined) observations.push({ ...observation, outcome: "scored", score });
+			else {
+				const state = test.result().state;
+				const outcome = state === "passed" ? "unscored" : state === "failed" ? "errored" : state;
+				observations.push({ ...observation, outcome });
+			}
+		}
+	}
+	return observations;
+}
+
+export default class EvalHarnessReporter implements Reporter {
+	private vitest: Vitest | undefined;
+
+	onInit(vitest: Vitest): void {
+		this.vitest = vitest;
+	}
+
+	async onTestCaseResult(test: TestCase): Promise<void> {
+		await appendHarnessRunReport(test);
+	}
+
+	onTestRunEnd(
+		modules: ReadonlyArray<TestModule>,
+		_errors: ReadonlyArray<SerializedError>,
+		reason: TestRunEndReason,
+	): void {
+		if (reason === "interrupted") {
+			this.vitest?.logger.log("\nEval comparisons unavailable: test run interrupted.");
+			return;
+		}
+		const report = summarizeHarnessComparisons(collectHarnessObservations(modules));
+		const formatted = formatHarnessComparisonReport(report);
+		if (formatted) this.vitest?.logger.log(`\n${formatted}`);
+	}
+}