Improve MemFuse recall via hybrid ranking and candidate generation.

Add RRF fusion with English word-level lexical scoring, tiered keyword fetch, and vector margin expansion (0.15/200) to fix pre-rank truncation; wire DashScope embedding bench path and update baseline to 28.8% recall@20. Co-authored-by: Cursor <cursoragent@cursor.com>
2026-09-02 13:41:12 +08:00
parent b923e54eff
commit fb3a442e73
11 changed files with 1436 additions and 97 deletions
@@ -1,5 +1,14 @@
 const DEFAULT_TABLE = 'memory_embeddings';
 const DEFAULT_LIMIT = 8;
+const VECTOR_CANDIDATE_LIMIT_MAX = 150;
+const VECTOR_SCORE_MARGIN = Math.max(
+  0,
+  Number(process.env.MEMIND_VECTOR_SCORE_MARGIN ?? 0.15) || 0.15,
+);
+const VECTOR_EXPAND_CAP = Math.max(
+  50,
+  Number(process.env.MEMIND_VECTOR_EXPAND_CAP ?? 200) || 200,
+);

 function isSafeIdentifier(value) {
  return /^[a-zA-Z_][a-zA-Z0-9_]*$/.test(String(value ?? ''));
@@ -42,7 +51,7 @@ function buildCharacterNgrams(value, size = 2) {
  return grams;
 }

-function lexicalQueryCoverage(query, text) {
+function characterNgramCoverage(query, text) {
  const queryGrams = buildCharacterNgrams(query);
  if (queryGrams.size === 0) return 0;
  const textGrams = buildCharacterNgrams(text);
@@ -53,6 +62,63 @@ function lexicalQueryCoverage(query, text) {
  return overlap / queryGrams.size;
 }

+function latinWordTokens(value) {
+  return new Set(
+    String(value ?? '')
+      .normalize('NFKC')
+      .toLowerCase()
+      .match(/[a-z0-9]{2,}/g) ?? [],
+  );
+}
+
+const ENGLISH_KEYWORD_STOP_TERMS = new Set([
+  'what', 'when', 'where', 'which', 'who', 'whom', 'whose', 'why', 'how',
+  'did', 'does', 'do', 'was', 'were', 'are', 'is', 'am', 'be', 'been', 'being',
+  'the', 'and', 'for', 'with', 'from', 'that', 'this', 'those', 'these',
+  'have', 'has', 'had', 'can', 'could', 'would', 'should', 'will', 'shall',
+  'about', 'after', 'before', 'during', 'into', 'onto', 'over', 'under',
+  'any', 'all', 'some', 'many', 'much', 'most', 'more', 'less', 'than',
+  'she', 'her', 'him', 'his', 'they', 'them', 'their', 'our', 'your', 'you',
+  'together', 'happened', 'piece', 'moment', 'until', 'timeline', 'events',
+  'everyone', 'doing', 'around', 'progress', 'based', 'want', 'know', 'full',
+  'please', 'tell', 'give', 'help', 'need', 'across', 'through', 'also',
+  'just', 'like', 'make', 'made', 'other', 'well', 'very', 'really',
+  'right', 'now', 'still', 'walk', 'currently', 'here', 'there', 'then',
+]);
+
+function latinWordQueryCoverage(query, text) {
+  const stop = ENGLISH_KEYWORD_STOP_TERMS;
+  let queryTokens = [...latinWordTokens(query)].filter(
+    (token) => token.length >= 3 && !stop.has(token),
+  );
+  if (queryTokens.length === 0) {
+    queryTokens = [...latinWordTokens(query)].filter((token) => token.length >= 3);
+  }
+  if (queryTokens.length === 0) return 0;
+  const textTokens = latinWordTokens(text);
+  let overlap = 0;
+  for (const token of queryTokens) {
+    if (textTokens.has(token)) overlap += 1;
+  }
+  return overlap / queryTokens.length;
+}
+
+function queryScriptProfile(query) {
+  const text = String(query ?? '');
+  const cjkChars = (text.match(/[\u4e00-\u9fff]/g) || []).length;
+  const latinChars = (text.match(/[a-z]/gi) || []).length;
+  if (latinChars > 0 && cjkChars === 0) return 'latin';
+  if (cjkChars > 0 && latinChars === 0) return 'cjk';
+  return 'mixed';
+}
+
+function lexicalQueryCoverage(query, text) {
+  const profile = queryScriptProfile(query);
+  if (profile === 'latin') return latinWordQueryCoverage(query, text);
+  if (profile === 'cjk') return characterNgramCoverage(query, text);
+  return Math.max(latinWordQueryCoverage(query, text), characterNgramCoverage(query, text));
+}
+
 const KEYWORD_STOP_TERMS = new Set([
  '记得', '忘记', '之前', '我们', '聊过', '讨论', '继续', '聊聊', '什么', '吗', '呢',
  '有没有', '是否', '告诉', '提到', '说过', '以前', '上次', '对话', '会话', '回忆',
@@ -63,6 +129,10 @@ export function extractKeywordTerms(query, { maxTerms = 8, minLength = 2 } = {})
  const normalized = String(query ?? '').normalize('NFKC').trim();
  if (!normalized) return [];
  const terms = new Set();
+  const properNouns = new Set();
+  for (const match of normalized.matchAll(/\b[A-Z][a-z]{2,}\b/g)) {
+    properNouns.add(match[0].toLowerCase());
+  }
  const cjkOnly = normalized.replace(/[^\u4e00-\u9fff]/gu, '');
  for (let index = 0; index < cjkOnly.length; index += 1) {
    for (const size of [4, 3, 2]) {
@@ -74,10 +144,18 @@ export function extractKeywordTerms(query, { maxTerms = 8, minLength = 2 } = {})
    }
  }
  for (const match of normalized.matchAll(/[a-z0-9]{3,}/gi)) {
-    terms.add(match[0].toLowerCase());
+    const term = match[0].toLowerCase();
+    if (!ENGLISH_KEYWORD_STOP_TERMS.has(term)) {
+      terms.add(term);
+    }
  }
  return [...terms]
-    .sort((left, right) => right.length - left.length)
+    .sort((left, right) => {
+      const leftProper = properNouns.has(left) ? 1 : 0;
+      const rightProper = properNouns.has(right) ? 1 : 0;
+      if (leftProper !== rightProper) return rightProper - leftProper;
+      return right.length - left.length;
+    })
    .slice(0, maxTerms);
 }

@@ -86,35 +164,29 @@ function dedupeContentPrefix(text, length = 96) {
 }

 const KEYWORD_FETCH_CAP = 500;
+const KEYWORD_PRIORITY_TERM_CAP = 120;
+const KEYWORD_RETURN_CAP = 100;

-async function fetchKeywordCandidates(pool, {
-  userId,
-  query,
-  tableName,
-  limit = 20,
-  maxTerms = 8,
-} = {}) {
-  const terms = extractKeywordTerms(query, { maxTerms });
-  if (!terms.length) return [];
-  const clauses = terms.map((_term, index) => `content ILIKE $${index + 2}`);
-  const params = [userId, ...terms.map((term) => `%${term}%`)];
-  const safeLimit = Math.max(1, Math.min(50, Number(limit) || 20));
-  const fetchCap = Math.min(KEYWORD_FETCH_CAP, Math.max(safeLimit, safeLimit * 10));
-  const sql = `
-    SELECT id, content, type, created_at, updated_at, 1.0 AS score
-    FROM ${tableName}
-    WHERE user_id = $1
-      AND (${clauses.join(' OR ')})
-    LIMIT $${params.length + 1}
-  `;
-  params.push(fetchCap);
-  const result = await pool.query(sql, params);
-  const rows = result?.rows ?? [];
+function splitPriorityKeywordTerms(query, terms) {
+  const priority = new Set();
+  for (const match of String(query ?? '').matchAll(/\b[A-Z][a-z]{2,}\b/g)) {
+    priority.add(match[0].toLowerCase());
+  }
+  for (const match of String(query ?? '').matchAll(/\b[A-Z]{2,}\b/g)) {
+    priority.add(match[0].toLowerCase());
+  }
+  const priorityTerms = terms.filter((term) => priority.has(term));
+  const generalTerms = terms.filter((term) => !priority.has(term));
+  return { priorityTerms, generalTerms };
+}
+
+function rankKeywordCandidateRows(rows, query, returnLimit) {
+  const safeLimit = Math.max(1, Number(returnLimit) || 20);
  return rows
    .map((row) => ({
      row,
-      lexicalScore: lexicalQueryCoverage(query, row.content ?? ''),
-      updatedAt: timestampValue(row.updated_at ?? row.created_at),
+      lexicalScore: lexicalQueryCoverage(query, row.content ?? row.text ?? ''),
+      updatedAt: timestampValue(row.updated_at ?? row.updatedAt ?? row.created_at ?? row.createdAt),
    }))
    .sort((left, right) => {
      if (left.lexicalScore !== right.lexicalScore) {
@@ -123,7 +195,156 @@ async function fetchKeywordCandidates(pool, {
      return right.updatedAt - left.updatedAt;
    })
    .slice(0, safeLimit)
-    .map(({ row, lexicalScore }) => ({ ...row, score: lexicalScore }));
+    .map(({ row }) => ({ ...row, score: null }));
+}
+
+function capRowsByLexical(rows, query, cap) {
+  if (rows.length <= cap) return rows;
+  return rows
+    .map((row) => ({
+      row,
+      lexicalScore: lexicalQueryCoverage(query, row.content ?? row.text ?? ''),
+    }))
+    .sort((left, right) => right.lexicalScore - left.lexicalScore)
+    .slice(0, cap)
+    .map(({ row }) => row);
+}
+
+/**
+ * Offline / in-memory keyword candidate selection. Priority terms (proper nouns)
+ * are fetched in separate buckets so a broad OR query cannot truncate them out
+ * before lexical ranking — the root cause of MemFuseBench candidate misses.
+ */
+export function selectKeywordCandidateRows(rows, query, terms, {
+  fetchCap = KEYWORD_FETCH_CAP,
+  returnLimit = 100,
+  priorityTermCap = 120,
+} = {}) {
+  if (!terms.length) return [];
+  const { priorityTerms, generalTerms } = splitPriorityKeywordTerms(query, terms);
+  const byId = new Map();
+  const haystacks = rows.map((row) => ({
+    row,
+    text: String(row.content ?? row.text ?? '').toLowerCase(),
+  }));
+
+  function addTermMatches(termSubset, cap) {
+    if (!termSubset.length) return;
+    const matched = haystacks
+      .filter(({ text }) => termSubset.some((term) => text.includes(term)))
+      .map(({ row }) => row);
+    for (const row of capRowsByLexical(matched, query, cap)) {
+      byId.set(String(row.id), row);
+    }
+  }
+
+  for (const term of priorityTerms.slice(0, 6)) {
+    addTermMatches([term], priorityTermCap);
+  }
+  if (generalTerms.length > 0) {
+    addTermMatches(generalTerms, fetchCap);
+  } else if (priorityTerms.length > 0) {
+    addTermMatches(priorityTerms, fetchCap);
+  }
+
+  return rankKeywordCandidateRows([...byId.values()], query, returnLimit);
+}
+
+async function queryKeywordTermSet(pool, {
+  userId,
+  tableName,
+  terms,
+  cap,
+}) {
+  if (!terms.length) return [];
+  const clauses = terms.map((_term, index) => `content ILIKE $${index + 2}`);
+  const params = [userId, ...terms.map((term) => `%${term}%`), cap];
+  const sql = `
+    SELECT id, content, type, created_at, updated_at, 1.0 AS score
+    FROM ${tableName}
+    WHERE user_id = $1
+      AND (${clauses.join(' OR ')})
+    LIMIT $${params.length}
+  `;
+  const result = await pool.query(sql, params);
+  return result?.rows ?? [];
+}
+
+async function fetchKeywordCandidates(pool, {
+  userId,
+  query,
+  tableName,
+  limit = 20,
+  maxTerms = 12,
+} = {}) {
+  const terms = extractKeywordTerms(query, { maxTerms });
+  if (!terms.length) return [];
+  const safeLimit = Math.max(1, Math.min(KEYWORD_RETURN_CAP, Number(limit) || 20));
+  const fetchCap = Math.min(KEYWORD_FETCH_CAP, Math.max(safeLimit, safeLimit * 10));
+  const { priorityTerms, generalTerms } = splitPriorityKeywordTerms(query, terms);
+  const byId = new Map();
+
+  for (const term of priorityTerms.slice(0, 6)) {
+    const rows = await queryKeywordTermSet(pool, {
+      userId,
+      tableName,
+      terms: [term],
+      cap: KEYWORD_PRIORITY_TERM_CAP,
+    });
+    for (const row of capRowsByLexical(rows, query, KEYWORD_PRIORITY_TERM_CAP)) {
+      byId.set(String(row.id), row);
+    }
+  }
+
+  const generalQueryTerms = generalTerms.length > 0 ? generalTerms : priorityTerms;
+  if (generalQueryTerms.length > 0) {
+    const rows = await queryKeywordTermSet(pool, {
+      userId,
+      tableName,
+      terms: generalQueryTerms,
+      cap: fetchCap,
+    });
+    for (const row of capRowsByLexical(rows, query, fetchCap)) {
+      byId.set(String(row.id), row);
+    }
+  }
+
+  return rankKeywordCandidateRows([...byId.values()], query, safeLimit);
+}
+
+/**
+ * Vector candidate union: top-N by cosine plus any row within `margin` of the
+ * best score (capped at expandCap), then optional recency rows for cold-start.
+ */
+export function selectVectorCandidateRows(scoredEntries, {
+  baseLimit = 100,
+  recentRows = [],
+  margin = VECTOR_SCORE_MARGIN,
+  expandCap = VECTOR_EXPAND_CAP,
+} = {}) {
+  const sorted = [...scoredEntries].sort((left, right) => right.score - left.score);
+  const topScore = sorted[0]?.score ?? 0;
+  const scoreFloor = topScore - margin;
+  const merged = new Map();
+
+  for (const entry of sorted.slice(0, baseLimit)) {
+    merged.set(String(entry.row.id), entry);
+  }
+  for (const entry of sorted) {
+    if (merged.size >= expandCap) break;
+    if (entry.score < scoreFloor) break;
+    merged.set(String(entry.row.id), entry);
+  }
+
+  for (const row of recentRows.slice(0, baseLimit)) {
+    const key = String(row.id);
+    if (!merged.has(key)) {
+      const existing = sorted.find((entry) => String(entry.row.id) === key);
+      merged.set(key, existing ?? { row, score: 0 });
+    }
+  }
+
+  return [...merged.values()];
 }

 function timestampValue(value) {
@@ -147,6 +368,55 @@ function normalizeRow(row) {
  };
 }

+const RRF_RANK_CONSTANT = 60;
+
+function rankEntries(entries, compareFn) {
+  const order = [...entries].sort(compareFn);
+  const ranks = new Map();
+  for (let index = 0; index < order.length; index += 1) {
+    ranks.set(order[index].key, index + 1);
+  }
+  return ranks;
+}
+
+function recallVectorSpread(entries) {
+  const vectorScores = entries
+    .map((entry) => entry.vectorScore)
+    .filter((score) => score >= 0);
+  if (vectorScores.length < 2) return null;
+  const sorted = [...vectorScores].sort((left, right) => right - left);
+  const max = sorted[0];
+  const median = sorted[Math.floor(sorted.length / 2)];
+  const spread = max - median;
+  return {
+    max,
+    median,
+    spread,
+    semantic: spread >= 0.20 && max >= 0.50,
+  };
+}
+
+function recallRankingWeights(query, entries) {
+  if (queryScriptProfile(query) !== 'latin') {
+    return { lexical: 1, vector: 1 };
+  }
+  const spreadInfo = recallVectorSpread(entries);
+  if (spreadInfo?.semantic) {
+    return { lexical: 0.45, vector: 1.55 };
+  }
+  return { lexical: 1, vector: 1 };
+}
+
+function recallVectorTermWeight(query, entry, spreadInfo, vectorWeight) {
+  if (!spreadInfo?.semantic || queryScriptProfile(query) !== 'latin') {
+    return vectorWeight;
+  }
+  if (entry.vectorScore < 0) return vectorWeight;
+  // Pure vector matches with zero lexical overlap often beat weak-overlap gold in RRF.
+  if (entry.lexicalScore <= 0) return vectorWeight * 0.6;
+  return vectorWeight;
+}
+
 function rankHybridCandidates(rows, query, limit) {
  const byId = new Map();
  const byPrefix = new Set();
@@ -159,24 +429,67 @@ function rankHybridCandidates(rows, query, limit) {
    const key = memory.id ?? `${memory.label}:${memory.text}`;
    if (!byId.has(key)) byId.set(key, memory);
  }
-  return [...byId.values()]
-    .map((memory) => ({
-      memory,
-      lexicalScore: lexicalQueryCoverage(query, memory.text),
-      vectorScore: Number.isFinite(memory.score) ? memory.score : -1,
-      updatedAt: timestampValue(memory.updatedAt),
-    }))
-    .sort((left, right) => {
-      if (left.lexicalScore !== right.lexicalScore) {
-        return right.lexicalScore - left.lexicalScore;
-      }
-      if (left.lexicalScore > 0 && left.updatedAt !== right.updatedAt) {
-        return right.updatedAt - left.updatedAt;
-      }
+  const entries = [...byId.values()].map((memory) => ({
+    memory,
+    key: memory.id ?? `${memory.label}:${memory.text}`,
+    lexicalScore: lexicalQueryCoverage(query, memory.text),
+    vectorScore: Number.isFinite(memory.score) ? memory.score : -1,
+    updatedAt: timestampValue(memory.updatedAt),
+  }));
+  if (entries.length === 0) return [];
+
+  const lexicalRanks = rankEntries(entries, (left, right) => {
+    if (left.lexicalScore !== right.lexicalScore) {
+      return right.lexicalScore - left.lexicalScore;
+    }
+    if (left.lexicalScore > 0 && left.updatedAt !== right.updatedAt) {
+      return right.updatedAt - left.updatedAt;
+    }
+    return right.vectorScore - left.vectorScore;
+  });
+  const vectorEligible = entries.filter((entry) => entry.vectorScore >= 0);
+  const vectorRanks = rankEntries(vectorEligible, (left, right) => {
+    if (left.vectorScore !== right.vectorScore) {
      return right.vectorScore - left.vectorScore;
+    }
+    if (left.lexicalScore !== right.lexicalScore) {
+      return right.lexicalScore - left.lexicalScore;
+    }
+    return right.updatedAt - left.updatedAt;
+  });
+
+  const spreadInfo = recallVectorSpread(entries);
+  const { lexical: lexicalWeight, vector: vectorWeight } = recallRankingWeights(query, entries);
+
+  return entries
+    .map((entry) => {
+      const lexicalRank = lexicalRanks.get(entry.key);
+      const lexicalTerm = lexicalWeight / (RRF_RANK_CONSTANT + lexicalRank);
+      const effectiveVectorWeight = recallVectorTermWeight(query, entry, spreadInfo, vectorWeight);
+      const vectorTerm = entry.vectorScore >= 0
+        ? effectiveVectorWeight / (RRF_RANK_CONSTANT + vectorRanks.get(entry.key))
+        : vectorWeight / (RRF_RANK_CONSTANT + lexicalRank);
+      return {
+        entry,
+        fusedScore: lexicalTerm + vectorTerm,
+      };
+    })
+    .sort((left, right) => {
+      if (left.fusedScore !== right.fusedScore) {
+        return right.fusedScore - left.fusedScore;
+      }
+      const leftEntry = left.entry;
+      const rightEntry = right.entry;
+      if (leftEntry.lexicalScore !== rightEntry.lexicalScore) {
+        return rightEntry.lexicalScore - leftEntry.lexicalScore;
+      }
+      if (leftEntry.vectorScore !== rightEntry.vectorScore) {
+        return rightEntry.vectorScore - leftEntry.vectorScore;
+      }
+      return rightEntry.updatedAt - leftEntry.updatedAt;
    })
    .slice(0, limit)
-    .map(({ memory }) => memory);
+    .map(({ entry }) => entry.memory);
 }

 export function createPgvectorMemoryBackend({
@@ -220,17 +533,25 @@ export function createPgvectorMemoryBackend({
      const limit = Math.max(1, Math.min(50, Number(input.limit ?? defaultLimit) || defaultLimit));
      const candidateLimit = Math.max(
        limit,
-        Math.min(100, Number(input.candidateLimit ?? 50) || 50),
+        Math.min(VECTOR_CANDIDATE_LIMIT_MAX, Number(input.candidateLimit ?? 50) || 50),
      );
      const sql = `
-        WITH vector_candidates AS (
-          SELECT id, content, type, created_at, updated_at,
-                 1 - (embedding <=> $2::vector) AS score,
-                 0 AS source_priority
+        WITH top_score AS (
+          SELECT (1 - (embedding <=> $2::vector))::float8 AS best_score
          FROM ${resolvedTableName}
          WHERE user_id = $1
          ORDER BY embedding <=> $2::vector
-          LIMIT $3
+          LIMIT 1
+        ), vector_candidates AS (
+          SELECT m.id, m.content, m.type, m.created_at, m.updated_at,
+                 (1 - (m.embedding <=> $2::vector))::float8 AS score,
+                 0 AS source_priority
+          FROM ${resolvedTableName} m
+          CROSS JOIN top_score t
+          WHERE m.user_id = $1
+            AND (1 - (m.embedding <=> $2::vector)) >= (t.best_score - $4::float8)
+          ORDER BY m.embedding <=> $2::vector
+          LIMIT $5
        ), recent_candidates AS (
          SELECT id, content, type, created_at, updated_at,
                 1 - (embedding <=> $2::vector) AS score,
@@ -249,7 +570,13 @@ export function createPgvectorMemoryBackend({
        ORDER BY id, source_priority
      `;
      const [result, keywordRows] = await Promise.all([
-        pool.query(sql, [userId, vectorLiteral(embedding), candidateLimit]),
+        pool.query(sql, [
+          userId,
+          vectorLiteral(embedding),
+          candidateLimit,
+          VECTOR_SCORE_MARGIN,
+          VECTOR_EXPAND_CAP,
+        ]),
        fetchKeywordCandidates(pool, {
          userId,
          query: input.query,
@@ -272,7 +599,12 @@ export function createPgvectorMemoryBackend({

 export const pgvectorMemoryBackendInternals = {
  lexicalQueryCoverage,
+  latinWordQueryCoverage,
+  queryScriptProfile,
+  recallRankingWeights,
  rankHybridCandidates,
  extractKeywordTerms,
+  selectKeywordCandidateRows,
+  selectVectorCandidateRows,
  dedupeContentPrefix,
 };