127 lines
4.3 KiB
TypeScript
127 lines
4.3 KiB
TypeScript
import { existsSync, mkdirSync, writeFileSync, appendFileSync } from "node:fs";
|
|
import { dirname, resolve } from "node:path";
|
|
import { parseArgs } from "node:util";
|
|
import { agentmemoryAdapter } from "./adapters/agentmemory.js";
|
|
import { grepAdapter } from "./adapters/grep.js";
|
|
import { vectorAdapter } from "./adapters/vector.js";
|
|
import { loadLongMemEval, stratifySample } from "./load.js";
|
|
import { aggregate, scoreQuestion } from "./score.js";
|
|
import type { Adapter, ScoreRow } from "./types.js";
|
|
|
|
const ADAPTERS: Record<string, Adapter> = {
|
|
grep: grepAdapter as unknown as Adapter,
|
|
vector: vectorAdapter as unknown as Adapter,
|
|
agentmemory: agentmemoryAdapter as unknown as Adapter,
|
|
};
|
|
|
|
interface CliOptions {
|
|
data: string;
|
|
adapters: string;
|
|
k: string;
|
|
limit?: string;
|
|
stratify?: string;
|
|
out: string;
|
|
}
|
|
|
|
function parse(): CliOptions {
|
|
const { values } = parseArgs({
|
|
options: {
|
|
data: { type: "string", default: process.env.LONGMEMEVAL_PATH ?? "" },
|
|
adapters: { type: "string", default: "grep,vector,agentmemory" },
|
|
k: { type: "string", default: "5" },
|
|
limit: { type: "string" },
|
|
stratify: { type: "string" },
|
|
out: { type: "string", default: "eval/reports/longmemeval" },
|
|
},
|
|
});
|
|
return values as unknown as CliOptions;
|
|
}
|
|
|
|
async function main(): Promise<void> {
|
|
const opts = parse();
|
|
if (!opts.data) {
|
|
console.error("--data <path/to/longmemeval_s.json> required (or LONGMEMEVAL_PATH env)");
|
|
process.exit(2);
|
|
}
|
|
const k = Number(opts.k);
|
|
if (!Number.isInteger(k) || k <= 0) {
|
|
console.error(`--k must be a positive integer, got: ${opts.k}`);
|
|
process.exit(2);
|
|
}
|
|
let limit: number | undefined;
|
|
if (opts.limit !== undefined) {
|
|
limit = Number(opts.limit);
|
|
if (!Number.isInteger(limit) || limit <= 0) {
|
|
console.error(`--limit must be a positive integer, got: ${opts.limit}`);
|
|
process.exit(2);
|
|
}
|
|
}
|
|
let perType: number | undefined;
|
|
if (opts.stratify !== undefined) {
|
|
perType = Number(opts.stratify);
|
|
if (!Number.isInteger(perType) || perType <= 0) {
|
|
console.error(`--stratify must be a positive integer, got: ${opts.stratify}`);
|
|
process.exit(2);
|
|
}
|
|
}
|
|
const adapterNames = opts.adapters.split(",").map((s) => s.trim()).filter(Boolean);
|
|
for (const a of adapterNames) {
|
|
if (!ADAPTERS[a]) {
|
|
console.error(`unknown adapter: ${a}. options: ${Object.keys(ADAPTERS).join(",")}`);
|
|
process.exit(2);
|
|
}
|
|
}
|
|
let questions = loadLongMemEval(resolve(opts.data), limit);
|
|
if (perType) questions = stratifySample(questions, perType);
|
|
console.log(
|
|
`loaded ${questions.length} questions, adapters: ${adapterNames.join(",")}, k=${k}`,
|
|
);
|
|
|
|
const outDir = resolve(opts.out);
|
|
mkdirSync(outDir, { recursive: true });
|
|
const ndjsonPath = `${outDir}/scores.ndjson`;
|
|
if (existsSync(ndjsonPath)) writeFileSync(ndjsonPath, "");
|
|
mkdirSync(dirname(ndjsonPath), { recursive: true });
|
|
|
|
const rows: ScoreRow[] = [];
|
|
for (const adapterName of adapterNames) {
|
|
const adapter = ADAPTERS[adapterName];
|
|
console.log(`\n== ${adapter.name} ==`);
|
|
for (const q of questions) {
|
|
const t0 = performance.now();
|
|
const state = await adapter.init(q.haystack);
|
|
try {
|
|
const ranked = await adapter.query(q.question, state, k);
|
|
const latencyMs = performance.now() - t0;
|
|
const row = scoreQuestion(q, ranked, k, adapter.name, latencyMs);
|
|
rows.push(row);
|
|
appendFileSync(ndjsonPath, JSON.stringify(row) + "\n");
|
|
const mark = row.hit ? "+" : "-";
|
|
console.log(
|
|
` ${mark} ${q.id} [${q.type}] R@${k}=${row.recallAtK.toFixed(2)} (${Math.round(latencyMs)}ms)`,
|
|
);
|
|
} finally {
|
|
if (adapter.teardown) await adapter.teardown(state);
|
|
}
|
|
}
|
|
}
|
|
|
|
const agg = aggregate(rows);
|
|
const summaryPath = `${outDir}/summary.json`;
|
|
writeFileSync(summaryPath, JSON.stringify(agg, null, 2));
|
|
|
|
console.log("\n=== Summary ===");
|
|
for (const [adapter, stats] of Object.entries(agg.byAdapter)) {
|
|
console.log(
|
|
` ${adapter.padEnd(22)} P@${k}=${stats.p.toFixed(3)} R@${k}=${stats.r.toFixed(3)} hit=${stats.hit}/${stats.n} p50=${Math.round(stats.latencyP50)}ms`,
|
|
);
|
|
}
|
|
console.log(`\nwrote ${ndjsonPath}`);
|
|
console.log(`wrote ${summaryPath}`);
|
|
}
|
|
|
|
main().catch((err) => {
|
|
console.error(err);
|
|
process.exit(1);
|
|
});
|