343 lines
16 KiB
TypeScript
343 lines
16 KiB
TypeScript
#!/usr/bin/env tsx
|
|
/**
|
|
* BEAM 1M — write-time OBSERVATION DISTILLATION (mem0/Mastra/LongMemEval-style).
|
|
*
|
|
* Ported from D:\Projects\hive-mind\benchmarks\longmemeval\34-run-observations.mjs
|
|
* (the pattern that added ~+10pp on LongMemEval). Instead of retrieving raw
|
|
* conversation turns (~900 tok each), we distill each conversation into dense,
|
|
* dated, atomic, pronoun-resolved facts (~35 tok each) with a windowed
|
|
* gpt-4o-mini pass, and store those as `agent_inferred` frames in a SEPARATE
|
|
* per-conversation mind cache (`minds-1M-obs/`). Answering then retrieves top-k
|
|
* FACTS ≈ true mem0-parity semantics (~7K-tok prompts, cheap) with our
|
|
* extraction quality — this is what lets a top-200 run cost ~$20 instead of
|
|
* ~$470 while (hypothesis) lifting accuracy on the harder abilities.
|
|
*
|
|
* DISTILL_SYSTEM is the LongMemEval prompt verbatim. Windows are ~WIN chars of
|
|
* dated turn text; BEAM carries per-MESSAGE `time_anchor` dates ("March-01-2024"),
|
|
* carried forward so each window is tagged with its session date.
|
|
*
|
|
* RESUMABLE: per-conv `.done.json` marker in minds-1M-obs; skip-if-complete; a
|
|
* partial (crashed) mind without a marker is rebuilt.
|
|
*
|
|
* COST-SAFE: `--estimate` builds windows only (NO gpt-4o-mini calls) and reports
|
|
* window/token counts + a projected full-35 cost. Use it before any paid run.
|
|
* Embedding of the resulting facts is local ollama (free).
|
|
*
|
|
* Usage:
|
|
* tsx benchmarks/harness/scripts/beam-distill-1m.ts --estimate # free
|
|
* tsx benchmarks/harness/scripts/beam-distill-1m.ts --convs 1-35 [--win 8000] [--conc 6]
|
|
*/
|
|
|
|
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 { createSubstrate } from '../src/substrate.js';
|
|
import { createBeamOpenAiClient, loadDotEnv } from '../src/beam-openai-client.js';
|
|
import type { BeamOpenAiClient } from '../src/beam-openai-client.js';
|
|
|
|
const DISTILL_SYSTEM =
|
|
'You extract durable, atomic facts about the USER from a slice of their conversation with an assistant. ' +
|
|
'Output one fact per line, each starting with "[YYYY-MM-DD] " using the date the fact/event pertains to ' +
|
|
'(use the session date shown if no other date). Cover: preferences and dislikes, possessions/brands/tools, ' +
|
|
'decisions, events (what happened, when), plans, personal attributes, relationships, numbers/quantities. ' +
|
|
'Be specific and self-contained (resolve pronouns to the entity). Only facts grounded in the text. ' +
|
|
'No preamble, no bullets, no blank lines. If nothing durable, output nothing.';
|
|
|
|
interface Args {
|
|
convs: number[];
|
|
win: number;
|
|
conc: number;
|
|
distillModel: string;
|
|
beamChats: string;
|
|
mindsDir: string;
|
|
estimate: boolean;
|
|
force: boolean;
|
|
maxTokens: number;
|
|
}
|
|
|
|
function parseConvSpec(spec: string): number[] {
|
|
const out = new Set<number>();
|
|
for (const part of spec.split(',')) {
|
|
const m = part.match(/^(\d+)-(\d+)$/);
|
|
if (m) { for (let i = +m[1]; i <= +m[2]; i++) out.add(i); }
|
|
else if (/^\d+$/.test(part.trim())) out.add(+part.trim());
|
|
}
|
|
return [...out].sort((a, b) => a - b);
|
|
}
|
|
|
|
function parseArgs(): Args {
|
|
const argv = process.argv.slice(2);
|
|
const here = url.fileURLToPath(import.meta.url);
|
|
const repoRoot = path.resolve(path.dirname(here), '..', '..', '..');
|
|
const a: Args = {
|
|
convs: parseConvSpec('1-35'),
|
|
win: 8000,
|
|
conc: 6,
|
|
distillModel: 'gpt-4o-mini',
|
|
beamChats: path.resolve(repoRoot, '..', 'BEAM', 'chats'),
|
|
mindsDir: path.join(repoRoot, 'benchmarks', 'data', 'beam', 'minds-1M-obs'),
|
|
estimate: false,
|
|
force: false,
|
|
maxTokens: 700,
|
|
};
|
|
for (let i = 0; i < argv.length; i++) {
|
|
const f = argv[i]; const next = argv[i + 1];
|
|
if (f === '--convs' && next) { a.convs = parseConvSpec(next); i++; }
|
|
else if (f === '--win' && next) { a.win = parseInt(next, 10); i++; }
|
|
else if (f === '--conc' && next) { a.conc = Math.max(1, parseInt(next, 10)); i++; }
|
|
else if (f === '--distill-model' && next) { a.distillModel = next; i++; }
|
|
else if (f === '--beam-chats' && next) { a.beamChats = path.resolve(next); i++; }
|
|
else if (f === '--minds-dir' && next) { a.mindsDir = path.resolve(next); i++; }
|
|
else if (f === '--estimate') { a.estimate = true; }
|
|
else if (f === '--force') { a.force = true; }
|
|
}
|
|
return a;
|
|
}
|
|
|
|
// ── Dated turn flattening + windowing ────────────────────────────────────────
|
|
|
|
interface RawMsg { role?: string; content?: string; time_anchor?: string | null }
|
|
interface RawBatch { turns?: RawMsg[][] }
|
|
|
|
/** Parse BEAM's "March-01-2024" (or ISO) message time_anchor to YYYY-MM-DD. */
|
|
function normBeamDate(s?: string | null): string | null {
|
|
if (!s) return null;
|
|
const raw = String(s).trim();
|
|
let t = Date.parse(raw);
|
|
if (!Number.isFinite(t)) t = Date.parse(raw.replace(/-/g, ' '));
|
|
if (!Number.isFinite(t)) return null;
|
|
return new Date(t).toISOString().slice(0, 10);
|
|
}
|
|
|
|
interface DatedTurn { role: string; content: string; date: string | null }
|
|
|
|
function flattenDatedTurns(chatJsonPath: string): DatedTurn[] {
|
|
const batches = JSON.parse(fs.readFileSync(chatJsonPath, 'utf-8')) as RawBatch[];
|
|
const out: DatedTurn[] = [];
|
|
let lastDate: string | null = null;
|
|
for (const batch of batches) {
|
|
if (!Array.isArray(batch.turns)) continue;
|
|
for (const group of batch.turns) {
|
|
if (!Array.isArray(group)) continue;
|
|
for (const msg of group) {
|
|
const d = normBeamDate(msg.time_anchor);
|
|
if (d) lastDate = d;
|
|
const content = String(msg.content ?? '').trim();
|
|
if (content) out.push({ role: String(msg.role ?? 'unknown').toLowerCase(), content, date: lastDate });
|
|
}
|
|
}
|
|
}
|
|
return out;
|
|
}
|
|
|
|
interface Window { text: string; date: string | null }
|
|
|
|
function buildWindows(turns: DatedTurn[], win: number): Window[] {
|
|
const out: Window[] = [];
|
|
let cur = ''; let curDate: string | null = null;
|
|
for (const t of turns) {
|
|
const line = `${t.date ? `[${t.date}] ` : ''}${t.role}: ${t.content}\n`;
|
|
if (cur.length + line.length > win && cur) { out.push({ text: cur, date: curDate }); cur = ''; }
|
|
if (!cur) curDate = t.date;
|
|
cur += line;
|
|
}
|
|
if (cur) out.push({ text: cur, date: curDate });
|
|
return out;
|
|
}
|
|
|
|
function toIso(d: string | null): string | undefined {
|
|
if (!d) return undefined;
|
|
const t = Date.parse(d);
|
|
return Number.isFinite(t) ? new Date(t).toISOString() : undefined;
|
|
}
|
|
|
|
function approxTokens(s: string): number { return Math.max(1, Math.ceil(s.length / 4)); }
|
|
|
|
// ── Marker helpers ───────────────────────────────────────────────────────────
|
|
|
|
interface DistillMarker {
|
|
conv: number; gop_id: string; facts: number; windows: number;
|
|
distill_ms: number; index_ms: number; cost_usd: number; win: number;
|
|
distill_model: string; built_at: string;
|
|
}
|
|
function markerPath(dir: string, conv: number): string { return path.join(dir, `beam_1M_${conv}.done.json`); }
|
|
function mindPath(dir: string, conv: number): string { return path.join(dir, `beam_1M_${conv}.mind`); }
|
|
function isComplete(dir: string, conv: number): DistillMarker | null {
|
|
const mp = markerPath(dir, conv);
|
|
if (!fs.existsSync(mp)) return null;
|
|
try { const m = JSON.parse(fs.readFileSync(mp, 'utf-8')) as DistillMarker; if (m && m.facts >= 0 && fs.existsSync(mindPath(dir, conv))) return m; } catch { /* */ }
|
|
return null;
|
|
}
|
|
function logProgress(dir: string, line: string): void {
|
|
const stamped = `${new Date().toISOString()} ${line}`;
|
|
process.stdout.write(stamped + '\n');
|
|
try { fs.appendFileSync(path.join(dir, '_distill-progress.log'), stamped + '\n'); } catch { /* */ }
|
|
}
|
|
|
|
// Poll the local ollama server until it answers /api/tags (or ~2min elapses).
|
|
// Called between index-batch retries so we resume only once the server is live.
|
|
async function waitForOllama(dir: string, conv: number): Promise<void> {
|
|
const host = process.env.OLLAMA_HOST || 'http://localhost:11434';
|
|
const url = `${host.replace(/\/$/, '')}/api/tags`;
|
|
for (let i = 0; i < 24; i++) {
|
|
try {
|
|
const ac = new AbortController();
|
|
const t = setTimeout(() => ac.abort(), 5000);
|
|
const res = await fetch(url, { signal: ac.signal });
|
|
clearTimeout(t);
|
|
if (res.ok) return;
|
|
} catch { /* server not up yet */ }
|
|
await new Promise(r => setTimeout(r, 5000));
|
|
}
|
|
logProgress(dir, `[distill][conv ${conv}] ollama still unresponsive after ~2min wait — retrying batch anyway`);
|
|
}
|
|
|
|
// ── Estimate mode (no gpt-4o-mini spend) ─────────────────────────────────────
|
|
|
|
function runEstimate(args: Args): void {
|
|
let totalWindows = 0; let totalInputChars = 0; let convsSeen = 0;
|
|
for (const conv of args.convs) {
|
|
const cj = path.join(args.beamChats, '1M', String(conv), 'chat.json');
|
|
if (!fs.existsSync(cj)) continue;
|
|
const turns = flattenDatedTurns(cj);
|
|
const wins = buildWindows(turns, args.win);
|
|
totalWindows += wins.length;
|
|
for (const w of wins) totalInputChars += w.text.length;
|
|
convsSeen++;
|
|
}
|
|
const SYS_TOK = approxTokens(DISTILL_SYSTEM) + 20;
|
|
const inputToks = Math.ceil(totalInputChars / 4) + totalWindows * SYS_TOK;
|
|
const outToksEst = totalWindows * 350; // ~350 output tokens/window (facts)
|
|
// gpt-4o-mini pricing.
|
|
const IN = 0.15 / 1e6, OUT = 0.6 / 1e6;
|
|
const cost = inputToks * IN + outToksEst * OUT;
|
|
const scale = convsSeen > 0 ? 35 / convsSeen : 1;
|
|
console.log('\n════════ DISTILL COST ESTIMATE (gpt-4o-mini) ════════');
|
|
console.log(`convs measured: ${convsSeen} (window size ${args.win} chars)`);
|
|
console.log(`windows: ${totalWindows} (~${(totalWindows / (convsSeen || 1)).toFixed(0)}/conv)`);
|
|
console.log(`input tokens: ${(inputToks / 1e6).toFixed(2)}M`);
|
|
console.log(`est output tokens: ${(outToksEst / 1e6).toFixed(2)}M (~350/window)`);
|
|
console.log(`cost (measured ${convsSeen} conv): $${cost.toFixed(2)}`);
|
|
console.log(`──────────────────────────────────────────────────`);
|
|
console.log(`PROJECTED FULL 35: $${(cost * scale).toFixed(2)} windows≈${Math.round(totalWindows * scale)}`);
|
|
console.log(`(embedding the facts is local ollama = free; distill is gpt-4o-mini only)`);
|
|
}
|
|
|
|
// ── Distill one conversation ─────────────────────────────────────────────────
|
|
|
|
interface FactRow { text: string; date: string | null }
|
|
|
|
async function distillWindows(client: BeamOpenAiClient, wins: Window[], now: string | null, conc: number, maxTokens: number): Promise<{ facts: FactRow[]; costUsd: number }> {
|
|
const facts: FactRow[] = [];
|
|
let costUsd = 0;
|
|
for (let w = 0; w < wins.length; w += conc) {
|
|
const batch = wins.slice(w, w + conc);
|
|
const results = await Promise.all(batch.map(win =>
|
|
client.chat({ system: DISTILL_SYSTEM, user: `Session date: ${win.date || now || 'unknown'}\n\n${win.text}`, maxTokens })
|
|
.catch(() => ({ text: '', inputTokens: 0, outputTokens: 0, costUsd: 0, latencyMs: 0, failureMode: 'error' })),
|
|
));
|
|
results.forEach((r, bi) => {
|
|
costUsd += r.costUsd;
|
|
for (const line of r.text.split('\n')) {
|
|
const s = line.trim();
|
|
if (s.length > 8) {
|
|
const m = s.match(/\[(\d{4}-\d{2}-\d{2})\]/);
|
|
facts.push({ text: s, date: (m ? m[1] : null) ?? batch[bi].date ?? now });
|
|
}
|
|
}
|
|
});
|
|
}
|
|
return { facts, costUsd };
|
|
}
|
|
|
|
async function distillOne(args: Args, client: BeamOpenAiClient, conv: number): Promise<DistillMarker> {
|
|
const gopId = `beam_${conv}`;
|
|
const cj = path.join(args.beamChats, '1M', String(conv), 'chat.json');
|
|
if (!fs.existsSync(cj)) throw new Error(`chat.json not found: ${cj}`);
|
|
|
|
const mp = mindPath(args.mindsDir, conv);
|
|
for (const s of ['', '-wal', '-shm']) if (fs.existsSync(mp + s)) fs.rmSync(mp + s, { force: true });
|
|
|
|
const turns = flattenDatedTurns(cj);
|
|
const now = (() => { const ds = turns.map(t => t.date).filter(Boolean).sort() as string[]; return ds.length ? ds[ds.length - 1] : null; })();
|
|
const wins = buildWindows(turns, args.win);
|
|
logProgress(args.mindsDir, `[distill][conv ${conv}] START windows=${wins.length} turns=${turns.length}`);
|
|
|
|
const tD = Date.now();
|
|
const { facts, costUsd } = await distillWindows(client, wins, now, args.conc, args.maxTokens);
|
|
const distillMs = Date.now() - tD;
|
|
|
|
const embedder = createOllamaEmbedder();
|
|
const substrate = createSubstrate({ dbPath: mp, embedder });
|
|
let indexMs = 0;
|
|
try {
|
|
substrate.sessions.ensure(gopId, 'beam-obs', `BEAM obs ${gopId}`);
|
|
const toIndex: Array<{ id: number; content: string }> = [];
|
|
const seen = new Set<number>();
|
|
for (const fct of facts) {
|
|
const frame = substrate.frames.createIFrame(gopId, fct.text, 'important', 'agent_inferred', toIso(fct.date));
|
|
if (seen.has(frame.id)) continue;
|
|
seen.add(frame.id);
|
|
toIndex.push({ id: frame.id, content: fct.text });
|
|
}
|
|
const tI = Date.now();
|
|
for (let b = 0; b < toIndex.length; b += 200) {
|
|
const batch = toIndex.slice(b, b + 200);
|
|
// ollama periodically becomes unresponsive under sustained multi-hour load
|
|
// (mid-embed AbortError, or a hard connect-timeout when the server stalls).
|
|
// A resumable run must not die on either: wait for the server to come back,
|
|
// then retry the batch. Up to 8 attempts, backoff to 60s (~4min window).
|
|
for (let attempt = 1; ; attempt++) {
|
|
try { await substrate.search.indexFramesBatch(batch); break; }
|
|
catch (err) {
|
|
if (attempt >= 8) throw err;
|
|
const waitMs = Math.min(60000, 8000 * attempt);
|
|
logProgress(args.mindsDir, `[distill][conv ${conv}] index batch @${b} failed (attempt ${attempt}/8): ${(err as Error).message} — waiting for ollama, retry in ${waitMs / 1000}s`);
|
|
await new Promise(r => setTimeout(r, waitMs));
|
|
await waitForOllama(args.mindsDir, conv);
|
|
}
|
|
}
|
|
}
|
|
indexMs = Date.now() - tI;
|
|
|
|
const marker: DistillMarker = {
|
|
conv, gop_id: gopId, facts: toIndex.length, windows: wins.length,
|
|
distill_ms: distillMs, index_ms: indexMs, cost_usd: Math.round(costUsd * 1e4) / 1e4,
|
|
win: args.win, distill_model: args.distillModel, built_at: new Date().toISOString(),
|
|
};
|
|
fs.writeFileSync(markerPath(args.mindsDir, conv), JSON.stringify(marker, null, 2) + '\n', 'utf-8');
|
|
return marker;
|
|
} finally {
|
|
substrate.close();
|
|
}
|
|
}
|
|
|
|
async function main(): Promise<void> {
|
|
const args = parseArgs();
|
|
if (args.estimate) { runEstimate(args); return; }
|
|
|
|
fs.mkdirSync(args.mindsDir, { recursive: true });
|
|
loadDotEnv();
|
|
const client = createBeamOpenAiClient({ model: args.distillModel });
|
|
logProgress(args.mindsDir, `[distill] start convs=${args.convs[0]}..${args.convs[args.convs.length - 1]} (n=${args.convs.length}) win=${args.win} conc=${args.conc} model=${args.distillModel}`);
|
|
|
|
let done = 0; let totalCost = 0; const total = args.convs.length; const t0 = Date.now();
|
|
for (const conv of args.convs) {
|
|
const existing = args.force ? null : isComplete(args.mindsDir, conv);
|
|
if (existing) { done++; logProgress(args.mindsDir, `[distill][conv ${conv}] SKIP (facts=${existing.facts}) [${done}/${total}]`); continue; }
|
|
const c0 = Date.now();
|
|
const m = await distillOne(args, client, conv);
|
|
totalCost += m.cost_usd; done++;
|
|
const secs = ((Date.now() - c0) / 1000).toFixed(0);
|
|
const rate = done / ((Date.now() - t0) / 60000);
|
|
const eta = rate > 0 ? ((total - done) / rate).toFixed(1) : '?';
|
|
logProgress(args.mindsDir, `[distill][conv ${conv}] DONE facts=${m.facts} windows=${m.windows} cost=$${m.cost_usd} took=${secs}s [${done}/${total}] cum=$${totalCost.toFixed(2)} eta=${eta}m`);
|
|
}
|
|
logProgress(args.mindsDir, `[distill] ALL DONE ${done}/${total} convs, total cost=$${totalCost.toFixed(2)} in ${((Date.now() - t0) / 60000).toFixed(1)}m`);
|
|
}
|
|
|
|
main().catch(err => { console.error('[beam-distill-1m] FATAL:', err); process.exit(1); });
|