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>
This commit is contained in:
@@ -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=<uuid> 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
|
||||
|
||||
+42
-43
@@ -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();
|
||||
|
||||
+383
-51
@@ -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,
|
||||
};
|
||||
|
||||
+130
-2
@@ -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' }),
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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';
|
||||
@@ -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;
|
||||
@@ -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');
|
||||
});
|
||||
@@ -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();
|
||||
}
|
||||
}
|
||||
@@ -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');
|
||||
});
|
||||
@@ -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) => {
|
||||
|
||||
Reference in New Issue
Block a user