john
94 lines
2.9 KiB
JavaScript
94 lines
2.9 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 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,
|
|
};
|
|
}
|
|
|
|
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 sql = `
|
|
SELECT id, content, type, created_at, 1 - (embedding <=> $2::vector) AS score
|
|
FROM ${resolvedTableName}
|
|
WHERE user_id = $1
|
|
ORDER BY embedding <=> $2::vector
|
|
LIMIT $3
|
|
`;
|
|
const result = await pool.query(sql, [userId, vectorLiteral(embedding), limit]);
|
|
const memories = (result?.rows ?? []).map((row) => normalizeRow(row)).filter(Boolean);
|
|
return {
|
|
semanticMemories: memories.map((item) => item.text),
|
|
memories,
|
|
};
|
|
},
|
|
};
|
|
}
|