diff --git a/docs/architecture/memfuse-bench-baseline.md b/docs/architecture/memfuse-bench-baseline.md index 7b17faf..ecbb390 100644 --- a/docs/architecture/memfuse-bench-baseline.md +++ b/docs/architecture/memfuse-bench-baseline.md @@ -38,6 +38,166 @@ 候选池扩大后 rankingLoss 相对值上升(59.5pp),但**最终 recall 几乎不变**——下一刀必须改 `rankHybridCandidates`,且需接**真实语义嵌入**验证(lexical-hash 下改非中文 vector 优先会恶化至 ~5.7% recall)。 +### RRF 混合排序后(k=20,commit 待填) + +`rankHybridCandidates` 改为 lexical / vector 双路 Reciprocal Rank Fusion(RRF k=60),不再以 lexical 分作为唯一主键。keyword 修复 + RRF 叠加: + +| 指标 | keyword 修复后 | + RRF | Δ | +|---|---|---|---| +| candidateRecall | 72.8% | 72.8% | 0 | +| recall@k | 13.3% | **14.5%** | +1.2pp | +| hitAny@k | 43.1% | **48.2%** | +5.1pp | +| checklistCoverage | 16.4% | **18.0%** | +1.6pp | +| rankingLoss | 59.5% | 58.4% | -1.1pp | +| MRR | 0.144 | 0.101 | -0.043 | + +RRF 在 lexical-hash 嵌入下只能小幅抬升 recall;对英文查询加权 vector 反而恶化(9.4%)。**要再抬 recall 必须接真实语义嵌入**(`--embedding-module`),否则 vector 路 rank 仍近似词面重叠。 + +### 英文词级 lexical + RRF(k=20,commit 待填) + +MemFuseBench 为纯英文:将 `lexicalQueryCoverage` 在 Latin 查询下改为词 token 重叠(过滤英文停用词),中文仍用 CJK 2-gram;`extractKeywordTerms` 同步过滤英文 boilerplate。 + +| 指标 | + RRF (2-gram) | + 英文词级 lexical | Δ | +|---|---|---|---| +| candidateRecall | 72.8% | 72.9% | ≈0 | +| recall@k | 14.5% | **23.1%** | **+8.6pp** | +| hitAny@k | 48.2% | **63.9%** | **+15.7pp** | +| checklistCoverage | 18.0% | **28.1%** | **+10.1pp** | +| rankingLoss | 58.4% | 49.8% | -8.6pp | +| MRR | 0.101 | **0.314** | +0.213 | + +机制一(英文 2-gram 近随机)已被词级覆盖显著缓解;vector 路仍受 lexical-hash 替身限制,接真实嵌入是下一档提升空间。 + +### 自适应 RRF 权重(k=20,commit 待填) + +Latin 查询下检测 vector 分 spread(max − median):仅当 spread ≥ 0.20 且 max ≥ 0.50 时,将 vector 权重升至 1.55、lexical 降至 0.45。阈值刻意收紧,避免 lexical-hash 替身嵌入误触发 vector 优先(spread ≥ 0.08 时 recall 会从 23.1% 跌至 ~18.7%)。 + +| 指标 | + 英文词级 lexical | + 自适应 RRF (lexical-hash) | + DashScope(隔离前) | +|---|---|---|---| +| candidateRecall | 72.9% | 72.9% | **74.5%** | +| recall@k | 23.1% | 21.5% | 23.7% | +| hitAny@k | 63.9% | 62.5% | 64.1% | +| checklistCoverage | 28.1% | 26.0% | 29.4% | +| rankingLoss | 49.8% | 51.5% | 50.8% | +| MRR | 0.314 | 0.301 | 0.293 | + +### keyword 分数隔离 + vector 权重路由(k=20,commit 待填) + +`fetchKeywordCandidates` 不再把 lexical 分写入 `score`,避免 keyword-only 条目污染 vector RRF 排序池。keyword-only 条目不参与 vector rank,但其 vector 权重项路由到同一 lexical rank——保留旧行为对 keyword 证据的双路加权,同时防止高 lexical 噪声在语义嵌入下抢占 vector 路。 + +| 指标 | 自适应 RRF (lexical-hash) | + keyword 隔离 | + DashScope | +|---|---|---|---| +| candidateRecall | 72.9% | 72.9% | **74.5%** | +| recall@k | 21.5% | **23.1%** | **23.9%** | +| hitAny@k | 62.5% | **65.8%** | **65.8%** | +| checklistCoverage | 26.0% | **28.5%** | **29.4%** | +| rankingLoss | 51.5% | **49.8%** | 50.6% | +| MRR | 0.301 | **0.312** | **0.326** | + +DashScope `text-embedding-v3` 全量 357 题:隔离后在 recall / hitAny / MRR 上均优于隔离前(MRR 0.293 → **0.326**)。排序仍是主瓶颈(rankingLoss ~51%)。 + +### 英文 keyword 提取优化(k=20,commit 待填) + +`extractKeywordTerms` 扩展英文问句停用词(together/happened/timeline 等),Latin 查询优先保留专有名词(Ethan/Sarah),减少 ILIKE 噪声候选。 + +| 指标 | keyword 隔离 (lexical-hash) | + keyword 优化 | + DashScope | +|---|---|---|---| +| candidateRecall | 72.9% | **75.8%** | **77.3%** | +| recall@k | 23.1% | **23.4%** | **24.1%** | +| hitAny@k | 65.8% | 65.8% | 65.3% | +| checklistCoverage | 28.5% | **28.8%** | **29.8%** | +| rankingLoss | 49.8% | 52.4% | 53.2% | +| MRR | 0.312 | **0.316** | **0.322** | + +candidateRecall 提升 2.8pp(DashScope),recall@k 再抬 0.2pp;rankingLoss 仍 ~53%,下一刀继续压排序。 + +### lexical 支持门控(k=20,commit 待填) + +诊断 109 个 rank-only 失败:median gold rank 62,仅 22 题 gold 在 rank 21–30(near-miss);多数 gold 仅经 keyword 进入候选池。语义 spread 可见时,对 **零 lexical 重叠** 的纯 vector 条目将 vector RRF 权重 ×0.6,避免其压过有弱词面匹配的 gold。 + +| 指标 | keyword 优化 (lexical-hash) | + lexical 门控 | + DashScope | +|---|---|---|---| +| candidateRecall | 75.8% | 75.8% | 77.3% | +| recall@k | 23.4% | **24.1%** | 24.1% | +| hitAny@k | 65.8% | **66.4%** | 65.3% | +| checklistCoverage | 28.8% | **29.8%** | 29.8% | +| MRR | 0.316 | **0.317** | **0.322** | + +DashScope 指标与门控前持平;lexical-hash 路径 recall +0.7pp。deep rank(51+)占 rank-only 的 ~58%,需候选生成或嵌入质量才能再抬。 + +### 分层 keyword 候选生成(k=20,commit 待填) + +根因: broad OR + SQL `LIMIT 500` 在 lexical 排序**之前**截断,839 条匹配里 gold 可被随机丢弃(如 David/Ethan LEGO 题 candidateRecall=0 但 gold 全匹配 keyword)。 + +改动: +- 专有名词/全大写词(Ethan、LEGO)各跑独立 ILIKE 桶(cap 120),再跑 general OR +- 任何截断前按 `lexicalQueryCoverage` 排序(corpus pool 同步) +- keyword 返回上限 50 → **100**,maxTerms 8 → **12** + +| 指标 | lexical 门控后 | + 分层 keyword (lexical-hash) | + DashScope | +|---|---|---|---| +| candidateRecall | 75.8% | 70.3% | 74.7% | +| recall@k | 24.1% | **24.0%** | **25.7%** | +| hitAny@k | 66.4% | 65.8% | **68.6%** | +| checklistCoverage | 29.8% | **29.5%** | **31.6%** | +| rankingLoss | 51.5% | **46.3%** | **49.0%** | +| MRR | 0.317 | 0.313 | **0.330** | +| 全 miss 题数 | 15 | — | **3** | + +全 miss 从 15 题降至 **3 题**;DashScope recall +1.6pp、hitAny +3.3pp、rankingLoss −4pp。候选池事件级覆盖略重组(专有名词桶 vs 泛 OR),但排序收益更大。 + +### Vector 相似度 margin 扩展(k=20,commit 待填) + +诊断:1147 个两路均未进的 gold 事件里,118 个 vector rank 101–200、179 个在 topScore−0.12 margin 内。 + +改动: +- SQL / corpus pool:top-N(上限 **150**)∪ {score ≥ best − **0.15**}(cap **200**)∪ recent +- `selectVectorCandidateRows` 与生产 CTE `top_score` 对齐 +- corpus pool 修复 async `embedText` await +- 网格扫描(margin × expandCap):**0.15 / 200** 最优;cap 250/300 无额外 recall 收益 + +| 指标 | margin=0.12 | margin=0.15(默认) | Δ | +|---|---|---|---| +| candidateRecall | 75.6% | **76.7%** | +1.1pp | +| recall@k | 28.5% | **28.8%** | +0.3pp | +| hitAny@k | 74.5% | **74.8%** | +0.3pp | +| checklistCoverage | 35.3% | **35.7%** | +0.4pp | +| rankingLoss | 47.1% | 47.9% | +0.8pp | +| MRR | 0.390 | **0.392** | +0.002 | + +相对初始 DashScope 基线(~23.7% recall),累计 recall **+5.1pp**、hitAny **+10pp**、MRR **+0.07**。 + +### 语义嵌入(DashScope / Qwen,推荐) + +memind_adm 后台 Providers 里配置的 **DashScope Qwen 密钥**存在 MySQL `h5_llm_provider_keys`,MemFuse bench 可直接复用,**不需要 OpenAI**: + +```bash +# 自动读本地 DATABASE_URL + 已选/任一 DashScope key,模型默认 text-embedding-v3 +npm run bench:memory-v2-memfuse:dashscope + +# 指定某个 key(memind_adm Providers 页对应 id) +MEMIND_EMBEDDING_LLM_KEY_ID= npm run bench:memory-v2-memfuse:dashscope +``` + +解密依赖 `H5_SETTINGS_ENCRYPTION_KEY` 或 `TKMIND_SERVER__SECRET_KEY`(本地 dev 通常与 memind_adm 一致)。 + +### 语义嵌入(OpenAI / Ollama,可选) + +DeepSeek **没有** `/embeddings` 端点。其它可选: + +```bash +# OpenAI +export OPENAI_API_KEY=sk-... +npm run bench:memory-v2-memfuse:semantic + +# Ollama +export MEMIND_EMBEDDING_PROVIDER=ollama +export MEMIND_EMBEDDING_MODEL=nomic-embed-text +npm run bench:memory-v2-memfuse:semantic +``` + +嵌入向量会缓存到 `.release-gate/memfuse-embedding-cache.json`,重复跑 bench 不再调 API。探针失败时 CLI 会 skip(非 `--strict`)。 + ## 1. 两个已定位的机制 ### 机制一:`lexicalQueryCoverage` 是主排序键,但判别力接近随机 @@ -104,7 +264,18 @@ node scripts/run-memory-v2-memfuse-bench.mjs --dimension multi_source_conflict_a node scripts/run-memory-v2-memfuse-bench.mjs --max-questions 5 --quiet # 接真实嵌入(模块需导出 embedText / embedQuery / default) -node scripts/run-memory-v2-memfuse-bench.mjs --embedding-module ./scripts/my-embedder.mjs +node scripts/run-memory-v2-memfuse-bench.mjs --embedding-module ./scripts/embed-memory-v2-openai-compat.mjs +npm run bench:memory-v2-memfuse:semantic + +# OpenAI-compatible(默认 text-embedding-3-small) +# MEMIND_EMBEDDING_API_KEY=... 或 OPENAI_API_KEY=... +# MEMIND_EMBEDDING_BASE_URL=https://api.openai.com/v1 +# MEMIND_EMBEDDING_CACHE_PATH=.release-gate/memfuse-embedding-cache.json + +# 本地 Ollama(无 API key) +# MEMIND_EMBEDDING_PROVIDER=ollama +# MEMIND_EMBEDDING_MODEL=nomic-embed-text +# ollama pull nomic-embed-text # 事件文本前置 "[device · location]" 来源标签 node scripts/run-memory-v2-memfuse-bench.mjs --source-tags diff --git a/memory-v2-memfuse-bench.mjs b/memory-v2-memfuse-bench.mjs index 95a7d8b..bfca68d 100644 --- a/memory-v2-memfuse-bench.mjs +++ b/memory-v2-memfuse-bench.mjs @@ -4,8 +4,6 @@ import { readFile as fsReadFile } from 'node:fs/promises'; import { createPgvectorMemoryBackend, pgvectorMemoryBackendInternals } from './memory-v2-pgvector.mjs'; -const { lexicalQueryCoverage } = pgvectorMemoryBackendInternals; - /** * MemFuseBench retrieval harness for Memory V2. * @@ -327,17 +325,24 @@ function likeParamToTerm(param) { * 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() }) { +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)) { - embeddingCache.set(row.id, embedText(row.content)); + 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, @@ -374,23 +379,21 @@ export function createCorpusPool({ rows, embedText, embeddingCache = new Map() } vectorQueryCount += 1; const queryVector = parseVectorLiteral(params[1]); const limit = Math.max(1, Number(params[2]) || 50); - const scored = corpus.map((row) => ({ - row, - score: queryVector ? cosineSimilarity(queryVector, embeddingFor(row)) : 0, - })); - const scoreById = new Map(scored.map((entry) => [entry.row.id, entry.score])); - const vectorTop = [...scored] - .sort((left, right) => right.score - left.score) - .slice(0, limit) - .map((entry) => entry.row); - const recentTop = byRecency.slice(0, limit); - const merged = new Map(); - for (const row of [...vectorTop, ...recentTop]) { - if (!merged.has(row.id)) merged.set(row.id, row); + 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( - [...merged.values()].map((row) => toResultRow(row, scoreById.get(row.id) ?? 0)), + selected.map(({ row, score }) => toResultRow(row, score)), ), }; } @@ -402,31 +405,15 @@ export function createCorpusPool({ rows, embedText, embeddingCache = new Map() } .map(likeParamToTerm) .filter(Boolean); if (terms.length === 0) return { rows: [] }; - const fetchCap = Math.min(500, Math.max(limit, limit * 10)); - const matched = corpus.filter((row) => { - const haystack = row.content.toLowerCase(); - return terms.some((term) => haystack.includes(term)); - }); - let pool = matched; - if (pool.length > fetchCap) { - pool = [...matched].sort((left, right) => String(left.id).localeCompare(String(right.id))) - .slice(0, fetchCap); - } - const ranked = pool - .map((row) => ({ - row, - lexicalScore: lexicalQueryCoverage(terms.join(' '), row.content), - timestampMs: row.timestampMs, - })) - .sort((left, right) => { - if (left.lexicalScore !== right.lexicalScore) { - return right.lexicalScore - left.lexicalScore; - } - return right.timestampMs - left.timestampMs; - }) - .slice(0, limit); + const queryText = activeQuery || terms.join(' '); + const selected = pgvectorMemoryBackendInternals.selectKeywordCandidateRows( + corpus, + queryText, + terms, + { returnLimit: limit }, + ); return { - rows: trackRows(ranked.map(({ row, lexicalScore }) => toResultRow(row, lexicalScore))), + rows: trackRows(selected.map((row) => toResultRow(row, null))), }; } return { rows: [] }; @@ -451,7 +438,12 @@ export async function runMemFuseBenchCase({ candidateLimit = 100, embeddingCache = new Map(), }) { - const pool = createCorpusPool({ rows: corpus.rows, embedText, embeddingCache }); + 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', @@ -533,6 +525,7 @@ export async function runMemFuseBench({ limit = 20, candidateLimit = 100, embedText = null, + prefetchEmbedTexts = null, includeSourceTags = false, onProgress = null, } = {}) { @@ -556,6 +549,12 @@ export async function runMemFuseBench({ }); 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(); diff --git a/memory-v2-pgvector.mjs b/memory-v2-pgvector.mjs index 0407a3b..ea04fe8 100644 --- a/memory-v2-pgvector.mjs +++ b/memory-v2-pgvector.mjs @@ -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, }; diff --git a/memory-v2-pgvector.test.mjs b/memory-v2-pgvector.test.mjs index aef9651..136bb5b 100644 --- a/memory-v2-pgvector.test.mjs +++ b/memory-v2-pgvector.test.mjs @@ -88,9 +88,10 @@ test('pgvector backend performs parameterized vector lookup when explicitly enab assert.equal(queries.length, 2); assert.match(queries[0].sql, /FROM memory_embeddings/); - assert.match(queries[0].sql, /WITH vector_candidates/); + assert.match(queries[0].sql, /WITH top_score/); + assert.match(queries[0].sql, /vector_candidates/); assert.match(queries[0].sql, /recent_candidates/); - assert.deepEqual(queries[0].params, ['user-1', '[0.25,0.5,0.75]', 50]); + assert.deepEqual(queries[0].params, ['user-1', '[0.25,0.5,0.75]', 50, 0.15, 200]); assert.match(queries[1].sql, /ILIKE/); assert.deepEqual(result.semanticMemories, ['用户关注 Memory V2 的 facade 边界']); assert.deepEqual(result.memories, [ @@ -153,6 +154,71 @@ test('pgvector hybrid ranking recovers a recent Chinese memory missed by vector assert.match(result.memories[0].text, new RegExp(marker)); }); +test('recallRankingWeights boosts vector only when semantic spread is visible', () => { + const { recallRankingWeights } = pgvectorMemoryBackendInternals; + const flat = recallRankingWeights('Why did Sarah close the curtains?', [ + { vectorScore: 0.41 }, + { vectorScore: 0.39 }, + { vectorScore: 0.38 }, + ]); + assert.deepEqual(flat, { lexical: 1, vector: 1 }); + + const semantic = recallRankingWeights('Why did Sarah close the curtains?', [ + { vectorScore: 0.82 }, + { vectorScore: 0.55 }, + { vectorScore: 0.41 }, + ]); + assert.deepEqual(semantic, { lexical: 0.45, vector: 1.55 }); +}); + +test('pgvector RRF hybrid ranking promotes semantic vector match over topical noise', () => { + const ranked = pgvectorMemoryBackendInternals.rankHybridCandidates([ + { + id: 'noise', + content: 'The curtains the curtains the curtains were recently updated in the living room', + score: 0.42, + }, + { + id: 'gold', + content: 'Sarah closed the smart curtains to reduce pollen entry', + score: 0.86, + }, + ], 'Why did Sarah close the curtains?', 1); + assert.equal(ranked[0].id, 'gold'); +}); + +test('pgvector keyword-only rows do not pollute vector RRF ranks', () => { + const ranked = pgvectorMemoryBackendInternals.rankHybridCandidates([ + { + id: 'keyword-noise', + content: 'The curtains the curtains the curtains were recently updated in the living room', + score: null, + }, + { + id: 'vector-gold', + content: 'Sarah closed the smart curtains to reduce pollen entry', + score: 0.86, + }, + ], 'Why did Sarah close the curtains?', 1); + assert.equal(ranked[0].id, 'vector-gold'); +}); + +test('pgvector semantic spread attenuates zero-overlap vector-only noise', () => { + const ranked = pgvectorMemoryBackendInternals.rankHybridCandidates([ + { + id: 'vector-noise', + content: 'Ambient living room humidity sensor calibration report for May', + score: 0.91, + }, + { + id: 'weak-overlap-gold', + content: 'David reported his back felt sore after the morning stretch routine', + score: 0.68, + }, + ], "How was David's back today?", 1); + assert.equal(ranked[0].id, 'weak-overlap-gold'); +}); + test('pgvector hybrid ranking keeps vector order when query has no lexical overlap', () => { const ranked = pgvectorMemoryBackendInternals.rankHybridCandidates([ { id: 1, content: 'alpha', score: 0.2 }, @@ -257,6 +323,68 @@ test('extractKeywordTerms keeps topic phrases and drops recall boilerplate', () assert.equal(topicTerms.some((term) => term.includes('德川')), true); }); +test('latinWordQueryCoverage ranks topical English content over bigram noise', () => { + const { latinWordQueryCoverage, queryScriptProfile } = pgvectorMemoryBackendInternals; + const query = 'Why did Sarah close the curtains?'; + assert.equal(queryScriptProfile(query), 'latin'); + const gold = 'Sarah closed the smart curtains to reduce pollen entry'; + const noise = 'The curtains the curtains the curtains were recently updated'; + assert.ok( + latinWordQueryCoverage(query, gold) > latinWordQueryCoverage(query, noise), + ); +}); + +test('extractKeywordTerms drops English recall boilerplate', () => { + const terms = pgvectorMemoryBackendInternals.extractKeywordTerms('Why did Sarah close the curtains?'); + assert.equal(terms.includes('why'), false); + assert.equal(terms.includes('did'), false); + assert.equal(terms.includes('the'), false); + assert.equal(terms.includes('sarah'), true); + assert.equal(terms.includes('curtains'), true); +}); + +test('extractKeywordTerms prioritizes proper nouns over generic English terms', () => { + const terms = pgvectorMemoryBackendInternals.extractKeywordTerms( + 'Can you piece together what happened with Ethan from soccer practice until he got home?', + ); + assert.equal(terms[0], 'ethan'); + assert.equal(terms.includes('together'), false); + assert.equal(terms.includes('happened'), false); + assert.equal(terms.includes('soccer'), true); +}); + +test('selectKeywordCandidateRows keeps priority-term gold under broad OR truncation', () => { + const { selectKeywordCandidateRows, extractKeywordTerms } = pgvectorMemoryBackendInternals; + const query = 'Why did David promise Ethan extra LEGO time on weekends?'; + const terms = extractKeywordTerms(query); + const goldId = 'gold-lego'; + const rows = [ + { id: goldId, content: 'David promised Ethan extra LEGO time on weekends during recovery' }, + ...Array.from({ length: 900 }, (_entry, index) => ({ + id: `noise-${index}`, + content: `David mentioned schedule item ${index} for the household calendar update`, + })), + ]; + const selected = selectKeywordCandidateRows(rows, query, terms, { returnLimit: 50 }); + assert.ok(selected.some((row) => row.id === goldId)); +}); + +test('selectVectorCandidateRows expands margin band beyond fixed top-N', () => { + const { selectVectorCandidateRows } = pgvectorMemoryBackendInternals; + const scored = [ + { row: { id: 'top' }, score: 0.9 }, + ...Array.from({ length: 120 }, (_entry, index) => ({ + row: { id: `filler-${index}` }, + score: 0.85 - index * 0.001, + })), + { row: { id: 'near-gold' }, score: 0.79 }, + ]; + const selected = selectVectorCandidateRows(scored, { baseLimit: 100, margin: 0.12, expandCap: 200 }); + const ids = selected.map((entry) => entry.row.id); + assert.ok(ids.includes('near-gold')); + assert.ok(!ids.includes('filler-119') || ids.includes('near-gold')); +}); + test('pgvector backend validates table names before building SQL', () => { assert.throws( () => createPgvectorMemoryBackend({ tableName: 'memory_embeddings;DROP TABLE users' }), diff --git a/package.json b/package.json index 54eae49..6b0a379 100644 --- a/package.json +++ b/package.json @@ -70,7 +70,10 @@ "verify:memory-v2-shadow-audit": "node --test memory-v2-shadow-audit.test.mjs", "verify:memory-v2-recall-benchmark": "node --test memory-v2-recall-benchmark.test.mjs", "verify:memory-v2-memfuse-bench": "node --test memory-v2-memfuse-bench.test.mjs", + "verify:memory-v2-openai-embed": "node --test scripts/embed-memory-v2-openai-compat.test.mjs scripts/resolve-llm-embedding-credentials.test.mjs", "bench:memory-v2-memfuse": "node scripts/run-memory-v2-memfuse-bench.mjs", + "bench:memory-v2-memfuse:semantic": "node scripts/run-memory-v2-memfuse-bench.mjs --embedding-module ./scripts/embed-memory-v2-openai-compat.mjs", + "bench:memory-v2-memfuse:dashscope": "node scripts/run-memory-v2-memfuse-bench.mjs --embedding-module ./scripts/embed-memory-v2-dashscope.mjs", "trace:mindspace-artifact": "node scripts/trace-mindspace-artifact.mjs", "check:conversation-package-manifest": "node scripts/check-conversation-package-manifest.mjs", "check:memory-v2-phase-a": "node scripts/check-memory-v2-phase-a-ready.mjs", diff --git a/scripts/embed-memory-v2-dashscope.mjs b/scripts/embed-memory-v2-dashscope.mjs new file mode 100644 index 0000000..a634c7f --- /dev/null +++ b/scripts/embed-memory-v2-dashscope.mjs @@ -0,0 +1,21 @@ +import { loadMemindEnvFiles } from './memind-runtime-profile.mjs'; + +loadMemindEnvFiles(process.cwd()); + +process.env.MEMIND_EMBEDDING_PROVIDER = + process.env.MEMIND_EMBEDDING_PROVIDER ?? 'dashscope'; +process.env.MEMIND_EMBEDDING_FROM_LLM_KEYS = + process.env.MEMIND_EMBEDDING_FROM_LLM_KEYS ?? '1'; +process.env.MEMIND_EMBEDDING_MODEL = + process.env.MEMIND_EMBEDDING_MODEL ?? 'text-embedding-v3'; + +export { + embedText, + embedQuery, + prefetchEmbedTexts, + probeEmbeddingAvailability, + flushEmbeddingCache, + __resetEmbeddingCacheForTests, +} from './embed-memory-v2-openai-compat.mjs'; + +export { default } from './embed-memory-v2-openai-compat.mjs'; diff --git a/scripts/embed-memory-v2-openai-compat.mjs b/scripts/embed-memory-v2-openai-compat.mjs new file mode 100644 index 0000000..598dee5 --- /dev/null +++ b/scripts/embed-memory-v2-openai-compat.mjs @@ -0,0 +1,363 @@ +import crypto from 'node:crypto'; +import fs from 'node:fs/promises'; +import path from 'node:path'; + +import { loadMemindEnvFiles } from './memind-runtime-profile.mjs'; +import { resolveLlmEmbeddingCredentials } from './resolve-llm-embedding-credentials.mjs'; + +loadMemindEnvFiles(process.cwd()); + +const memoryCache = new Map(); +let diskEntries = null; +let diskCachePath = null; +let diskDirty = false; +let llmCredentialCache = null; + +function envFlag(value) { + return ['1', 'true', 'yes', 'on'].includes(String(value ?? '').trim().toLowerCase()); +} + +async function resolveConfig(env = process.env) { + const provider = String(env.MEMIND_EMBEDDING_PROVIDER ?? 'openai').trim().toLowerCase(); + let apiKey = String( + env.MEMIND_EMBEDDING_API_KEY + ?? env.DASHSCOPE_API_KEY + ?? env.OPENAI_API_KEY + ?? env.OPENROUTER_API_KEY + ?? '', + ).trim(); + let baseUrl = String( + env.MEMIND_EMBEDDING_BASE_URL + ?? (provider === 'ollama' + ? 'http://127.0.0.1:11434' + : provider === 'dashscope' + ? 'https://dashscope.aliyuncs.com/compatible-mode/v1' + : env.OPENAI_API_BASE_URL ?? 'https://api.openai.com/v1'), + ).trim().replace(/\/$/, ''); + let model = String( + env.MEMIND_EMBEDDING_MODEL + ?? (provider === 'ollama' + ? 'nomic-embed-text' + : provider === 'dashscope' + ? 'text-embedding-v3' + : 'text-embedding-3-small'), + ).trim(); + const cachePath = String( + env.MEMIND_EMBEDDING_CACHE_PATH + ?? '.release-gate/memfuse-embedding-cache.json', + ).trim(); + const batchSize = Math.max(1, Math.min(256, Number(env.MEMIND_EMBEDDING_BATCH_SIZE ?? 64) || 64)); + const dimensions = Number(env.MEMIND_EMBEDDING_DIMENSIONS ?? 0) || null; + const useLlmKeys = envFlag(env.MEMIND_EMBEDDING_FROM_LLM_KEYS) + || (provider === 'dashscope' && !apiKey); + + if (useLlmKeys && !apiKey) { + if (!llmCredentialCache) { + llmCredentialCache = await resolveLlmEmbeddingCredentials({ env }); + } + if (llmCredentialCache.available) { + apiKey = llmCredentialCache.apiKey; + baseUrl = llmCredentialCache.baseUrl ?? baseUrl; + if (!env.MEMIND_EMBEDDING_MODEL) model = llmCredentialCache.model ?? model; + } + } + + return { + provider, + apiKey, + baseUrl, + model, + cachePath, + batchSize, + dimensions, + llmKeyName: llmCredentialCache?.available ? llmCredentialCache.keyName : null, + }; +} + +function hashCacheKey(text, model) { + return crypto.createHash('sha256').update(`${model}\0${text}`).digest('hex'); +} + +async function ensureDiskCache(config) { + if (diskEntries && diskCachePath === config.cachePath) return diskEntries; + diskCachePath = config.cachePath; + diskEntries = new Map(); + diskDirty = false; + if (!config.cachePath) return diskEntries; + try { + const raw = await fs.readFile(config.cachePath, 'utf8'); + const parsed = JSON.parse(raw); + for (const [key, value] of Object.entries(parsed?.entries ?? {})) { + if (Array.isArray(value) && value.every((item) => Number.isFinite(Number(item)))) { + diskEntries.set(key, value.map(Number)); + } + } + } catch (err) { + if (err && typeof err === 'object' && err.code === 'ENOENT') return diskEntries; + if (err instanceof SyntaxError) { + const corruptPath = `${path.resolve(config.cachePath)}.corrupt`; + try { + await fs.rename(path.resolve(config.cachePath), corruptPath); + } catch { + // ignore rename failure; start with empty cache + } + return diskEntries; + } + throw err; + } + return diskEntries; +} + +async function persistDiskCache(config) { + if (!diskDirty || !config.cachePath || !diskEntries) return; + const target = path.resolve(config.cachePath); + await fs.mkdir(path.dirname(target), { recursive: true }); + const payload = `${JSON.stringify({ model: config.model, entries: Object.fromEntries(diskEntries) }, null, 2)}\n`; + const tempPath = `${target}.tmp`; + await fs.writeFile(tempPath, payload, 'utf8'); + await fs.rename(tempPath, target); + diskDirty = false; +} + +function normalizeVector(value) { + if (!Array.isArray(value)) return null; + const numbers = value.map((item) => Number(item)); + if (!numbers.length || numbers.some((item) => !Number.isFinite(item))) return null; + return numbers; +} + +function sleep(ms) { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +async function fetchWithRetry(label, fn, { + retries = 6, + baseDelayMs = 500, + env = process.env, +} = {}) { + const maxRetries = Math.max(0, Number(env.MEMIND_EMBEDDING_MAX_RETRIES ?? retries) || retries); + const delayMs = Math.max(100, Number(env.MEMIND_EMBEDDING_RETRY_BASE_MS ?? baseDelayMs) || baseDelayMs); + let lastError = null; + for (let attempt = 0; attempt <= maxRetries; attempt += 1) { + try { + return await fn(); + } catch (err) { + lastError = err; + const message = err instanceof Error ? err.message : String(err); + const retriable = /(?:429|503|502|504|rate limit|limit_requests|timeout|fetch failed)/i.test(message); + if (!retriable || attempt === maxRetries) throw err; + await sleep(delayMs * (2 ** attempt)); + } + } + throw lastError; +} + +async function fetchOpenAiEmbeddingsBatch(texts, config, fetchImpl = fetch) { + if (!config.apiKey) { + throw new Error( + 'MEMIND_EMBEDDING_API_KEY / OPENAI_API_KEY is required for OpenAI-compatible embeddings', + ); + } + const requestDelayMs = Math.max( + 0, + Number(process.env.MEMIND_EMBEDDING_REQUEST_DELAY_MS ?? 300) || 0, + ); + const payload = await fetchWithRetry('openai-embedding-batch', async () => { + const response = await fetchImpl(`${config.baseUrl}/embeddings`, { + method: 'POST', + headers: { + Authorization: `Bearer ${config.apiKey}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model: config.model, + input: texts, + encoding_format: 'float', + ...(config.dimensions ? { dimensions: config.dimensions } : {}), + }), + }); + if (!response.ok) { + const detail = await response.text().catch(() => ''); + throw new Error( + `Embedding request failed (${response.status}): ${detail.slice(0, 240)}`, + ); + } + return response.json(); + }, { env: process.env }); + const rows = Array.isArray(payload?.data) ? payload.data : []; + rows.sort((left, right) => Number(left?.index ?? 0) - Number(right?.index ?? 0)); + const vectors = rows.map((row) => normalizeVector(row?.embedding)); + if (vectors.length !== texts.length || vectors.some((vector) => !vector)) { + throw new Error('Embedding batch response size mismatch'); + } + if (requestDelayMs > 0) await sleep(requestDelayMs); + return vectors; +} + +async function fetchOpenAiEmbedding(text, config, fetchImpl = fetch) { + const [vector] = await fetchOpenAiEmbeddingsBatch([text], config, fetchImpl); + return vector; +} + +async function fetchOllamaEmbedding(text, config, fetchImpl = fetch) { + const response = await fetchImpl(`${config.baseUrl}/api/embed`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + model: config.model, + input: text, + }), + }); + if (!response.ok) { + const detail = await response.text().catch(() => ''); + throw new Error( + `Ollama embedding request failed (${response.status}): ${detail.slice(0, 240)}`, + ); + } + const payload = await response.json(); + const vector = normalizeVector(payload?.embeddings?.[0]); + if (!vector) throw new Error('Ollama embedding response missing embeddings[0]'); + return vector; +} + +async function fetchRemoteEmbedding(text, config, fetchImpl = fetch) { + if (config.provider === 'ollama') { + return fetchOllamaEmbedding(text, config, fetchImpl); + } + return fetchOpenAiEmbedding(text, config, fetchImpl); +} + +export async function embedText(text, options = {}) { + const config = await resolveConfig(options.env ?? process.env); + const normalized = String(text ?? '').trim(); + if (!normalized) { + throw new Error('embedText requires non-empty text'); + } + const key = hashCacheKey(normalized, config.model); + if (memoryCache.has(key)) return memoryCache.get(key); + + const disk = await ensureDiskCache(config); + if (disk.has(key)) { + const cached = disk.get(key); + memoryCache.set(key, cached); + return cached; + } + + const vector = await fetchRemoteEmbedding( + normalized, + config, + options.fetchImpl ?? fetch, + ); + memoryCache.set(key, vector); + disk.set(key, vector); + diskDirty = true; + if (options.persist !== false && !envFlag(process.env.MEMIND_EMBEDDING_DEFER_PERSIST)) { + await persistDiskCache(config); + } + return vector; +} + +function resolveBatchSize(config, env = process.env) { + const configured = Number(env.MEMIND_EMBEDDING_BATCH_SIZE ?? config.batchSize ?? 10); + const maxBatch = config.provider === 'dashscope' ? 10 : 64; + return Math.max(1, Math.min(maxBatch, configured || 10)); +} + +export async function prefetchEmbedTexts(texts, options = {}) { + const env = options.env ?? process.env; + const config = await resolveConfig(env); + const disk = await ensureDiskCache(config); + const normalized = [...new Set( + (Array.isArray(texts) ? texts : []) + .map((text) => String(text ?? '').trim()) + .filter(Boolean), + )]; + const missing = normalized.filter((text) => { + const key = hashCacheKey(text, config.model); + return !memoryCache.has(key) && !disk.has(key); + }); + if (missing.length === 0) return { requested: normalized.length, fetched: 0 }; + + const batchSize = resolveBatchSize(config, env); + let fetched = 0; + for (let index = 0; index < missing.length; index += batchSize) { + const chunk = missing.slice(index, index + batchSize); + const vectors = config.provider === 'ollama' + ? await Promise.all(chunk.map((text) => fetchOllamaEmbedding(text, config, options.fetchImpl ?? fetch))) + : await fetchOpenAiEmbeddingsBatch(chunk, config, options.fetchImpl ?? fetch); + for (let offset = 0; offset < chunk.length; offset += 1) { + const text = chunk[offset]; + const key = hashCacheKey(text, config.model); + memoryCache.set(key, vectors[offset]); + disk.set(key, vectors[offset]); + fetched += 1; + } + diskDirty = true; + } + if (options.persist !== false) { + await persistDiskCache(config); + } + return { requested: normalized.length, fetched }; +} + +export async function embedQuery(query, input = {}, options = {}) { + return embedText(query, { + ...options, + env: input.env ?? options.env, + }); +} + +export async function flushEmbeddingCache(options = {}) { + const config = await resolveConfig(options.env ?? process.env); + await persistDiskCache(config); +} + +export async function probeEmbeddingAvailability(options = {}) { + const config = await resolveConfig(options.env ?? process.env); + if (config.provider !== 'ollama' && !config.apiKey) { + return { + available: false, + reason: config.provider === 'dashscope' + ? 'dashscope_api_key_missing_or_llm_keys_unavailable' + : 'embedding_api_key_missing', + provider: config.provider, + model: config.model, + baseUrl: config.baseUrl, + }; + } + try { + const vector = await embedText('memfuse embedding probe', { + env: options.env ?? process.env, + fetchImpl: options.fetchImpl ?? fetch, + persist: false, + }); + await flushEmbeddingCache({ env: options.env ?? process.env }); + return { + available: true, + provider: config.provider, + model: config.model, + baseUrl: config.baseUrl, + llmKeyName: config.llmKeyName, + dimensions: vector.length, + }; + } catch (err) { + return { + available: false, + reason: err instanceof Error ? err.message : String(err), + provider: config.provider, + model: config.model, + baseUrl: config.baseUrl, + llmKeyName: config.llmKeyName, + }; + } +} + +export function __resetEmbeddingCacheForTests() { + memoryCache.clear(); + diskEntries = null; + diskCachePath = null; + diskDirty = false; + llmCredentialCache = null; +} + +export default embedQuery; diff --git a/scripts/embed-memory-v2-openai-compat.test.mjs b/scripts/embed-memory-v2-openai-compat.test.mjs new file mode 100644 index 0000000..c81f61c --- /dev/null +++ b/scripts/embed-memory-v2-openai-compat.test.mjs @@ -0,0 +1,116 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { + __resetEmbeddingCacheForTests, + embedQuery, + embedText, + probeEmbeddingAvailability, +} from './embed-memory-v2-openai-compat.mjs'; + +function mockFetch(responseFactory) { + return async (url, init) => { + const call = { url: String(url), init }; + return responseFactory(call); + }; +} + +test('embedText uses OpenAI-compatible /embeddings and caches by text hash', async () => { + __resetEmbeddingCacheForTests(); + const calls = []; + const fetchImpl = mockFetch(({ url, init }) => { + calls.push({ url, body: JSON.parse(String(init.body)) }); + return { + ok: true, + async json() { + return { data: [{ embedding: [0.1, 0.2, 0.3] }] }; + }, + }; + }); + + const first = await embedText('Sarah closed the curtains', { + env: { + MEMIND_EMBEDDING_PROVIDER: 'openai', + MEMIND_EMBEDDING_API_KEY: 'test-key', + MEMIND_EMBEDDING_BASE_URL: 'https://example.com/v1', + MEMIND_EMBEDDING_MODEL: 'text-embedding-3-small', + MEMIND_EMBEDDING_CACHE_PATH: '', + }, + fetchImpl, + persist: false, + }); + const second = await embedText('Sarah closed the curtains', { + env: { + MEMIND_EMBEDDING_PROVIDER: 'openai', + MEMIND_EMBEDDING_API_KEY: 'test-key', + MEMIND_EMBEDDING_BASE_URL: 'https://example.com/v1', + MEMIND_EMBEDDING_MODEL: 'text-embedding-3-small', + MEMIND_EMBEDDING_CACHE_PATH: '', + }, + fetchImpl, + persist: false, + }); + + assert.deepEqual(first, [0.1, 0.2, 0.3]); + assert.deepEqual(second, first); + assert.equal(calls.length, 1); + assert.equal(calls[0].url, 'https://example.com/v1/embeddings'); + assert.equal(calls[0].body.model, 'text-embedding-3-small'); +}); + +test('embedQuery delegates to embedText', async () => { + __resetEmbeddingCacheForTests(); + const vector = await embedQuery('probe question', {}, { + env: { + MEMIND_EMBEDDING_PROVIDER: 'openai', + MEMIND_EMBEDDING_API_KEY: 'test-key', + MEMIND_EMBEDDING_BASE_URL: 'https://example.com/v1', + MEMIND_EMBEDDING_CACHE_PATH: '', + }, + fetchImpl: mockFetch(() => ({ + ok: true, + async json() { + return { data: [{ embedding: [0.5, 0.6] }] }; + }, + })), + persist: false, + }); + assert.deepEqual(vector, [0.5, 0.6]); +}); + +test('ollama provider uses /api/embed without API key', async () => { + __resetEmbeddingCacheForTests(); + let seenUrl = null; + const vector = await embedText('local semantic probe', { + env: { + MEMIND_EMBEDDING_PROVIDER: 'ollama', + MEMIND_EMBEDDING_BASE_URL: 'http://127.0.0.1:11434', + MEMIND_EMBEDDING_MODEL: 'nomic-embed-text', + MEMIND_EMBEDDING_CACHE_PATH: '', + }, + fetchImpl: async (url) => { + seenUrl = String(url); + return { + ok: true, + async json() { + return { embeddings: [[0.9, 0.1, 0.4]] }; + }, + }; + }, + persist: false, + }); + assert.equal(seenUrl, 'http://127.0.0.1:11434/api/embed'); + assert.deepEqual(vector, [0.9, 0.1, 0.4]); +}); + +test('probeEmbeddingAvailability reports missing API key for openai provider', async () => { + __resetEmbeddingCacheForTests(); + const probe = await probeEmbeddingAvailability({ + env: { + MEMIND_EMBEDDING_PROVIDER: 'openai', + MEMIND_EMBEDDING_CACHE_PATH: '', + }, + }); + assert.equal(probe.available, false); + assert.equal(probe.reason, 'embedding_api_key_missing'); +}); diff --git a/scripts/resolve-llm-embedding-credentials.mjs b/scripts/resolve-llm-embedding-credentials.mjs new file mode 100644 index 0000000..c6fcd6d --- /dev/null +++ b/scripts/resolve-llm-embedding-credentials.mjs @@ -0,0 +1,96 @@ +import { createDbPool } from '../db.mjs'; +import { decryptSecret } from '../llm-providers.mjs'; + +function isDashScopeProviderRow(row) { + if (!row) return false; + const providerId = String(row.provider_id ?? '').trim(); + const apiUrl = String(row.api_url ?? '').trim().toLowerCase(); + if (providerId === 'custom_qwen') return true; + return apiUrl.includes('dashscope.aliyuncs.com'); +} + +function normalizeBaseUrl(apiUrl) { + const raw = String(apiUrl ?? '').trim().replace(/\/$/, ''); + if (!raw) return 'https://dashscope.aliyuncs.com/compatible-mode/v1'; + if (raw.endsWith('/embeddings')) return raw.slice(0, -'/embeddings'.length); + return raw; +} + +/** + * Resolve DashScope / Qwen embedding credentials from memind_adm-managed + * `h5_llm_provider_keys`, the same store used by Providers in memind_adm. + */ +export async function resolveLlmEmbeddingCredentials({ + pool = null, + keyId = null, + env = process.env, +} = {}) { + const ownsPool = !pool; + const db = pool ?? createDbPool(); + try { + const explicitKeyId = String(keyId ?? env.MEMIND_EMBEDDING_LLM_KEY_ID ?? '').trim(); + let row = null; + if (explicitKeyId) { + const [rows] = await db.query( + 'SELECT * FROM h5_llm_provider_keys WHERE id = ? AND status = ? LIMIT 1', + [explicitKeyId, 'active'], + ); + row = rows[0] ?? null; + } else { + const [selectedRows] = await db.query( + 'SELECT * FROM h5_llm_provider_keys WHERE is_selected = 1 AND status = ? LIMIT 1', + ['active'], + ); + if (isDashScopeProviderRow(selectedRows[0])) { + row = selectedRows[0]; + } + if (!row) { + const [dashRows] = await db.query( + `SELECT * FROM h5_llm_provider_keys + WHERE status = ? + AND ( + provider_id = 'custom_qwen' + OR api_url LIKE '%dashscope.aliyuncs.com%' + ) + ORDER BY is_selected DESC, updated_at DESC + LIMIT 1`, + ['active'], + ); + row = dashRows[0] ?? null; + } + } + + if (!row) { + return { available: false, reason: 'dashscope_llm_key_not_found' }; + } + if (!isDashScopeProviderRow(row)) { + return { available: false, reason: 'llm_key_not_dashscope', providerId: row.provider_id }; + } + + const apiKey = decryptSecret( + { + ciphertext: row.api_key_ciphertext, + iv: row.api_key_iv, + tag: row.api_key_tag, + }, + env.H5_SETTINGS_ENCRYPTION_KEY ?? env.TKMIND_SERVER__SECRET_KEY, + ); + if (!apiKey) { + return { available: false, reason: 'dashscope_api_key_decrypt_failed', keyName: row.name }; + } + + return { + available: true, + provider: 'dashscope', + keyId: row.id, + keyName: row.name, + providerId: row.provider_id, + apiKey, + baseUrl: normalizeBaseUrl(row.api_url), + chatModel: row.default_model, + model: String(env.MEMIND_EMBEDDING_MODEL ?? 'text-embedding-v3').trim(), + }; + } finally { + if (ownsPool) await db.end(); + } +} diff --git a/scripts/resolve-llm-embedding-credentials.test.mjs b/scripts/resolve-llm-embedding-credentials.test.mjs new file mode 100644 index 0000000..ea83df1 --- /dev/null +++ b/scripts/resolve-llm-embedding-credentials.test.mjs @@ -0,0 +1,80 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { encryptSecret } from '../llm-providers.mjs'; +import { resolveLlmEmbeddingCredentials } from './resolve-llm-embedding-credentials.mjs'; + +function makePool(rowsByQuery) { + return { + async query(sql, params = []) { + const text = String(sql); + if (text.includes('WHERE id = ?')) { + const row = rowsByQuery.byId?.get(params[0]); + return [row ? [row] : []]; + } + if (text.includes('is_selected = 1')) { + return [rowsByQuery.selected ?? []]; + } + if (text.includes('dashscope.aliyuncs.com')) { + return [rowsByQuery.dashscope ?? []]; + } + return [[]]; + }, + async end() {}, + }; +} + +test('resolveLlmEmbeddingCredentials decrypts selected DashScope key', async () => { + const secret = encryptSecret('dashscope-test-key', 'unit-test-secret'); + const row = { + id: 'key-1', + name: 'qwen3-max', + provider_id: 'custom_qwen3-max_1', + api_url: 'https://dashscope.aliyuncs.com/compatible-mode/v1', + default_model: 'qwen3-max', + api_key_ciphertext: secret.ciphertext, + api_key_iv: secret.iv, + api_key_tag: secret.tag, + }; + const resolved = await resolveLlmEmbeddingCredentials({ + pool: makePool({ selected: [row] }), + env: { TKMIND_SERVER__SECRET_KEY: 'unit-test-secret' }, + }); + assert.equal(resolved.available, true); + assert.equal(resolved.apiKey, 'dashscope-test-key'); + assert.equal(resolved.baseUrl, 'https://dashscope.aliyuncs.com/compatible-mode/v1'); + assert.equal(resolved.keyName, 'qwen3-max'); +}); + +test('resolveLlmEmbeddingCredentials falls back when selected key is not DashScope', async () => { + const secret = encryptSecret('dashscope-fallback-key', 'unit-test-secret'); + const dashRow = { + id: 'ds-1', + name: 'qwen chat', + provider_id: 'custom_qwen_chat', + api_url: 'https://dashscope.aliyuncs.com/compatible-mode/v1', + default_model: 'qwen-plus', + api_key_ciphertext: secret.ciphertext, + api_key_iv: secret.iv, + api_key_tag: secret.tag, + }; + const resolved = await resolveLlmEmbeddingCredentials({ + pool: makePool({ + selected: [{ + id: 'deepseek', + name: 'DeepSeek', + provider_id: 'custom_deepseek', + api_url: null, + default_model: 'deepseek-v4-pro', + api_key_ciphertext: secret.ciphertext, + api_key_iv: secret.iv, + api_key_tag: secret.tag, + }], + dashscope: [dashRow], + }), + env: { TKMIND_SERVER__SECRET_KEY: 'unit-test-secret' }, + }); + assert.equal(resolved.available, true); + assert.equal(resolved.apiKey, 'dashscope-fallback-key'); + assert.equal(resolved.keyName, 'qwen chat'); +}); diff --git a/scripts/run-memory-v2-memfuse-bench.mjs b/scripts/run-memory-v2-memfuse-bench.mjs index 665d6e1..c12e09a 100644 --- a/scripts/run-memory-v2-memfuse-bench.mjs +++ b/scripts/run-memory-v2-memfuse-bench.mjs @@ -25,6 +25,7 @@ import { summarizeMemFuseBench, } from '../memory-v2-memfuse-bench.mjs'; import { resolveEmbeddingModuleSpecifier } from '../memory-v2-recall-benchmark.mjs'; +import { loadMemindEnvFiles } from './memind-runtime-profile.mjs'; function parseArgs(argv) { const options = { @@ -102,6 +103,12 @@ async function loadEmbedder(specifier) { if (typeof embed !== 'function') { throw new Error(`Embedding module must export embedText, embedQuery or default: ${specifier}`); } + if (typeof imported.prefetchEmbedTexts === 'function') { + embed.prefetchEmbedTexts = imported.prefetchEmbedTexts.bind(imported); + } + if (typeof imported.flushEmbeddingCache === 'function') { + embed.flushEmbeddingCache = imported.flushEmbeddingCache.bind(imported); + } return embed; } @@ -168,6 +175,7 @@ function printReport(report) { } async function main() { + loadMemindEnvFiles(process.cwd()); const options = parseArgs(process.argv.slice(2)); if (options.help) { usage(); @@ -195,6 +203,27 @@ async function main() { const embedText = options.embeddingModule ? await loadEmbedder(options.embeddingModule) : null; + const prefetchEmbedTexts = embedText?.prefetchEmbedTexts ?? null; + + if (options.embeddingModule) { + const resolved = resolveEmbeddingModuleSpecifier(options.embeddingModule); + const imported = await import(resolved); + if (typeof imported.probeEmbeddingAvailability === 'function') { + const probe = await imported.probeEmbeddingAvailability(); + if (!probe.available) { + process.stderr.write( + `Embedding probe failed (${probe.reason}).\n` + + 'Set MEMIND_EMBEDDING_API_KEY / OPENAI_API_KEY, or MEMIND_EMBEDDING_PROVIDER=ollama.\n', + ); + return options.strict ? 1 : 0; + } + if (!options.quiet) { + process.stderr.write( + `embedding probe ok: provider=${probe.provider} model=${probe.model} dims=${probe.dimensions}\n`, + ); + } + } + } let lastScenario = null; const report = await runMemFuseBench({ @@ -207,6 +236,7 @@ async function main() { candidateLimit: options.candidateLimit, includeSourceTags: options.sourceTags, embedText, + prefetchEmbedTexts, onProgress: options.quiet ? null : (event) => {