Files
memind/memory-v2-pgvector.mjs
T
john d52aab8ab0
Memind CI / Test, build, and release guards (push) Successful in 3m52s
fix(memory): add bounded hybrid recall
2026-07-22 11:15:00 +08:00

188 lines
5.8 KiB
JavaScript

const DEFAULT_TABLE = 'memory_embeddings';
const DEFAULT_LIMIT = 8;
function isSafeIdentifier(value) {
return /^[a-zA-Z_][a-zA-Z0-9_]*$/.test(String(value ?? ''));
}
function resolveTableName(tableName) {
const normalized = String(tableName ?? DEFAULT_TABLE).trim();
if (!isSafeIdentifier(normalized)) {
throw new Error(`Invalid pgvector table name: ${normalized}`);
}
return normalized;
}
function normalizeEmbedding(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 vectorLiteral(embedding) {
return `[${embedding.join(',')}]`;
}
function normalizeSearchText(value) {
return String(value ?? '')
.normalize('NFKC')
.toLowerCase()
.replace(/[^a-z0-9\u4e00-\u9fff]+/gu, '');
}
function buildCharacterNgrams(value, size = 2) {
const text = normalizeSearchText(value);
if (!text) return new Set();
if (text.length <= size) return new Set([text]);
const grams = new Set();
for (let index = 0; index <= text.length - size; index += 1) {
grams.add(text.slice(index, index + size));
}
return grams;
}
function lexicalQueryCoverage(query, text) {
const queryGrams = buildCharacterNgrams(query);
if (queryGrams.size === 0) return 0;
const textGrams = buildCharacterNgrams(text);
let overlap = 0;
for (const gram of queryGrams) {
if (textGrams.has(gram)) overlap += 1;
}
return overlap / queryGrams.size;
}
function timestampValue(value) {
if (value == null) return 0;
const numeric = Number(value);
if (Number.isFinite(numeric)) return numeric;
const parsed = Date.parse(String(value));
return Number.isFinite(parsed) ? parsed : 0;
}
function normalizeRow(row) {
const text = String(row?.content ?? row?.memory_text ?? row?.text ?? '').trim();
if (!text) return null;
return {
id: row?.id == null ? null : String(row.id),
label: row?.type ?? row?.label ?? 'semantic',
text,
score: row?.score == null ? null : Number(row.score),
createdAt: row?.created_at ?? row?.createdAt ?? null,
updatedAt: row?.updated_at ?? row?.updatedAt ?? row?.created_at ?? row?.createdAt ?? null,
};
}
function rankHybridCandidates(rows, query, limit) {
const byId = new Map();
for (const row of rows ?? []) {
const memory = normalizeRow(row);
if (!memory) continue;
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;
}
return right.vectorScore - left.vectorScore;
})
.slice(0, limit)
.map(({ memory }) => memory);
}
export function createPgvectorMemoryBackend({
pool = null,
enabled = false,
tableName = DEFAULT_TABLE,
embedQuery = null,
defaultLimit = DEFAULT_LIMIT,
unavailableReason = 'not_configured',
} = {}) {
const resolvedTableName = resolveTableName(tableName);
async function resolveEmbedding(input) {
const explicit = normalizeEmbedding(input?.embedding);
if (explicit) return explicit;
if (typeof embedQuery !== 'function' || !input?.query) return null;
return normalizeEmbedding(await embedQuery(input.query, input));
}
return {
name: 'pgvector',
category: 'semantic',
role: 'primary-vector-store',
flag: 'MEMORY_VECTOR_ENABLED',
unavailableReason,
isAvailable() {
return Boolean(enabled && pool?.query);
},
getUnavailableReason() {
return this.isAvailable() ? null : unavailableReason;
},
async resolve(input = {}) {
if (!this.isAvailable()) return { memories: [], semanticMemories: [] };
const userId = String(input.userId ?? '').trim();
if (!userId) return { memories: [], semanticMemories: [] };
const embedding = await resolveEmbedding(input);
if (!embedding) return { memories: [], semanticMemories: [] };
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),
);
const sql = `
WITH vector_candidates AS (
SELECT id, content, type, created_at, updated_at,
1 - (embedding <=> $2::vector) AS score,
0 AS source_priority
FROM ${resolvedTableName}
WHERE user_id = $1
ORDER BY embedding <=> $2::vector
LIMIT $3
), recent_candidates AS (
SELECT id, content, type, created_at, updated_at,
1 - (embedding <=> $2::vector) AS score,
1 AS source_priority
FROM ${resolvedTableName}
WHERE user_id = $1
ORDER BY updated_at DESC
LIMIT $3
)
SELECT DISTINCT ON (id) id, content, type, created_at, updated_at, score
FROM (
SELECT * FROM vector_candidates
UNION ALL
SELECT * FROM recent_candidates
) AS candidates
ORDER BY id, source_priority
`;
const result = await pool.query(sql, [userId, vectorLiteral(embedding), candidateLimit]);
const memories = rankHybridCandidates(result?.rows ?? [], input.query, limit);
return {
semanticMemories: memories.map((item) => item.text),
memories,
};
},
};
}
export const pgvectorMemoryBackendInternals = {
lexicalQueryCoverage,
rankHybridCandidates,
};