Files
waggle-os/benchmarks/harness/scripts/beam-e2-cells.ts
Oleg Maslov b20b138fe4 moving
2026-09-02 10:14:22 +02:00

394 lines
20 KiB
TypeScript
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env tsx
/**
* E2 — BEAM 2×2 store×prompt causal ablation on the contradiction ability.
*
* Isolates whether our +23pp BEAM contradiction win comes from (a) retaining raw
* dated turns in the store, (b) the conflict-aware v2 answer prompt, or their
* interaction. Four cells over the SAME 70 contradiction_resolution questions:
*
* | cell | store | prompt |
* | A | raw dated turns (ours) | conflict-aware v2 | published run (reused)
* | B | raw dated turns (ours) | incumbent-base (v1) | this script
* | C | reconciled current-only | conflict-aware v2 | this script
* | D | reconciled current-only | incumbent-base (v1) | this script
*
* Retrieval is held CONSTANT across all cells: OUR retriever, top-30, regenerated
* locally (deterministic) and cached once. Cells B/C/D read the identical cache.
* The "store" dimension changes ONLY whether the retrieved dated turns are passed
* raw (B) or first collapsed by a read-time reconciliation pass (C/D). The
* "prompt" dimension changes ONLY the rule set (v2 vs v1), dating held constant.
*
* Reconciliation is a read-time SIMULATION of write-time reconciliation: a single
* gpt-5-mini ADD/UPDATE/DELETE pass over the top-30 dated turns → a current-state
* dated fact list where the later statement wins and no contradiction survives.
* It is cached per instance_id so cells C and D answer a byte-identical store.
*
* Modes (one command each; resumable):
* --mode retrieval build beam-e2-retrieval.json (needs ollama; no API spend)
* --mode repro-a re-answer N cell-A questions over the cache, compare scores
* --mode cellB answer incumbent over dated raw turns → beam-e2-cellB.jsonl
* --mode reconcile build beam-e2-reconciled.json (gpt-5-mini)
* --mode cellC answer v2 over reconciled store → beam-e2-cellC.jsonl
* --mode cellD answer incumbent over reconciled store → beam-e2-cellD.jsonl
* --mode confound classify every cell's answers surface-both/pick-latest/abstain
*/
import fs from 'node:fs';
import path from 'node:path';
import url from 'node:url';
import process from 'node:process';
import { createOllamaEmbedder } from '@waggle/core';
import type { SearchResult } from '@waggle/core';
import { createSubstrate } from '../src/substrate.js';
import { createBeamOpenAiClient } from '../src/beam-openai-client.js';
import type { BeamOpenAiClient } from '../src/beam-openai-client.js';
import {
buildAnswerGenerationPrompt, // v1 / incumbent-base rule set
buildAnswerGenerationPromptV2, // conflict-aware v2 rule set
judgeQuestion,
} from '../src/beam-nugget-judge.js';
import type { BeamLlmResult } from '../src/beam-nugget-judge.js';
import { buildConvDateMap, renderMemories } from '../src/beam-date-map.js';
const ABILITY = 'contradiction_resolution';
const TOP_K = 30;
const CONVS = Array.from({ length: 35 }, (_, i) => i + 1);
const here = url.fileURLToPath(import.meta.url);
const repoRoot = path.resolve(path.dirname(here), '..', '..', '..');
const beamChats = path.resolve(repoRoot, '..', 'BEAM', 'chats');
const mindsDir = path.join(repoRoot, 'benchmarks', 'data', 'beam', 'minds-1M');
const outDir = path.join(repoRoot, 'benchmarks', 'results', 'beam');
const RETRIEVAL_CACHE = path.join(outDir, 'beam-e2-retrieval.json');
const RECONCILED_CACHE = path.join(outDir, 'beam-e2-reconciled.json');
const PUBLISHED_A = path.join(outDir, 'beam-1m-FULL700-gpt5-retv2.jsonl');
interface Question {
instanceId: string; conv: number; gopId: string; question: string; rubric: string[];
}
function extractRubric(pq: Record<string, unknown>): string[] {
const raw = pq.rubric;
if (Array.isArray(raw)) return raw.map(String).map(s => s.trim()).filter(Boolean);
if (raw && typeof raw === 'object') {
const n = (raw as Record<string, unknown>).nuggets;
if (Array.isArray(n)) return n.map(String).map(s => s.trim()).filter(Boolean);
}
if (raw) return [String(raw).trim()];
return [];
}
/** Load only the contradiction_resolution questions for a conv, in dataset order. */
function loadContradictionQuestions(conv: number): Question[] {
const pqPath = path.join(beamChats, '1M', String(conv), 'probing_questions', 'probing_questions.json');
if (!fs.existsSync(pqPath)) return [];
const data = JSON.parse(fs.readFileSync(pqPath, 'utf-8')) as Record<string, Record<string, unknown>[]>;
const out: Question[] = [];
const questions = data[ABILITY];
if (!Array.isArray(questions)) return [];
questions.forEach((pq, qi) => {
const q = typeof pq.question === 'string' ? pq.question : '';
if (!q) return;
out.push({
instanceId: `beam_1M_${conv}_${ABILITY}_q${qi}`,
conv, gopId: `beam_${conv}`, question: q, rubric: extractRubric(pq),
});
});
return out;
}
function allQuestions(): Question[] {
return CONVS.flatMap(loadContradictionQuestions);
}
function mindPath(conv: number): string { return path.join(mindsDir, `beam_1M_${conv}.mind`); }
/** Oldest-first (frame id asc) — mirrors the harness's memoriesFromResults. */
function memoriesFromResults(results: readonly SearchResult[]): string[] {
return [...results].sort((a, b) => a.frame.id - b.frame.id).map(r => r.frame.content);
}
// ── Retrieval cache (deterministic, no API) ──────────────────────────────────
interface RetrievalEntry { instanceId: string; conv: number; question: string; rubric: string[]; datedMemories: string[]; }
async function buildRetrievalCache(): Promise<void> {
const embedder = createOllamaEmbedder();
const byConv = new Map<number, Question[]>();
for (const q of allQuestions()) {
if (!byConv.has(q.conv)) byConv.set(q.conv, []);
byConv.get(q.conv)!.push(q);
}
const cache: Record<string, RetrievalEntry> = {};
let total = 0;
for (const conv of CONVS) {
const qs = byConv.get(conv) ?? [];
if (qs.length === 0) continue;
const substrate = createSubstrate({ dbPath: mindPath(conv), embedder });
try {
const dateMap = buildConvDateMap(path.join(beamChats, '1M', String(conv), 'chat.json'));
for (const q of qs) {
const results = await substrate.search.search(q.question, { limit: TOP_K, gopId: q.gopId });
const memories = memoriesFromResults(results);
const datedMemories = renderMemories(memories, dateMap, 'v2');
cache[q.instanceId] = { instanceId: q.instanceId, conv, question: q.question, rubric: q.rubric, datedMemories };
total++;
process.stdout.write(` [conv ${conv}] ${q.instanceId}${datedMemories.length} dated turns\n`);
}
} finally { substrate.close(); }
}
fs.writeFileSync(RETRIEVAL_CACHE, JSON.stringify(cache, null, 0) + '\n', 'utf-8');
console.log(`\n[retrieval] cached ${total} questions → ${RETRIEVAL_CACHE}`);
}
function loadRetrievalCache(): Record<string, RetrievalEntry> {
if (!fs.existsSync(RETRIEVAL_CACHE)) { console.error(`[e2] missing ${RETRIEVAL_CACHE}; run --mode retrieval first`); process.exit(2); }
return JSON.parse(fs.readFileSync(RETRIEVAL_CACHE, 'utf-8')) as Record<string, RetrievalEntry>;
}
// ── Reconciliation cache (gpt-5-mini; simulate write-time ADD/UPDATE/DELETE) ──
const RECONCILE_SYSTEM =
'You maintain a running fact store from a user\'s chat history, mem0-style. You are given dated ' +
'conversation excerpts in chronological order (oldest first). Produce the CURRENT state of the ' +
'user\'s facts as a bulleted list. Apply ADD/UPDATE/DELETE semantics: when a later statement ' +
'changes or contradicts an earlier one, KEEP ONLY the later (current) value and discard the ' +
'earlier one — the superseded value must NOT appear anywhere. Each surviving fact is a single ' +
'line prefixed with the date it was last affirmed: [YYYY-MM-DD] fact. Do NOT include ' +
'contradictions, history, or "previously X now Y" phrasing — only the current resolved state. ' +
'Output ONLY the bulleted list, one fact per line.';
function buildReconcilePrompt(datedMemories: string[]): string {
const body = datedMemories.map((m, i) => `${i + 1}. ${m}`).join('\n');
return `## Dated conversation excerpts (oldest first)\n${body}\n\n## Current-state fact list (later statement wins, no contradiction survives):`;
}
interface ReconciledEntry { instanceId: string; reconciledMemories: string[]; }
function readCache(p: string): Record<string, ReconciledEntry> {
if (!fs.existsSync(p)) return {};
try { return JSON.parse(fs.readFileSync(p, 'utf-8')) as Record<string, ReconciledEntry>; }
catch { console.warn(`[reconcile] ${path.basename(p)} unparseable — treating as empty`); return {}; }
}
/** Reconcile a conv-range shard into its own cache file. Skips ids already
* present in the MAIN cache (from an earlier partial run) or this shard's file,
* so parallel shards never redo each other's or the prior run's work. */
async function buildReconciledCache(convRange: number[] | null, cachePath: string): Promise<void> {
const retrieval = loadRetrievalCache();
const client = createBeamOpenAiClient({ model: 'gpt-5-mini' });
const shard = readCache(cachePath);
const mainDone = new Set(Object.keys(readCache(RECONCILED_CACHE)));
const inRange = (conv: number): boolean => !convRange || convRange.includes(conv);
let cost = 0, done = 0;
const ids = Object.keys(retrieval).sort();
for (const id of ids) {
const entry = retrieval[id];
if (!inRange(entry.conv)) continue;
if (shard[id] || mainDone.has(id)) continue;
const res = await client.chat({ system: RECONCILE_SYSTEM, user: buildReconcilePrompt(entry.datedMemories), maxTokens: 4096 });
cost += res.costUsd;
const lines = res.text.split('\n').map(l => l.replace(/^\s*[-*•]\s*/, '').trim()).filter(Boolean);
shard[id] = { instanceId: id, reconciledMemories: lines };
fs.writeFileSync(cachePath, JSON.stringify(shard, null, 0) + '\n', 'utf-8');
done++;
process.stdout.write(` reconciled ${id}: ${entry.datedMemories.length} turns → ${lines.length} facts $${cost.toFixed(3)}\n`);
}
console.log(`\n[reconcile] shard ${path.basename(cachePath)} +${done} (${Object.keys(shard).length} in shard) cost=$${cost.toFixed(4)}`);
}
/** Merge the main cache + all beam-e2-reconciled.shard-*.json into the main cache. */
function mergeReconciledShards(): void {
const merged = readCache(RECONCILED_CACHE);
for (const f of fs.readdirSync(outDir)) {
if (!/^beam-e2-reconciled\.shard-.*\.json$/.test(f)) continue;
const shard = readCache(path.join(outDir, f));
for (const [id, v] of Object.entries(shard)) merged[id] = v;
}
fs.writeFileSync(RECONCILED_CACHE, JSON.stringify(merged, null, 0) + '\n', 'utf-8');
const retrieval = loadRetrievalCache();
const missing = Object.keys(retrieval).filter(id => !merged[id]);
console.log(`[reconcile-merge] ${Object.keys(merged).length}/${Object.keys(retrieval).length} reconciled → ${RECONCILED_CACHE}`);
if (missing.length) console.warn(`[reconcile-merge] MISSING ${missing.length}: ${missing.slice(0, 10).join(', ')}${missing.length > 10 ? '…' : ''}`);
}
function loadReconciledCache(): Record<string, ReconciledEntry> {
if (!fs.existsSync(RECONCILED_CACHE)) { console.error(`[e2] missing ${RECONCILED_CACHE}; run --mode reconcile first`); process.exit(2); }
return JSON.parse(fs.readFileSync(RECONCILED_CACHE, 'utf-8')) as Record<string, ReconciledEntry>;
}
// ── Answer + judge cells ─────────────────────────────────────────────────────
type Store = 'raw' | 'reconciled';
type Prompt = 'v2' | 'incumbent';
function stripAns(text: string): string {
return text.includes('ANSWER:') ? text.split('ANSWER:').pop()!.trim() : text.trim();
}
function buildPrompt(prompt: Prompt, question: string, memories: string[]): string {
return prompt === 'v2'
? buildAnswerGenerationPromptV2(question, memories)
: buildAnswerGenerationPrompt(question, memories);
}
async function runCell(cellLabel: string, store: Store, prompt: Prompt, outPath: string, budget: number): Promise<void> {
const retrieval = loadRetrievalCache();
const reconciled = store === 'reconciled' ? loadReconciledCache() : null;
const client = createBeamOpenAiClient({ model: 'gpt-5' });
const doneIds = new Set<string>();
if (fs.existsSync(outPath)) {
for (const line of fs.readFileSync(outPath, 'utf-8').split('\n')) {
const t = line.trim(); if (!t) continue;
try { const r = JSON.parse(t) as { instance_id?: string }; if (r.instance_id) doneIds.add(r.instance_id); } catch { /* skip */ }
}
console.log(`[${cellLabel}] resume: ${doneIds.size} already answered`);
}
const out = fs.createWriteStream(outPath, { flags: 'a' });
const ids = Object.keys(retrieval).sort();
let cost = 0, n = 0;
for (const id of ids) {
if (doneIds.has(id)) continue;
if (cost >= budget) { console.warn(`[${cellLabel}] budget $${budget} hit`); break; }
const entry = retrieval[id];
const memories = store === 'raw' ? entry.datedMemories : (reconciled![id]?.reconciledMemories ?? []);
const ans = await client.chat({ system: '', user: buildPrompt(prompt, entry.question, memories), maxTokens: 4096 });
cost += ans.costUsd;
const answer = stripAns(ans.text);
const { judgement, llmResults } = await judgeQuestion(
client, { question: entry.question, rubric: entry.rubric, memoryAbility: ABILITY, answer },
);
for (const r of llmResults) cost += r.costUsd;
out.write(JSON.stringify({
instance_id: id, conv: entry.conv, memory_ability: ABILITY, question: entry.question,
answer, score: judgement.score, judgment: judgement.judgment, nugget_scores: judgement.nuggetScores,
n_nuggets: entry.rubric.length, cell: cellLabel, store, prompt, top_k: TOP_K,
n_memories: memories.length,
note: store === 'reconciled' ? 'read-time simulation of write-time reconciliation' : undefined,
}) + '\n');
n++;
process.stdout.write(` [${cellLabel}] ${id} score=${judgement.score.toFixed(2)} $${cost.toFixed(3)}\n`);
}
out.end();
console.log(`\n[${cellLabel}] answered_now=${n} cost=$${cost.toFixed(4)}${outPath}`);
}
// ── Cell-A reproduction spot check ───────────────────────────────────────────
async function reproA(nSample: number): Promise<void> {
const retrieval = loadRetrievalCache();
const publishedA = new Map<string, number>();
for (const line of fs.readFileSync(PUBLISHED_A, 'utf-8').split('\n')) {
const t = line.trim(); if (!t) continue;
const r = JSON.parse(t) as { instance_id: string; memory_ability: string; score: number };
if (r.memory_ability === ABILITY && !publishedA.has(r.instance_id)) publishedA.set(r.instance_id, r.score);
}
const client = createBeamOpenAiClient({ model: 'gpt-5' });
const ids = Object.keys(retrieval).sort().slice(0, nSample);
let cost = 0;
const rows: Array<{ id: string; repro: number; published: number }> = [];
for (const id of ids) {
const entry = retrieval[id];
const ans = await client.chat({ system: '', user: buildAnswerGenerationPromptV2(entry.question, entry.datedMemories), maxTokens: 4096 });
cost += ans.costUsd;
const answer = stripAns(ans.text);
const { judgement, llmResults } = await judgeQuestion(
client, { question: entry.question, rubric: entry.rubric, memoryAbility: ABILITY, answer },
);
for (const r of llmResults) cost += r.costUsd;
rows.push({ id, repro: judgement.score, published: publishedA.get(id) ?? NaN });
process.stdout.write(` repro ${id}: repro=${judgement.score.toFixed(2)} published=${(publishedA.get(id) ?? NaN).toFixed(2)} $${cost.toFixed(3)}\n`);
}
const meanRepro = rows.reduce((s, r) => s + r.repro, 0) / rows.length;
const meanPub = rows.reduce((s, r) => s + r.published, 0) / rows.length;
console.log(`\n[repro-a] n=${rows.length} mean repro=${meanRepro.toFixed(4)} vs published=${meanPub.toFixed(4)} cost=$${cost.toFixed(4)}`);
}
// ── Confound classification: surface-both vs pick-latest vs abstain ──────────
const CONFOUND_SYSTEM =
'You classify how an assistant answer handled a question about a fact the user stated ' +
'inconsistently over time. Respond with STRICT JSON: {"label":"<one of>"} where <one of> is:\n' +
'- surface_both: the answer explicitly flags that the stored information is contradictory / ' +
'conflicting, OR presents more than one of the conflicting values (optionally asking the user which is correct).\n' +
'- pick_latest: the answer commits to a single value/state as the current answer without flagging any contradiction.\n' +
'- abstain: the answer declines, saying it does not have enough information.\n' +
'Judge ONLY the answer text\'s behaviour, not correctness.';
async function classifyConfound(cellFile: string): Promise<Record<string, number> | null> {
const p = path.join(outDir, cellFile);
if (!fs.existsSync(p)) { console.log(`[confound] ${cellFile}: missing`); return null; }
const client = createBeamOpenAiClient({ model: 'gpt-5-mini' });
const seen = new Set<string>();
const counts: Record<string, number> = { surface_both: 0, pick_latest: 0, abstain: 0, other: 0 };
let cost = 0;
for (const line of fs.readFileSync(p, 'utf-8').split('\n')) {
const t = line.trim(); if (!t) continue;
const r = JSON.parse(t) as { instance_id: string; question: string; answer: string };
if (seen.has(r.instance_id)) continue; seen.add(r.instance_id);
const res = await client.chat({
system: CONFOUND_SYSTEM,
user: `QUESTION:\n${r.question}\n\nANSWER:\n${r.answer}`,
jsonMode: true, maxTokens: 200,
});
cost += res.costUsd;
let label = 'other';
try {
const o = JSON.parse(res.text) as { label?: string };
if (o.label && ['surface_both', 'pick_latest', 'abstain'].includes(o.label)) label = o.label;
} catch { /* keep other */ }
counts[label] = (counts[label] ?? 0) + 1;
}
const n = seen.size;
console.log(`[confound] ${cellFile}: n=${n} surface_both=${counts.surface_both} pick_latest=${counts.pick_latest} abstain=${counts.abstain} other=${counts.other} cost=$${cost.toFixed(3)}`);
return counts;
}
async function runConfound(): Promise<void> {
const results: Record<string, Record<string, number> | null> = {};
for (const [cell, file] of [['A','beam-e2-cellA.jsonl'],['B','beam-e2-cellB.jsonl'],['C','beam-e2-cellC.jsonl'],['D','beam-e2-cellD.jsonl']] as const) {
results[cell] = await classifyConfound(file);
}
fs.writeFileSync(path.join(outDir, 'beam-e2-confound.json'), JSON.stringify(results, null, 2) + '\n', 'utf-8');
console.log(`\n[confound] → ${path.join(outDir, 'beam-e2-confound.json')}`);
}
async function main(): Promise<void> {
const argv = process.argv.slice(2);
const mode = argv[argv.indexOf('--mode') + 1] ?? '';
const budgetArg = argv.indexOf('--budget');
const budget = budgetArg >= 0 ? parseFloat(argv[budgetArg + 1]) : 30;
const nArg = argv.indexOf('--n');
const nSample = nArg >= 0 ? parseInt(argv[nArg + 1], 10) : 5;
fs.mkdirSync(outDir, { recursive: true });
switch (mode) {
case 'retrieval': await buildRetrievalCache(); break;
case 'repro-a': await reproA(nSample); break;
case 'cellA': await runCell('A', 'raw', 'v2', path.join(outDir, 'beam-e2-cellA.jsonl'), budget); break;
case 'cellB': await runCell('B', 'raw', 'incumbent', path.join(outDir, 'beam-e2-cellB.jsonl'), budget); break;
case 'reconcile': {
const convArg = argv.indexOf('--convs');
const convRange = convArg >= 0
? argv[convArg + 1].split(',').flatMap(part => {
const m = part.match(/^(\d+)-(\d+)$/);
if (m) { const r: number[] = []; for (let i = +m[1]; i <= +m[2]; i++) r.push(i); return r; }
return [parseInt(part, 10)];
})
: null;
const cacheArg = argv.indexOf('--cache');
const cachePath = cacheArg >= 0 ? path.join(outDir, argv[cacheArg + 1]) : RECONCILED_CACHE;
await buildReconciledCache(convRange, cachePath);
break;
}
case 'reconcile-merge': mergeReconciledShards(); break;
case 'cellC': await runCell('C', 'reconciled', 'v2', path.join(outDir, 'beam-e2-cellC.jsonl'), budget); break;
case 'cellD': await runCell('D', 'reconciled', 'incumbent', path.join(outDir, 'beam-e2-cellD.jsonl'), budget); break;
case 'confound': await runConfound(); break;
default: console.error(`unknown --mode "${mode}" (retrieval|repro-a|cellB|reconcile|cellC|cellD)`); process.exit(2);
}
}
main().catch(e => { console.error(e); process.exit(1); });