memind/memory-v2-pgvector-smoke.mjs

import { ensurePgvectorMemorySchema } from './memory-v2-pgvector-schema.mjs';
import { createPgvectorMemoryBackend } from './memory-v2-pgvector.mjs';

const DEFAULT_TABLE = 'memory_embeddings';
const DEFAULT_DIMENSIONS = 3;
const SMOKE_SOURCE_ID = 'memory-v2-smoke-test';
const SMOKE_USER_ID = 'memory-v2-smoke-user';

function isSafeIdentifier(value) {
  return /^[a-zA-Z_][a-zA-Z0-9_]*$/.test(String(value ?? ''));
}

function quoteIdent(value) {
  const normalized = String(value ?? '').trim();
  if (!isSafeIdentifier(normalized)) {
    throw new Error(`Invalid PostgreSQL identifier: ${normalized}`);
  }
  return `"${normalized}"`;
}

function vectorLiteral(vector) {
  return `[${vector.map((item) => Number(item)).join(',')}]`;
}

function assertVector(vector, dimensions) {
  if (!Array.isArray(vector) || vector.length !== dimensions) {
    throw new Error(`Smoke vector must contain ${dimensions} dimensions`);
  }
  for (const item of vector) {
    if (!Number.isFinite(Number(item))) {
      throw new Error('Smoke vector contains a non-numeric value');
    }
  }
}

function resolveDimensions(value) {
  const dimensions = Number(value ?? DEFAULT_DIMENSIONS);
  if (!Number.isInteger(dimensions) || dimensions < 1 || dimensions > 16_384) {
    throw new Error(`Invalid pgvector smoke dimensions: ${value}`);
  }
  return dimensions;
}

export function buildPgvectorSmokeRows({ dimensions = DEFAULT_DIMENSIONS } = {}) {
  const resolvedDimensions = resolveDimensions(dimensions);
  const near = Array.from(
    { length: resolvedDimensions },
    (_, index) => (index === 0 ? 1 : 0),
  );
  const far = Array.from(
    { length: resolvedDimensions },
    (_, index) => (index === resolvedDimensions - 1 ? 1 : 0),
  );
  const query = [...near];
  return {
    userId: SMOKE_USER_ID,
    sourceMemoryPrefix: SMOKE_SOURCE_ID,
    query,
    expectedTopText: 'memory-v2 smoke near vector',
    rows: [
      {
        sourceMemoryId: `${SMOKE_SOURCE_ID}-near`,
        content: 'memory-v2 smoke near vector',
        embedding: near,
        type: 'smoke',
      },
      {
        sourceMemoryId: `${SMOKE_SOURCE_ID}-far`,
        content: 'memory-v2 smoke far vector',
        embedding: far,
        type: 'smoke',
      },
    ],
  };
}

export async function runPgvectorMemorySmoke({
  pool,
  tableName = DEFAULT_TABLE,
  dimensions = DEFAULT_DIMENSIONS,
  createSchema = false,
  createExtension = false,
  cleanup = true,
} = {}) {
  if (!pool?.query) {
    throw new Error('runPgvectorMemorySmoke requires a PostgreSQL pool with query(sql, params)');
  }
  const resolvedDimensions = resolveDimensions(dimensions);
  const table = quoteIdent(tableName);
  const fixture = buildPgvectorSmokeRows({ dimensions: resolvedDimensions });

  if (createSchema) {
    await ensurePgvectorMemorySchema(pool, {
      tableName,
      dimensions: resolvedDimensions,
      createExtension,
      createVectorIndex: false,
    });
  }

  try {
    for (const row of fixture.rows) {
      assertVector(row.embedding, resolvedDimensions);
      await pool.query(
        `INSERT INTO ${table}
          (user_id, content, embedding, type, source_memory_id, metadata)
         VALUES ($1, $2, $3::vector, $4, $5, $6::jsonb)
         ON CONFLICT (source_memory_id) WHERE source_memory_id IS NOT NULL
         DO UPDATE SET
          content = EXCLUDED.content,
          embedding = EXCLUDED.embedding,
          type = EXCLUDED.type,
          metadata = EXCLUDED.metadata,
          updated_at = NOW()`,
        [
          fixture.userId,
          row.content,
          vectorLiteral(row.embedding),
          row.type,
          row.sourceMemoryId,
          JSON.stringify({ smoke: true }),
        ],
      );
    }

    const backend = createPgvectorMemoryBackend({
      pool,
      enabled: true,
      tableName,
      defaultLimit: 2,
    });
    const resolved = await backend.resolve({
      userId: fixture.userId,
      embedding: fixture.query,
      limit: 2,
    });
    const top = resolved.memories[0] ?? null;
    const ok = top?.text === fixture.expectedTopText;
    return {
      ok,
      tableName,
      dimensions: resolvedDimensions,
      inserted: fixture.rows.length,
      expectedTopText: fixture.expectedTopText,
      topMemory: top,
      memories: resolved.memories,
    };
  } finally {
    if (cleanup) {
      await pool.query(
        `DELETE FROM ${table} WHERE source_memory_id LIKE $1`,
        [`${fixture.sourceMemoryPrefix}%`],
      );
    }
  }
}