memind/memory-v2-pgvector.mjs

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;
}

const KEYWORD_STOP_TERMS = new Set([
  '记得', '忘记', '之前', '我们', '聊过', '讨论', '继续', '聊聊', '什么', '吗', '呢',
  '有没有', '是否', '告诉', '提到', '说过', '以前', '上次', '对话', '会话', '回忆',
  '搜索', '帮助', '可以', '一下', '还是', '然后', '现在', '今天', '晚上', '你好',
]);

export function extractKeywordTerms(query, { maxTerms = 8, minLength = 2 } = {}) {
  const normalized = String(query ?? '').normalize('NFKC').trim();
  if (!normalized) return [];
  const terms = new Set();
  const cjkOnly = normalized.replace(/[^\u4e00-\u9fff]/gu, '');
  for (let index = 0; index < cjkOnly.length; index += 1) {
    for (const size of [4, 3, 2]) {
      if (index + size > cjkOnly.length) continue;
      const term = cjkOnly.slice(index, index + size);
      if (term.length >= minLength && !KEYWORD_STOP_TERMS.has(term)) {
        terms.add(term);
      }
    }
  }
  for (const match of normalized.matchAll(/[a-z0-9]{3,}/gi)) {
    terms.add(match[0].toLowerCase());
  }
  return [...terms]
    .sort((left, right) => right.length - left.length)
    .slice(0, maxTerms);
}

function dedupeContentPrefix(text, length = 96) {
  return normalizeSearchText(String(text ?? '').slice(0, length));
}

const KEYWORD_FETCH_CAP = 500;

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 ?? [];
  return rows
    .map((row) => ({
      row,
      lexicalScore: lexicalQueryCoverage(query, row.content ?? ''),
      updatedAt: timestampValue(row.updated_at ?? row.created_at),
    }))
    .sort((left, right) => {
      if (left.lexicalScore !== right.lexicalScore) {
        return right.lexicalScore - left.lexicalScore;
      }
      return right.updatedAt - left.updatedAt;
    })
    .slice(0, safeLimit)
    .map(({ row, lexicalScore }) => ({ ...row, score: lexicalScore }));
}

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();
  const byPrefix = new Set();
  for (const row of rows ?? []) {
    const memory = normalizeRow(row);
    if (!memory) continue;
    const prefixKey = dedupeContentPrefix(memory.text);
    if (prefixKey && byPrefix.has(prefixKey)) continue;
    if (prefixKey) byPrefix.add(prefixKey);
    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, keywordRows] = await Promise.all([
        pool.query(sql, [userId, vectorLiteral(embedding), candidateLimit]),
        fetchKeywordCandidates(pool, {
          userId,
          query: input.query,
          tableName: resolvedTableName,
          limit: Math.max(limit, candidateLimit),
        }).catch(() => []),
      ]);
      const memories = rankHybridCandidates(
        [...(result?.rows ?? []), ...keywordRows],
        input.query,
        limit,
      );
      return {
        semanticMemories: memories.map((item) => item.text),
        memories,
      };
    },
  };
}

export const pgvectorMemoryBackendInternals = {
  lexicalQueryCoverage,
  rankHybridCandidates,
  extractKeywordTerms,
  dedupeContentPrefix,
};