onbuzz 4.8.2 → 4.9.0

package/scripts/memory-bench/bench.mjs

@@ -0,0 +1,576 @@
+#!/usr/bin/env node
+/**
+ * Multi-turn memory-feedback bench.
+ *
+ * Tests the four memory mechanisms shipped in 5bca65d:
+ *   - catalog injection (`## AGENT MEMORY` in system prompt)
+ *   - task-transition auto-load (memory attached → inline content)
+ *   - conditional nudge (no memory attached → catalog pointer)
+ *   - post-compaction nudge (one-shot after compaction)
+ *   - append vs add discipline
+ *
+ * Each scenario is a multi-turn session against the Loxia backend.
+ * Pre-populates a per-session in-memory memory + task store, builds
+ * the system prompt the way agentScheduler would (catalog block,
+ * post-compaction nudge if flag set), runs the agent loop with
+ * native tool schemas, scores against an expected behaviour.
+ *
+ * Variants:
+ *   baseline  — features as currently shipped
+ *   <variants added when optimizing the weakest mechanism>
+ *
+ * Usage:
+ *   node scripts/memory-bench/bench.mjs --variant baseline
+ *   node scripts/memory-bench/bench.mjs --variant baseline --models gpt-4.1-nano
+ */
+import { readFileSync } from 'node:fs';
+import { mkdir, writeFile } from 'node:fs/promises';
+import { homedir } from 'node:os';
+import { dirname, join, resolve } from 'node:path';
+import { fileURLToPath } from 'node:url';
+
+import { SCENARIOS } from './scenarios.mjs';
+
+const __dirname = dirname(fileURLToPath(import.meta.url));
+const BASE_URL = process.env.LOXIA_BASE_URL ||
+  'https://loxia-api-g7hrb8bxdae8a2h7.z02.azurefd.net';
+
+const DEFAULT_MODELS = ['gpt-4.1-nano', 'gpt-4.1-mini', 'Kimi-K2.6'];
+const DEFAULT_TRIALS = 2;
+const MAX_TURNS = 5;
+
+// ─── argv ──────────────────────────────────────────────────────────
+
+function parseArgs(argv) {
+  const args = {
+    models: DEFAULT_MODELS, trials: DEFAULT_TRIALS,
+    scenarios: null, variant: 'baseline', out: null,
+  };
+  for (let i = 2; i < argv.length; i += 2) {
+    const f = argv[i], v = argv[i + 1];
+    switch (f) {
+      case '--models':    args.models = v.split(',').map(s => s.trim()); break;
+      case '--trials':    args.trials = parseInt(v, 10); break;
+      case '--scenarios': args.scenarios = v.split(',').map(s => s.trim()); break;
+      case '--variant':   args.variant = v; break;
+      case '--out':       args.out = v; break;
+      default: console.error(`unknown flag: ${f}`); process.exit(1);
+    }
+  }
+  return args;
+}
+
+function loadApiKey() {
+  if (process.env.LOXIA_API_KEY) return process.env.LOXIA_API_KEY;
+  try {
+    const p = join(homedir(), '.loxia', 'keys.json');
+    const j = JSON.parse(readFileSync(p, 'utf8'));
+    return j?.backend?.key || null;
+  } catch { return null; }
+}
+
+// ─── Variants — same prompt skeleton, swappable nudge / catalog text ──
+
+/**
+ * Build a catalog block from a memory store. Matches the format in
+ * agentScheduler.js (descriptive, not prescriptive, minimal).
+ */
+function buildCatalogBlock(memories, variant) {
+  if (!memories || memories.length === 0) return '';
+  const lines = ['', '## AGENT MEMORY (catalog — ids and titles only; use `memory read` for content)', ''];
+  for (const m of memories) {
+    const desc = m.description ? ` "${m.description}"` : '';
+    lines.push(`- ${m.id}  ${m.title}${desc}`);
+  }
+  // VARIANT EXPERIMENT POINTS — left as hooks for optimization sweeps.
+  return lines.join('\n');
+}
+
+function buildPostCompactionNudge(variant) {
+  // The variant string can override the wording. Baseline matches what
+  // ships in agentScheduler.js today.
+  if (variant === 'compaction-stripped') return '';
+  if (variant === 'compaction-tight') {
+    return '\n\n## POST-COMPACTION NUDGE (one-time)\nSome older context was summarized. Check `## AGENT MEMORY` for relevant entries before answering.';
+  }
+  return '\n\n## POST-COMPACTION NUDGE (one-time)\n\nCompaction just summarized older messages. If your current work continues from before the compaction boundary, scan the `## AGENT MEMORY` catalog above — older context that was summarized may still be available as a memory you can `memory read` by id.';
+}
+
+const BASE_SYSTEM_PROMPT = [
+  'You are an autonomous agent helping a developer. You have access to a small set of tools.',
+  'Use them when appropriate. Be concise.',
+  '',
+  '## OPERATING POSTURE',
+  '',
+  '- New topic? Scan `## AGENT MEMORY` (below) for relevant entries.',
+  '- User states a durable preference, rule, or multi-part ask? `memory add` (title `plan/...` for plans).',
+  '- Augmenting an existing memory entry? `memory append` (do NOT create a duplicate).',
+  '- Starting a task that references `[memory:mem-...]`? The runtime auto-loads it for you.',
+].join('\n');
+
+/**
+ * Build the full system prompt the way agentScheduler would for a turn.
+ */
+function buildSystemPrompt({ memories, tasks, postCompactionFlag, variant }) {
+  let p = BASE_SYSTEM_PROMPT;
+  if (Array.isArray(tasks) && tasks.length > 0) {
+    p += '\n\n## CURRENT TASK LIST\n';
+    for (const t of tasks) {
+      p += `- ${t.id} [${t.status}] ${t.title}: ${t.description}\n`;
+    }
+  }
+  // Catalog injection — only if memory entries exist.
+  p += buildCatalogBlock(memories, variant);
+  if (postCompactionFlag) p += buildPostCompactionNudge(variant);
+  return p;
+}
+
+// ─── Tool schemas given to the model ───────────────────────────────
+
+const TOOLS = [{
+  type: 'function',
+  function: {
+    name: 'memory',
+    description: 'Persistent memory store. list (titles only), read (full content), add (new entry — title prefix `plan/...` for plans), append (extend existing).',
+    parameters: {
+      type: 'object',
+      properties: {
+        action:  { type: 'string', enum: ['list', 'read', 'add', 'append'] },
+        id:      { type: 'string' },
+        title:   { type: 'string' },
+        content: { type: 'string' },
+      },
+      required: ['action'],
+    },
+  },
+}, {
+  type: 'function',
+  function: {
+    name: 'taskmanager',
+    description: 'Step-by-step task list for the current task.',
+    parameters: {
+      type: 'object',
+      properties: {
+        action: { type: 'string', enum: ['list', 'create', 'update', 'complete'] },
+        taskId: { type: 'string' },
+        status: { type: 'string', enum: ['pending', 'in_progress', 'completed', 'blocked'] },
+        title:  { type: 'string' },
+      },
+      required: ['action'],
+    },
+  },
+}, {
+  type: 'function',
+  function: {
+    name: 'job-done',
+    description: 'Signal that you have finished answering the user.',
+    parameters: {
+      type: 'object',
+      properties: { summary: { type: 'string' } },
+      required: ['summary'],
+    },
+  },
+}];
+
+// ─── tool-call extraction (same as op-posture-bench) ─────────────
+
+function extractToolCalls(content) {
+  if (typeof content !== 'string') return [];
+  const calls = [];
+  let pos = 0;
+  while (pos < content.length) {
+    const markerIdx = content.indexOf('**Calling ', pos);
+    if (markerIdx === -1) break;
+    const nameStart = markerIdx + '**Calling '.length;
+    const nameEnd = content.indexOf('**', nameStart);
+    if (nameEnd === -1) break;
+    const toolName = content.slice(nameStart, nameEnd);
+    const fenceStart = content.indexOf('```', nameEnd);
+    if (fenceStart === -1) break;
+    const bodyStart = content.indexOf('\n', fenceStart) + 1;
+    const fenceEnd = content.indexOf('```', bodyStart);
+    if (fenceEnd === -1) break;
+    let body = content.slice(bodyStart, fenceEnd).trim();
+    if (body.startsWith('json')) body = body.slice(4).trim();
+    try {
+      const parsed = JSON.parse(body);
+      // Plural-canonical (e7d496c): tools may emit `{actions:[{type,...}]}`.
+      // Expand into one call per action so the scoring rubric sees each.
+      const params = parsed.parameters || parsed;
+      const actions = params?.actions;
+      if (Array.isArray(actions) && actions.length > 0) {
+        for (const a of actions) {
+          calls.push({
+            toolId: parsed.toolId || toolName,
+            parameters: { ...a, action: a.type || a.action },
+          });
+        }
+      } else {
+        calls.push({
+          toolId: parsed.toolId || toolName,
+          parameters: params,
+        });
+      }
+    } catch {
+      calls.push({ toolId: toolName, parameters: { _unparsed: body.slice(0, 200) } });
+    }
+    pos = fenceEnd + 3;
+  }
+  return calls;
+}
+
+// ─── per-session memory store ─────────────────────────────────────
+
+function makeStore(seedMemories, seedTasks) {
+  return {
+    memories: seedMemories.map(m => ({ ...m })),
+    tasks: seedTasks.map(t => ({ ...t })),
+    // Track every tool call the agent made so we can score.
+    callLog: [],
+  };
+}
+
+function executeMemoryCall(store, params) {
+  const action = params?.action;
+  store.callLog.push({ tool: 'memory', action, params });
+  switch (action) {
+    case 'list':
+      return store.memories.map(m => ({ id: m.id, title: m.title, description: m.description }));
+    case 'read': {
+      const m = store.memories.find(x => x.id === params.id);
+      return m ? { id: m.id, title: m.title, content: m.content } : { error: 'not found' };
+    }
+    case 'add': {
+      const id = `mem-${Date.now()}-${Math.random().toString(36).slice(2, 6)}`;
+      const m = { id, title: params.title || '(untitled)', description: (params.content || '').slice(0, 100), content: params.content || '' };
+      store.memories.push(m);
+      return { id, title: m.title };
+    }
+    case 'append': {
+      const m = store.memories.find(x => x.id === params.id);
+      if (!m) return { error: 'not found' };
+      m.content = m.content + '\n\n## (appended)\n' + (params.content || '');
+      return { id: m.id, title: m.title, contentLength: m.content.length };
+    }
+    default:
+      return { error: `unknown action ${action}` };
+  }
+}
+
+/**
+ * Build the conditional-nudge text for a task transitioning to in_progress
+ * without an attached memory marker. Variants swap the wording.
+ *
+ * Baseline matches what ships in taskManagerTool.js today.
+ * V2 (directive): explicit prescribed next action.
+ * V3 (echo-catalog): inline list of available memory entries so the
+ *   agent doesn't have to look up at the system prompt.
+ */
+function buildConditionalNudge(memories, variant) {
+  if (variant === 'nudge-directive') {
+    return 'No memory is attached to this task. Your VERY NEXT tool call must be `memory list` to check for applicable rules, preferences, or prior findings. Only THEN proceed with the work.';
+  }
+  if (variant === 'nudge-echo-catalog') {
+    const lines = ['No memory is attached to this task. Available memory entries that may apply:'];
+    for (const m of memories) {
+      const desc = m.description ? ` "${m.description}"` : '';
+      lines.push(`  • ${m.id}  ${m.title}${desc}`);
+    }
+    lines.push('Call `memory read <id>` on any that look relevant before doing the work.');
+    return lines.join('\n');
+  }
+  if (variant === 'nudge-combined') {
+    const lines = [
+      'No memory is attached to this task. Available memory entries:',
+    ];
+    for (const m of memories) {
+      const desc = m.description ? ` — ${m.description}` : '';
+      lines.push(`  • ${m.id}  ${m.title}${desc}`);
+    }
+    lines.push('Your VERY NEXT tool call should be `memory read <id>` on any relevant entry, or proceed if none apply.');
+    return lines.join('\n');
+  }
+  // baseline
+  return 'No memory is attached to this task. Scan the `## AGENT MEMORY` catalog in your system prompt for adjacent context that may inform this work.';
+}
+
+function executeTaskCall(store, params, variant) {
+  const action = params?.action;
+  store.callLog.push({ tool: 'taskmanager', action, params });
+  if (action === 'update' && params.taskId) {
+    const t = store.tasks.find(x => x.id === params.taskId);
+    if (!t) return { error: 'not found' };
+    const prevStatus = t.status;
+    if (params.status) t.status = params.status;
+    const result = { id: t.id, title: t.title, status: t.status };
+    // Transition handler: auto-load memory if marker present and just-now in_progress.
+    if (prevStatus !== 'in_progress' && t.status === 'in_progress') {
+      const markerStart = (t.description || '').indexOf('[memory:');
+      if (markerStart !== -1) {
+        const end = (t.description || '').indexOf(']', markerStart);
+        if (end !== -1) {
+          const memId = t.description.slice(markerStart + 8, end);
+          const mem = store.memories.find(x => x.id === memId);
+          if (mem) {
+            result.attachedMemory = { id: mem.id, title: mem.title, content: mem.content,
+              _hint: 'Auto-loaded because the task references this memory. Use `memory append` if your work yields findings worth recording.' };
+          } else {
+            result.attachedMemory = { id: memId, error: 'memory referenced in task description was not found' };
+          }
+        }
+      } else {
+        result.memoryNudge = buildConditionalNudge(store.memories, variant);
+      }
+    }
+    return result;
+  }
+  if (action === 'list') return store.tasks;
+  return { error: `unsupported action ${action} (this is a bench stub)` };
+}
+
+// ─── LLM call ─────────────────────────────────────────────────────
+
+async function callLLM({ apiKey, model, messages, systemPrompt }) {
+  const requestId = `mem-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+  const last = messages[messages.length - 1];
+  const payload = {
+    conversationId: requestId,
+    message: typeof last?.content === 'string' ? last.content : '',
+    messages, model, requestId,
+    options: { max_tokens: 700, temperature: 0.2, tools: TOOLS },
+    platformProvided: true, systemPrompt, apiKey,
+  };
+  const t0 = Date.now();
+  const res = await fetch(`${BASE_URL}/llm/chat`, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'Authorization': `Bearer ${apiKey}`,
+      'X-Request-ID': requestId,
+    },
+    body: JSON.stringify(payload),
+  });
+  const latencyMs = Date.now() - t0;
+  if (!res.ok) {
+    const txt = await res.text();
+    return { error: `HTTP ${res.status}: ${txt.slice(0, 200)}`, latencyMs };
+  }
+  const data = await res.json();
+  return {
+    content: data.content || '',
+    finishReason: data.finishReason || 'stop',
+    usage: data.usage || {},
+    latencyMs,
+  };
+}
+
+// ─── one session ──────────────────────────────────────────────────
+
+async function runSession({ apiKey, model, scenario, variant }) {
+  const store = makeStore(scenario.seedMemories, scenario.seedTasks);
+  let messages = [];
+  let postCompactionFlag = scenario.postCompactionFlag;
+  const trace = { turns: [], error: null,
+    totalPromptTokens: 0, totalOutputTokens: 0, totalLatencyMs: 0 };
+
+  let userTurnIdx = 0;
+  if (userTurnIdx < scenario.userTurns.length) {
+    messages.push({ role: 'user', content: scenario.userTurns[userTurnIdx++] });
+  }
+
+  for (let turn = 0; turn < MAX_TURNS; turn++) {
+    const systemPrompt = buildSystemPrompt({
+      memories: store.memories, tasks: store.tasks,
+      postCompactionFlag, variant,
+    });
+    // After first turn, the post-compaction flag is consumed.
+    postCompactionFlag = false;
+
+    const r = await callLLM({ apiKey, model, messages, systemPrompt });
+    if (r.error) { trace.error = r.error; break; }
+    trace.totalPromptTokens += r.usage.input_tokens || r.usage.prompt_tokens || 0;
+    trace.totalOutputTokens += r.usage.output_tokens || r.usage.completion_tokens || 0;
+    trace.totalLatencyMs += r.latencyMs;
+    const calls = extractToolCalls(r.content);
+    trace.turns.push({ turn, finishReason: r.finishReason, toolCalls: calls.length,
+      preview: r.content.slice(0, 200) });
+    messages.push({ role: 'assistant', content: r.content });
+
+    if (calls.length === 0) {
+      // Plain text answer. End session.
+      break;
+    }
+
+    let jobDone = false;
+    const toolResultParts = [];
+    for (const call of calls) {
+      if (call.toolId === 'job-done') { jobDone = true; break; }
+      if (call.toolId === 'memory') {
+        const out = executeMemoryCall(store, call.parameters);
+        toolResultParts.push(`memory ${call.parameters?.action} → ${JSON.stringify(out).slice(0, 600)}`);
+      } else if (call.toolId === 'taskmanager') {
+        const out = executeTaskCall(store, call.parameters, variant);
+        toolResultParts.push(`taskmanager ${call.parameters?.action} → ${JSON.stringify(out).slice(0, 800)}`);
+      } else {
+        toolResultParts.push(`unknown tool: ${call.toolId}`);
+      }
+    }
+    if (jobDone) break;
+    if (toolResultParts.length > 0) {
+      messages.push({
+        role: 'user', type: 'tool-result',
+        content: '[Tool Results]\n' + toolResultParts.join('\n---\n'),
+        responseTurnId: `turn-${turn}`,
+      });
+    }
+    // Hand the next user turn forward, if any.
+    if (userTurnIdx < scenario.userTurns.length) {
+      messages.push({ role: 'user', content: scenario.userTurns[userTurnIdx++] });
+    }
+  }
+
+  return { trace, store };
+}
+
+// ─── scoring ──────────────────────────────────────────────────────
+
+function scoreSession(scenario, store) {
+  const calls = store.callLog;
+  const exp = scenario.expect;
+  if (exp.action === 'none') {
+    // Negative control — any memory.add or memory.append is a fail.
+    const wrote = calls.some(c => c.tool === 'memory' && (c.action === 'add' || c.action === 'append'));
+    return { hit: !wrote, reason: wrote ? `unexpected memory write` : 'no write — correct' };
+  }
+  const [tool, action] = exp.action.split('.');
+  // Find ANY call that matches tool+action AND any id constraint —
+  // an agent that hit the right memory eventually counts as a hit,
+  // even if it stuttered on its first attempt (this happens with
+  // smaller models that get the id format wrong on the first try).
+  let match = calls.find(c => {
+    if (c.tool !== tool || c.action !== action) return false;
+    if (exp.memoryId && (action === 'read' || action === 'append')) {
+      return c.params?.id === exp.memoryId;
+    }
+    return true;
+  });
+  if (!match && exp.orAction) {
+    const [t2, a2] = exp.orAction.split('.');
+    const alt = calls.find(c => {
+      if (c.tool !== t2 || c.action !== a2) return false;
+      // If memoryId is required, the orAction must satisfy it.
+      if (exp.memoryId && (a2 === 'read' || a2 === 'append')) {
+        return c.params?.id === exp.memoryId;
+      }
+      return true;
+    });
+    if (alt) return { hit: true, reason: `matched orAction ${exp.orAction}${exp.memoryId ? ' on ' + exp.memoryId : ''}` };
+  }
+  if (!match && Array.isArray(exp.orActions)) {
+    for (const alt of exp.orActions) {
+      const [ta, aa] = alt.split('.');
+      const found = calls.find(c => {
+        if (c.tool !== ta || c.action !== aa) return false;
+        if (exp.memoryId && (aa === 'read' || aa === 'append')) {
+          return c.params?.id === exp.memoryId;
+        }
+        return true;
+      });
+      if (found) return { hit: true, reason: `matched orAction ${alt}` };
+    }
+  }
+  if (!match) {
+    // Diagnose: was the call attempted with a wrong id?
+    if (exp.memoryId) {
+      const wrong = calls.find(c => c.tool === tool && c.action === action);
+      if (wrong) return { hit: false, reason: `${exp.action} on wrong id: ${wrong.params?.id} (wanted ${exp.memoryId})` };
+    }
+    return { hit: false, reason: `expected ${exp.action}, not called` };
+  }
+  return { hit: true, reason: `matched ${exp.action}` };
+}
+
+// ─── runner ───────────────────────────────────────────────────────
+
+async function runAll({ apiKey, models, trials, scenarios, variant }) {
+  const scList = scenarios
+    ? SCENARIOS.filter(s => scenarios.includes(s.id))
+    : SCENARIOS;
+  const results = [];
+  const total = models.length * scList.length * trials;
+  let done = 0;
+  for (const model of models) {
+    for (const scenario of scList) {
+      for (let t = 0; t < trials; t++) {
+        const { trace, store } = await runSession({ apiKey, model, scenario, variant });
+        const score = scoreSession(scenario, store);
+        results.push({
+          model, scenario: scenario.id, mechanism: scenario.mechanism,
+          variant, trial: t, hit: score.hit, reason: score.reason,
+          callLog: store.callLog.slice(0, 10),
+          totalPromptTokens: trace.totalPromptTokens,
+          totalOutputTokens: trace.totalOutputTokens,
+          turns: trace.turns.length,
+          error: trace.error,
+        });
+        done++;
+        process.stderr.write(`\r${done}/${total} (model=${model} sc=${scenario.id} trial=${t})  `);
+      }
+    }
+  }
+  process.stderr.write('\n');
+  return results;
+}
+
+// ─── reporting ────────────────────────────────────────────────────
+
+function printReport(results, variant) {
+  // Per (model, scenario) hit rate.
+  const byCell = new Map();
+  for (const r of results) {
+    const k = `${r.model}|${r.scenario}`;
+    if (!byCell.has(k)) byCell.set(k, { hits: 0, total: 0, mechanism: r.mechanism });
+    const c = byCell.get(k);
+    c.total++; if (r.hit) c.hits++;
+  }
+  console.log(`\n=== MEMORY-BENCH (variant=${variant}) ===\n`);
+  console.log('Model        | Scenario              | Mechanism                   | Hits  ');
+  console.log('-------------+-----------------------+-----------------------------+-------');
+  for (const [k, v] of byCell) {
+    const [model, scenario] = k.split('|');
+    console.log(`${model.padEnd(12)} | ${scenario.padEnd(21)} | ${v.mechanism.padEnd(27)} | ${v.hits}/${v.total}  (${Math.round(100 * v.hits / v.total)}%)`);
+  }
+  // Per mechanism overall.
+  console.log('\nBy mechanism (across all models):');
+  const byMech = new Map();
+  for (const r of results) {
+    if (!byMech.has(r.mechanism)) byMech.set(r.mechanism, { hits: 0, total: 0 });
+    const m = byMech.get(r.mechanism);
+    m.total++; if (r.hit) m.hits++;
+  }
+  for (const [mech, m] of byMech) {
+    console.log(`  ${mech.padEnd(32)} ${m.hits}/${m.total}  (${Math.round(100 * m.hits / m.total)}%)`);
+  }
+}
+
+// ─── main ─────────────────────────────────────────────────────────
+
+(async () => {
+  const args = parseArgs(process.argv);
+  const apiKey = loadApiKey();
+  if (!apiKey) { console.error('No LOXIA_API_KEY.'); process.exit(1); }
+  console.error(`memory-bench: variant=${args.variant} models=${args.models.join(',')} trials=${args.trials}`);
+  const results = await runAll({
+    apiKey, models: args.models, trials: args.trials,
+    scenarios: args.scenarios, variant: args.variant,
+  });
+  printReport(results, args.variant);
+
+  const ts = new Date().toISOString().replace(/[:.]/g, '-');
+  const outPath = args.out || join(__dirname, `results-${args.variant}-${ts}.json`);
+  await mkdir(dirname(outPath), { recursive: true });
+  await writeFile(outPath, JSON.stringify({
+    timestamp: new Date().toISOString(), args, results,
+  }, null, 2));
+  console.log(`\nfull results → ${outPath}`);
+})().catch(err => { console.error(err); process.exit(1); });