moving

benchmarks/harness/scripts/beam-distill-1m.ts

@@ -0,0 +1,342 @@
+#!/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); });