memind/memory-v2-memfuse-bench.mjs

import os from 'node:os';
import path from 'node:path';
import { readFile as fsReadFile } from 'node:fs/promises';

import { createPgvectorMemoryBackend, pgvectorMemoryBackendInternals } from './memory-v2-pgvector.mjs';

/**
 * MemFuseBench retrieval harness for Memory V2.
 *
 * MemFuseBench (Mi-Memory / Darwin Agent Team, MIT) ships 357 evidence-grounded
 * questions over 7,823 source-tagged events. This module turns those scenarios
 * into a corpus + case list and runs them through the *production* pgvector
 * backend so the measured numbers reflect `rankHybridCandidates` and
 * `extractKeywordTerms`, not a parallel scoring implementation.
 *
 * The dataset itself is never vendored into this repository. It is read from an
 * external checkout and every entry point degrades to
 * `{ available: false, reason }` when the file is absent, so CI stays green
 * without it.
 */

export const MEMFUSE_DATASET_ENV = 'MEMFUSE_BENCH_DATASET';

export const MEMFUSE_DIMENSIONS = Object.freeze([
  'cross_device_causal_reasoning',
  'cross_device_information_fusion',
  'multi_source_conflict_arbitration',
  'cross_user_information_synthesis',
  'cross_user_query',
  'perspective_difference',
]);

// Events tagged with these sources are adversarial by construction: they look
// topical but are never part of any answer's evidence set.
export const MEMFUSE_DISTRACTOR_SOURCES = Object.freeze(['noise', 'adversarial']);

const DEFAULT_RELATIVE_DATASET_PATH = path.join(
  'mi-memory',
  'MemFuse',
  'MemFuseBench',
  'memfusebench_dataset.json',
);

export function resolveMemFuseDatasetPath(env = process.env) {
  const explicit = String(env?.[MEMFUSE_DATASET_ENV] ?? '').trim();
  if (explicit) return path.resolve(explicit);
  const projectRoot = String(env?.MEMFUSE_PROJECT_ROOT ?? '').trim();
  const base = projectRoot || path.join(os.homedir(), 'Project');
  return path.join(base, DEFAULT_RELATIVE_DATASET_PATH);
}

function asArray(value) {
  return Array.isArray(value) ? value : [];
}

function nonEmptyString(value) {
  const text = String(value ?? '').trim();
  return text.length > 0 ? text : null;
}

export function normalizeMemFuseDataset(raw) {
  const scenarios = asArray(raw?.scenarios)
    .map((scenario) => {
      const scenarioId = nonEmptyString(scenario?.scenario_id);
      if (!scenarioId) return null;
      const episodes = asArray(scenario?.episodes)
        .map((episode) => {
          const episodeId = nonEmptyString(episode?.episode_id);
          const events = asArray(episode?.events)
            .map((event) => {
              const eventId = nonEmptyString(event?.event_id);
              const description = nonEmptyString(event?.description);
              if (!eventId || !description) return null;
              return {
                eventId,
                description,
                device: nonEmptyString(event?.device),
                modality: nonEmptyString(event?.modality) ?? 'event',
                location: nonEmptyString(event?.location),
                source: nonEmptyString(event?.source) ?? 'unknown',
                timestamp: nonEmptyString(event?.timestamp),
                characters: asArray(event?.characters).map(String),
              };
            })
            .filter(Boolean);
          if (!episodeId || events.length === 0) return null;
          return { episodeId, events };
        })
        .filter(Boolean);
      const questions = asArray(scenario?.questions)
        .map((question) => {
          const questionId = nonEmptyString(question?.question_id);
          const text = nonEmptyString(question?.question);
          const evidenceEventIds = [
            ...new Set(asArray(question?.evidence_event_ids).map(String).filter(Boolean)),
          ];
          if (!questionId || !text || evidenceEventIds.length === 0) return null;
          const checklist = asArray(question?.answer_checklist)
            .map((entry) => ({
              point: nonEmptyString(entry?.point) ?? '',
              sourceEvents: [
                ...new Set(asArray(entry?.source_events).map(String).filter(Boolean)),
              ],
            }))
            .filter((entry) => entry.sourceEvents.length > 0);
          return {
            questionId,
            scenarioId,
            question: text,
            evidenceEventIds,
            checklist,
            dimension: nonEmptyString(question?.dimension) ?? 'unknown',
            questionUser: nonEmptyString(question?.question_user),
            questionTime: nonEmptyString(question?.question_time),
            questionDevice: nonEmptyString(question?.question_device),
          };
        })
        .filter(Boolean);
      if (episodes.length === 0 || questions.length === 0) return null;
      return {
        scenarioId,
        description: nonEmptyString(scenario?.description),
        timeSpan: nonEmptyString(scenario?.time_span),
        episodes,
        questions,
      };
    })
    .filter(Boolean);

  if (scenarios.length === 0) {
    return { ok: false, reason: 'dataset_has_no_usable_scenarios' };
  }

  const eventCount = scenarios.reduce(
    (total, scenario) =>
      total + scenario.episodes.reduce((sum, episode) => sum + episode.events.length, 0),
    0,
  );
  const questionCount = scenarios.reduce(
    (total, scenario) => total + scenario.questions.length,
    0,
  );

  return {
    ok: true,
    dataset: {
      scenarios,
      stats: { scenarioCount: scenarios.length, eventCount, questionCount },
    },
  };
}

export async function loadMemFuseDataset({
  datasetPath = null,
  env = process.env,
  readFile = fsReadFile,
} = {}) {
  const resolvedPath = datasetPath ? path.resolve(datasetPath) : resolveMemFuseDatasetPath(env);
  let payload;
  try {
    payload = await readFile(resolvedPath, 'utf8');
  } catch (err) {
    return {
      available: false,
      path: resolvedPath,
      reason: err?.code === 'ENOENT' ? 'dataset_not_found' : 'dataset_read_failed',
      error: err instanceof Error ? err.message : String(err),
    };
  }
  let raw;
  try {
    raw = JSON.parse(payload);
  } catch (err) {
    return {
      available: false,
      path: resolvedPath,
      reason: 'dataset_parse_failed',
      error: err instanceof Error ? err.message : String(err),
    };
  }
  const normalized = normalizeMemFuseDataset(raw);
  if (!normalized.ok) {
    return { available: false, path: resolvedPath, reason: normalized.reason };
  }
  return { available: true, path: resolvedPath, ...normalized.dataset };
}

/**
 * Flatten a scenario into rows shaped like the `memory_embeddings` table so the
 * pgvector backend can consume them unchanged.
 *
 * `includeSourceTags` prepends `[device · location]`. It defaults to false so the
 * baseline measures retrieval over raw event text rather than over a prefix
 * format chosen here.
 */
export function buildScenarioCorpus(scenario, { includeSourceTags = false } = {}) {
  const rows = [];
  const byId = new Map();
  for (const episode of asArray(scenario?.episodes)) {
    for (const event of asArray(episode?.events)) {
      const tags = [event.device, event.location].filter(Boolean).join(' · ');
      const content =
        includeSourceTags && tags ? `[${tags}] ${event.description}` : event.description;
      const timestampMs = event.timestamp ? Date.parse(event.timestamp) : NaN;
      const isoTimestamp = Number.isFinite(timestampMs)
        ? new Date(timestampMs).toISOString()
        : null;
      const row = {
        id: event.eventId,
        content,
        type: event.modality,
        created_at: isoTimestamp,
        updated_at: isoTimestamp,
        episodeId: episode.episodeId,
        source: event.source,
        device: event.device,
        location: event.location,
        timestampMs: Number.isFinite(timestampMs) ? timestampMs : 0,
      };
      rows.push(row);
      byId.set(row.id, row);
    }
  }
  return { rows, byId };
}

export function buildScenarioCases(
  scenario,
  { dimensions = null, maxQuestions = null, questionIds = null } = {},
) {
  const dimensionFilter = dimensions?.length ? new Set(dimensions) : null;
  const idFilter = questionIds?.length ? new Set(questionIds) : null;
  const cases = asArray(scenario?.questions).filter((question) => {
    if (dimensionFilter && !dimensionFilter.has(question.dimension)) return false;
    if (idFilter && !idFilter.has(question.questionId)) return false;
    return true;
  });
  const capped = Number(maxQuestions);
  return Number.isFinite(capped) && capped > 0 ? cases.slice(0, capped) : cases;
}

const CJK_RANGE = /[\u4e00-\u9fff]/u;

function tokenizeForEmbedding(text) {
  const normalized = String(text ?? '')
    .normalize('NFKC')
    .toLowerCase();
  const tokens = [];
  for (const match of normalized.matchAll(/[a-z0-9]{2,}/g)) tokens.push(match[0]);
  if (CJK_RANGE.test(normalized)) {
    const cjk = normalized.replace(/[^\u4e00-\u9fff]/gu, '');
    for (let index = 0; index + 2 <= cjk.length; index += 1) {
      tokens.push(cjk.slice(index, index + 2));
    }
  }
  return tokens;
}

function hashToken(token, dimensions) {
  let hash = 2166136261;
  for (let index = 0; index < token.length; index += 1) {
    hash ^= token.charCodeAt(index);
    hash = Math.imul(hash, 16777619);
  }
  return Math.abs(hash) % dimensions;
}

/**
 * Deterministic offline stand-in for a sentence embedder: hashed bag-of-tokens,
 * L2 normalized. Cosine over these vectors approximates lexical overlap, NOT
 * semantic similarity, so reports built on it are labelled `lexical-hash`. Pass
 * a real `embedText` to measure semantic recall.
 */
export function createLexicalHashEmbedder({ dimensions = 256 } = {}) {
  const size = Math.max(16, Math.min(4096, Number(dimensions) || 256));
  return function embedText(text) {
    const vector = new Array(size).fill(0);
    const tokens = tokenizeForEmbedding(text);
    if (tokens.length === 0) return vector;
    for (const token of tokens) {
      vector[hashToken(token, size)] += 1;
    }
    let norm = 0;
    for (const value of vector) norm += value * value;
    norm = Math.sqrt(norm);
    if (norm === 0) return vector;
    for (let index = 0; index < size; index += 1) vector[index] /= norm;
    return vector;
  };
}

function cosineSimilarity(left, right) {
  const length = Math.min(left.length, right.length);
  let dot = 0;
  let leftNorm = 0;
  let rightNorm = 0;
  for (let index = 0; index < length; index += 1) {
    dot += left[index] * right[index];
    leftNorm += left[index] * left[index];
    rightNorm += right[index] * right[index];
  }
  if (leftNorm === 0 || rightNorm === 0) return 0;
  return dot / (Math.sqrt(leftNorm) * Math.sqrt(rightNorm));
}

function parseVectorLiteral(literal) {
  const text = String(literal ?? '').trim();
  if (!text.startsWith('[') || !text.endsWith(']')) return null;
  const inner = text.slice(1, -1);
  if (!inner) return [];
  const parts = inner.split(',').map((part) => Number(part));
  return parts.some((part) => !Number.isFinite(part)) ? null : parts;
}

function likeParamToTerm(param) {
  return String(param ?? '')
    .replace(/^%+/, '')
    .replace(/%+$/, '')
    .toLowerCase();
}

/**
 * Fake pg pool that emulates the two query shapes issued by
 * `createPgvectorMemoryBackend.resolve`: the vector/recent candidate union and
 * the `content ILIKE` keyword fallback. Emulating the SQL rather than bypassing
 * it keeps the production ranking path under test.
 */
export function createCorpusPool({ rows, embedText, embeddingCache = new Map(), benchQuery = '' } = {}) {
  const corpus = asArray(rows);
  const byRecency = [...corpus].sort((left, right) => right.timestampMs - left.timestampMs);
  const activeQuery = String(benchQuery ?? '');

  function embeddingFor(row) {
    if (!embeddingCache.has(row.id)) {
      const embedded = embedText(row.content);
      embeddingCache.set(row.id, embedded);
    }
    return embeddingCache.get(row.id);
  }

  async function resolvedEmbeddingFor(row) {
    const cached = embeddingFor(row);
    return cached instanceof Promise ? cached : cached;
  }

  function toResultRow(row, score) {
    return {
      id: row.id,
      content: row.content,
      type: row.type,
      created_at: row.created_at,
      updated_at: row.updated_at,
      score,
    };
  }

  let vectorQueryCount = 0;
  let keywordQueryCount = 0;
  // Every id handed back to the backend, i.e. the candidate pool that ranking
  // then narrows to top-k. Tracking it separates "candidate generation lost the
  // evidence" from "ranking lost the evidence".
  const returnedIds = new Set();

  function trackRows(rows) {
    for (const row of rows) returnedIds.add(row.id);
    return rows;
  }

  return {
    stats: () => ({
      vectorQueryCount,
      keywordQueryCount,
      corpusSize: corpus.length,
      candidateIds: new Set(returnedIds),
    }),
    async query(sql, params = []) {
      const text = String(sql ?? '');
      if (text.includes('vector_candidates')) {
        vectorQueryCount += 1;
        const queryVector = parseVectorLiteral(params[1]);
        const limit = Math.max(1, Number(params[2]) || 50);
        const scored = [];
        for (const row of corpus) {
          const embedding = await resolvedEmbeddingFor(row);
          scored.push({
            row,
            score: queryVector ? cosineSimilarity(queryVector, embedding) : 0,
          });
        }
        const selected = pgvectorMemoryBackendInternals.selectVectorCandidateRows(scored, {
          baseLimit: limit,
          recentRows: byRecency,
        });
        return {
          rows: trackRows(
            selected.map(({ row, score }) => toResultRow(row, score)),
          ),
        };
      }
      if (text.includes('ILIKE')) {
        keywordQueryCount += 1;
        const limit = Math.max(1, Number(params[params.length - 1]) || 20);
        const terms = params
          .slice(1, params.length - 1)
          .map(likeParamToTerm)
          .filter(Boolean);
        if (terms.length === 0) return { rows: [] };
        const queryText = activeQuery || terms.join(' ');
        const selected = pgvectorMemoryBackendInternals.selectKeywordCandidateRows(
          corpus,
          queryText,
          terms,
          { returnLimit: limit },
        );
        return {
          rows: trackRows(selected.map((row) => toResultRow(row, null))),
        };
      }
      return { rows: [] };
    },
  };
}

function checklistCoverage(checklist, retrievedIds) {
  if (!checklist?.length) return null;
  let covered = 0;
  for (const entry of checklist) {
    if (entry.sourceEvents.some((eventId) => retrievedIds.has(eventId))) covered += 1;
  }
  return covered / checklist.length;
}

export async function runMemFuseBenchCase({
  testCase,
  corpus,
  embedText,
  limit = 20,
  candidateLimit = 100,
  embeddingCache = new Map(),
}) {
  const pool = createCorpusPool({
    rows: corpus.rows,
    embedText,
    embeddingCache,
    benchQuery: testCase.question,
  });
  const backend = createPgvectorMemoryBackend({ enabled: true, embedQuery: embedText, pool });
  const result = await backend.resolve({
    userId: 'memfuse-bench-user',
    query: testCase.question,
    limit,
    candidateLimit,
  });

  const retrieved = result.memories.map((memory) => memory.id).filter(Boolean);
  const retrievedIds = new Set(retrieved);
  const gold = new Set(testCase.evidenceEventIds);
  const hits = retrieved.filter((id) => gold.has(id));
  const firstHitIndex = retrieved.findIndex((id) => gold.has(id));
  const distractors = retrieved.filter((id) =>
    MEMFUSE_DISTRACTOR_SOURCES.includes(corpus.byId.get(id)?.source),
  );

  const { candidateIds } = pool.stats();
  const candidateHits = [...gold].filter((id) => candidateIds.has(id));
  const candidateRecall = gold.size > 0 ? candidateHits.length / gold.size : 0;
  const recall = gold.size > 0 ? hits.length / gold.size : 0;

  return {
    questionId: testCase.questionId,
    scenarioId: testCase.scenarioId,
    dimension: testCase.dimension,
    goldCount: gold.size,
    returned: retrieved.length,
    hitCount: hits.length,
    recall,
    precision: retrieved.length > 0 ? hits.length / retrieved.length : 0,
    hitAny: hits.length > 0,
    candidateCount: candidateIds.size,
    candidateHitCount: candidateHits.length,
    candidateRecall,
    // Positive value = evidence reached the candidate pool but ranking dropped
    // it. This is the split that decides whether to fix retrieval or ranking.
    rankingLoss: Math.max(0, candidateRecall - recall),
    firstHitRank: firstHitIndex >= 0 ? firstHitIndex + 1 : null,
    reciprocalRank: firstHitIndex >= 0 ? 1 / (firstHitIndex + 1) : 0,
    checklistCoverage: checklistCoverage(testCase.checklist, retrievedIds),
    distractorCount: distractors.length,
  };
}

function mean(values) {
  const usable = values.filter((value) => Number.isFinite(value));
  if (usable.length === 0) return 0;
  return usable.reduce((total, value) => total + value, 0) / usable.length;
}

function aggregate(results) {
  const coverage = results
    .map((item) => item.checklistCoverage)
    .filter((value) => Number.isFinite(value));
  return {
    caseCount: results.length,
    recallAtK: mean(results.map((item) => item.recall)),
    precisionAtK: mean(results.map((item) => item.precision)),
    candidateRecall: mean(results.map((item) => item.candidateRecall)),
    rankingLoss: mean(results.map((item) => item.rankingLoss)),
    hitAnyRate: results.length
      ? results.filter((item) => item.hitAny).length / results.length
      : 0,
    checklistCoverage: coverage.length ? mean(coverage) : null,
    mrr: mean(results.map((item) => item.reciprocalRank)),
    distractorRate: results.length
      ? mean(results.map((item) => (item.returned ? item.distractorCount / item.returned : 0)))
      : 0,
  };
}

export async function runMemFuseBench({
  dataset,
  scenarioIds = null,
  dimensions = null,
  maxQuestionsPerScenario = null,
  questionIds = null,
  limit = 20,
  candidateLimit = 100,
  embedText = null,
  prefetchEmbedTexts = null,
  includeSourceTags = false,
  onProgress = null,
} = {}) {
  const scenarios = asArray(dataset?.scenarios);
  if (scenarios.length === 0) throw new Error('runMemFuseBench requires a loaded dataset');
  const embedder = typeof embedText === 'function' ? embedText : createLexicalHashEmbedder();
  const embeddingMode = typeof embedText === 'function' ? 'external' : 'lexical-hash';
  const scenarioFilter = scenarioIds?.length ? new Set(scenarioIds) : null;
  // The pgvector backend clamps its own limit to 50; mirror it so the reported
  // `limit` never overstates how many rows were actually considered.
  const effectiveLimit = Math.max(1, Math.min(50, Number(limit) || 20));

  const results = [];
  const scenarioReports = [];
  for (const scenario of scenarios) {
    if (scenarioFilter && !scenarioFilter.has(scenario.scenarioId)) continue;
    const cases = buildScenarioCases(scenario, {
      dimensions,
      maxQuestions: maxQuestionsPerScenario,
      questionIds,
    });
    if (cases.length === 0) continue;
    const corpus = buildScenarioCorpus(scenario, { includeSourceTags });
    if (typeof prefetchEmbedTexts === 'function') {
      await prefetchEmbedTexts([
        ...corpus.rows.map((row) => row.content),
        ...cases.map((testCase) => testCase.question),
      ]);
    }
    // One cache per scenario: corpus embeddings are reused across that
    // scenario's questions, which is where nearly all the cost sits.
    const embeddingCache = new Map();
    const scenarioResults = [];
    for (const testCase of cases) {
      const caseResult = await runMemFuseBenchCase({
        testCase,
        corpus,
        embedText: embedder,
        limit: effectiveLimit,
        candidateLimit,
        embeddingCache,
      });
      scenarioResults.push(caseResult);
      results.push(caseResult);
      if (typeof onProgress === 'function') {
        onProgress({ scenarioId: scenario.scenarioId, completed: results.length, ...caseResult });
      }
    }
    scenarioReports.push({
      scenarioId: scenario.scenarioId,
      corpusSize: corpus.rows.length,
      ...aggregate(scenarioResults),
    });
  }

  if (results.length === 0) throw new Error('runMemFuseBench selected zero questions');

  const byDimension = {};
  for (const dimension of new Set(results.map((item) => item.dimension))) {
    byDimension[dimension] = aggregate(results.filter((item) => item.dimension === dimension));
  }

  return {
    limit: effectiveLimit,
    candidateLimit,
    embeddingMode,
    includeSourceTags,
    overall: aggregate(results),
    byScenario: scenarioReports,
    byDimension,
    results,
  };
}

export function summarizeMemFuseBench(report) {
  const round = (value) => (Number.isFinite(value) ? Number(value.toFixed(4)) : null);
  return {
    limit: report.limit,
    embeddingMode: report.embeddingMode,
    caseCount: report.overall.caseCount,
    recallAtK: round(report.overall.recallAtK),
    precisionAtK: round(report.overall.precisionAtK),
    candidateRecall: round(report.overall.candidateRecall),
    rankingLoss: round(report.overall.rankingLoss),
    hitAnyRate: round(report.overall.hitAnyRate),
    checklistCoverage: round(report.overall.checklistCoverage),
    mrr: round(report.overall.mrr),
    distractorRate: round(report.overall.distractorRate),
    weakestDimension:
      Object.entries(report.byDimension)
        .sort((left, right) => left[1].recallAtK - right[1].recallAtK)
        .map(([dimension]) => dimension)[0] ?? null,
  };
}