Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
144 changes: 37 additions & 107 deletions src/orchestrator/batch.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import type { ProviderName } from "../types/provider"
import type { BenchmarkName } from "../types/benchmark"
import type { SamplingConfig } from "../types/checkpoint"
import type { BenchmarkResult } from "../types/unified"
import type { BenchmarkResult, RetrievalAggregates } from "../types/unified"
import { orchestrator, CheckpointManager } from "./index"
import { createBenchmark } from "../benchmarks"
import { logger } from "../utils/logger"
Expand Down Expand Up @@ -432,116 +432,46 @@ export class BatchManager {
const hasRetrieval = reports.some((r) => r.report.retrieval)
if (hasRetrieval) {
const k = reports.find((r) => r.report.retrieval)?.report.retrieval?.k || 10
console.log(`\nRETRIEVAL METRICS (K=${k})`)
console.log(
"┌" +
"─".repeat(17) +
"┬" +
"─".repeat(9) +
"┬" +
"─".repeat(11) +
"┬" +
"─".repeat(10) +
"┬" +
"─".repeat(9) +
"┬" +
"─".repeat(9) +
"┬" +
"─".repeat(9) +
"┐"
)
console.log(
"│ " +
pad("Provider", 15) +
" │ " +
pad("Hit@K", 7) +
" │ " +
pad("Precision", 9) +
" │ " +
pad("Recall", 8) +
" │ " +
pad("F1", 7) +
" │ " +
pad("MRR", 7) +
" │ " +
pad("NDCG", 7) +
" │"
)
console.log(
"├" +
"─".repeat(17) +
"┼" +
"─".repeat(9) +
"┼" +
"─".repeat(11) +
"┼" +
"─".repeat(10) +
"┼" +
"─".repeat(9) +
"┼" +
"─".repeat(9) +
"┼" +
"─".repeat(9) +
"┤"
)
console.log(`
RETRIEVAL METRICS (K=${k})`)

// Recall/F1/NDCG used to sit in this table but were degenerate without ground-truth
// relevance counts (see #67), so the columns are derived from one list rather than
// repeated across four border strings and two row branches.
const columns: Array<{
header: string
width: number
value: (r: RetrievalAggregates) => string
}> = [
{ header: "Provider", width: 15, value: () => "" },
{ header: "Hit@K", width: 7, value: (r) => padPct(r.hitAtK, 7) },
{ header: "Precision", width: 9, value: (r) => padPct(r.precisionAtK, 9) },
{ header: "MRR", width: 7, value: (r) => r.mrr.toFixed(3).padStart(7) },
]

const border = (left: string, mid: string, right: string) =>
left + columns.map((c) => "─".repeat(c.width + 2)).join(mid) + right
const row = (cells: string[]) => "│ " + cells.join(" │ ") + " │"

console.log(border("┌", "┬", "┐"))
console.log(row(columns.map((c) => pad(c.header, c.width))))
console.log(border("├", "┼", "┤"))

for (const { provider, report } of reports) {
if (report.retrieval) {
const r = report.retrieval
console.log(
"│ " +
pad(provider, 15) +
" │ " +
padPct(r.hitAtK, 7) +
" │ " +
padPct(r.precisionAtK, 9) +
" │ " +
padPct(r.recallAtK, 8) +
" │ " +
padPct(r.f1AtK, 7) +
" │ " +
r.mrr.toFixed(3).padStart(7) +
" │ " +
r.ndcg.toFixed(3).padStart(7) +
" │"
const retrieval = report.retrieval
console.log(
row(
columns.map((c, i) =>
i === 0
? pad(provider, c.width)
: retrieval
? c.value(retrieval)
: pad("N/A", c.width)
)
)
} else {
console.log(
"│ " +
pad(provider, 15) +
" │ " +
pad("N/A", 7) +
" │ " +
pad("N/A", 9) +
" │ " +
pad("N/A", 8) +
" │ " +
pad("N/A", 7) +
" │ " +
pad("N/A", 7) +
" │ " +
pad("N/A", 7) +
" │"
)
}
)
}
console.log(
"└" +
"─".repeat(17) +
"┴" +
"─".repeat(9) +
"┴" +
"─".repeat(11) +
"┴" +
"─".repeat(10) +
"┴" +
"─".repeat(9) +
"┴" +
"─".repeat(9) +
"┴" +
"─".repeat(9) +
"┘"
)
console.log(border("└", "┴", "┘"))
}

const allTypes = new Set<string>()
Expand Down
13 changes: 2 additions & 11 deletions src/orchestrator/phases/report.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,23 +22,17 @@ function aggregateRetrievalMetrics(metrics: RetrievalMetrics[]): RetrievalAggreg
(acc, m) => ({
hitAtK: acc.hitAtK + m.hitAtK,
precisionAtK: acc.precisionAtK + m.precisionAtK,
recallAtK: acc.recallAtK + m.recallAtK,
f1AtK: acc.f1AtK + m.f1AtK,
mrr: acc.mrr + m.mrr,
ndcg: acc.ndcg + m.ndcg,
k: m.k,
}),
{ hitAtK: 0, precisionAtK: 0, recallAtK: 0, f1AtK: 0, mrr: 0, ndcg: 0, k: 10 }
{ hitAtK: 0, precisionAtK: 0, mrr: 0, k: 10 }
)

const n = metrics.length
return {
hitAtK: sum.hitAtK / n,
precisionAtK: sum.precisionAtK / n,
recallAtK: sum.recallAtK / n,
f1AtK: sum.f1AtK / n,
mrr: sum.mrr / n,
ndcg: sum.ndcg / n,
k: sum.k,
}
}
Expand Down Expand Up @@ -341,10 +335,7 @@ export function printReport(result: BenchmarkResult): void {
console.log("\nRETRIEVAL QUALITY (K=" + result.retrieval.k + "):")
console.log(` Hit@K: ${(result.retrieval.hitAtK * 100).toFixed(1)}%`)
console.log(` Precision: ${(result.retrieval.precisionAtK * 100).toFixed(1)}%`)
console.log(` Recall: ${(result.retrieval.recallAtK * 100).toFixed(1)}%`)
console.log(` F1: ${(result.retrieval.f1AtK * 100).toFixed(1)}%`)
console.log(` MRR: ${result.retrieval.mrr.toFixed(3)}`)
console.log(` NDCG: ${result.retrieval.ndcg.toFixed(3)}`)
}

console.log("-".repeat(60))
Expand All @@ -361,7 +352,7 @@ export function printReport(result: BenchmarkResult): void {
)
if (stats.retrieval) {
console.log(
` Retrieval: Hit@${stats.retrieval.k}=${(stats.retrieval.hitAtK * 100).toFixed(0)}%, P=${(stats.retrieval.precisionAtK * 100).toFixed(0)}%, R=${(stats.retrieval.recallAtK * 100).toFixed(0)}%, MRR=${stats.retrieval.mrr.toFixed(2)}`
` Retrieval: Hit@${stats.retrieval.k}=${(stats.retrieval.hitAtK * 100).toFixed(0)}%, P=${(stats.retrieval.precisionAtK * 100).toFixed(0)}%, MRR=${stats.retrieval.mrr.toFixed(2)}`
)
}
}
Expand Down
121 changes: 121 additions & 0 deletions src/orchestrator/phases/retrieval-eval.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,121 @@
import { test, expect } from "bun:test"
import type { LanguageModel } from "ai"
import { calculateRetrievalMetrics } from "./retrieval-eval"

// A minimal LanguageModelV2 stub. `ai/test`'s MockLanguageModelV2 would do this for us but it
// pulls in msw, which isn't a dependency of this project — and a hand-rolled stub is 10 lines.
function stubJudge(respond: () => string): LanguageModel {
return {
specificationVersion: "v2",
provider: "stub",
modelId: "stub-judge",
supportedUrls: {},
doGenerate: async () => ({
content: [{ type: "text" as const, text: respond() }],
finishReason: "stop" as const,
usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 },
warnings: [],
}),
doStream: async () => {
throw new Error("not used")
},
} as unknown as LanguageModel
}

const judgeReturning = (text: string) => stubJudge(() => text)
const judgeThatFails = () =>
stubJudge(() => {
throw new Error("429 rate limited")
})

const RESULTS = [{ chunk: "a" }, { chunk: "b" }, { chunk: "c" }, { chunk: "d" }]

test("reports only metrics that are well defined without ground-truth relevance counts", async () => {
// 2nd and 4th results relevant.
const judge = judgeReturning(
JSON.stringify([
{ id: "result_1", relevant: 0 },
{ id: "result_2", relevant: 1 },
{ id: "result_3", relevant: 0 },
{ id: "result_4", relevant: 1 },
])
)

const metrics = await calculateRetrievalMetrics(judge, "q", "gt", RESULTS)

expect(metrics).toEqual({
hitAtK: 1,
precisionAtK: 0.5,
mrr: 0.5, // first relevant at rank 2
k: 4,
relevantRetrieved: 2,
})
// The degenerate trio must not come back: recall was identical to hitAtK by construction,
// F1 was a re-encoding of precision, and NDCG's ideal set was built from what was found.
for (const gone of ["recallAtK", "f1AtK", "ndcg", "totalRelevant"]) {
expect(metrics).not.toHaveProperty(gone)
}
})

test("precision and MRR distinguish rankings that the old NDCG scored identically", async () => {
// Old behaviour: IDCG was built from the retrieved relevant count, so "one relevant at rank 1"
// and "four relevant at ranks 1-4" both scored NDCG 1.0 and recall 100%.
const onlyFirst = await calculateRetrievalMetrics(
judgeReturning(
JSON.stringify([
{ id: "result_1", relevant: 1 },
{ id: "result_2", relevant: 0 },
{ id: "result_3", relevant: 0 },
{ id: "result_4", relevant: 0 },
])
),
"q",
"gt",
RESULTS
)
const allFour = await calculateRetrievalMetrics(
judgeReturning(
JSON.stringify([
{ id: "result_1", relevant: 1 },
{ id: "result_2", relevant: 1 },
{ id: "result_3", relevant: 1 },
{ id: "result_4", relevant: 1 },
])
),
"q",
"gt",
RESULTS
)

// Both "hit" and both rank a relevant result first, so those two agree...
expect(onlyFirst!.hitAtK).toBe(allFour!.hitAtK)
expect(onlyFirst!.mrr).toBe(allFour!.mrr)
// ...but precision now separates them, which is the only honest signal available here.
expect(onlyFirst!.precisionAtK).toBe(0.25)
expect(allFour!.precisionAtK).toBe(1)
})

test("a judge failure yields no metrics instead of a zero score", async () => {
// Coercing a rate-limit or timeout to relevant:0 was indistinguishable from "retrieved
// nothing useful" and quietly dragged the provider's retrieval numbers down.
const metrics = await calculateRetrievalMetrics(judgeThatFails(), "q", "gt", RESULTS)

expect(metrics).toBeUndefined()
})

test("an unparseable judge response yields no metrics", async () => {
const metrics = await calculateRetrievalMetrics(
judgeReturning("I'm afraid I can't help with that."),
"q",
"gt",
RESULTS
)

expect(metrics).toBeUndefined()
})

test("retrieving nothing is a real zero, not a missing measurement", async () => {
const metrics = await calculateRetrievalMetrics(judgeReturning("[]"), "q", "gt", [])

expect(metrics).toEqual({ hitAtK: 0, precisionAtK: 0, mrr: 0, k: 0, relevantRetrieved: 0 })
})
Loading