memind/scripts/backfill-memory-v2-pgvector.mjs

#!/usr/bin/env node
import path from 'node:path';
import { pathToFileURL } from 'node:url';
import { backfillLegacyMemoriesToPgvector } from '../memory-v2-pgvector-backfill.mjs';

const DEFAULT_MYSQL_URL_ENV = 'MEMORY_BACKFILL_MYSQL_URL';
const DEFAULT_PG_URL_ENV = 'MEMORY_PGVECTOR_DATABASE_URL';

function usage() {
  return `
Usage:
  node scripts/backfill-memory-v2-pgvector.mjs [options]

Default mode is dry-run: scan MySQL source rows, but do not embed or write PostgreSQL.

Options:
  --apply                    Write embeddings into pgvector.
  --limit <number>           Batch size. Default: 100.
  --cursor-updated-at <ms>   Checkpoint updated_at value. Default: 0.
  --cursor-id <id>           Checkpoint id value. Default: empty.
  --table <name>             pgvector table. Default: memory_embeddings.
  --mysql-url-env <name>     Env var containing MySQL URL. Default: ${DEFAULT_MYSQL_URL_ENV}.
  --pg-url-env <name>        Env var containing PostgreSQL URL. Default: ${DEFAULT_PG_URL_ENV}.
  --embedding-module <path>  Required with --apply. Must export embedMemory(memory) or default.
  -h, --help                 Show this help.
`.trim();
}

export function parseMemoryPgvectorBackfillArgs(argv = []) {
  const options = {
    apply: false,
    limit: 100,
    cursor: { updatedAt: 0, id: '' },
    tableName: 'memory_embeddings',
    mysqlUrlEnv: DEFAULT_MYSQL_URL_ENV,
    pgUrlEnv: DEFAULT_PG_URL_ENV,
    embeddingModule: null,
    help: false,
  };

  for (let i = 0; i < argv.length; i += 1) {
    const arg = argv[i];
    switch (arg) {
      case '--apply':
        options.apply = true;
        break;
      case '--limit':
        i += 1;
        options.limit = Number(argv[i]);
        break;
      case '--cursor-updated-at':
        i += 1;
        options.cursor.updatedAt = Number(argv[i]);
        break;
      case '--cursor-id':
        i += 1;
        options.cursor.id = String(argv[i] ?? '');
        break;
      case '--table':
        i += 1;
        options.tableName = argv[i];
        break;
      case '--mysql-url-env':
        i += 1;
        options.mysqlUrlEnv = argv[i];
        break;
      case '--pg-url-env':
        i += 1;
        options.pgUrlEnv = argv[i];
        break;
      case '--embedding-module':
        i += 1;
        options.embeddingModule = argv[i];
        break;
      case '-h':
      case '--help':
        options.help = true;
        break;
      default:
        throw new Error(`Unknown argument: ${arg}`);
    }
  }

  if (!Number.isFinite(options.limit)) throw new Error('--limit requires a number');
  if (!Number.isFinite(options.cursor.updatedAt)) {
    throw new Error('--cursor-updated-at requires a number');
  }
  if (!options.tableName) throw new Error('--table requires a value');
  if (!options.mysqlUrlEnv) throw new Error('--mysql-url-env requires a value');
  if (!options.pgUrlEnv) throw new Error('--pg-url-env requires a value');
  return options;
}

async function defaultCreateMysqlPool(connectionString) {
  const mysql = await import('mysql2/promise');
  return mysql.createPool(connectionString);
}

async function defaultCreatePgPool(connectionString) {
  const { default: pg } = await import('pg');
  return new pg.Pool({ connectionString, max: 1 });
}

async function defaultLoadEmbedMemory(modulePath) {
  if (!modulePath) return null;
  const resolved = path.isAbsolute(modulePath)
    ? modulePath
    : path.resolve(process.cwd(), modulePath);
  const mod = await import(pathToFileURL(resolved).href);
  const embedMemory = mod.embedMemory ?? mod.default;
  if (typeof embedMemory !== 'function') {
    throw new Error('--embedding-module must export embedMemory(memory) or default function');
  }
  return embedMemory;
}

export async function runMemoryPgvectorBackfillCli(
  argv = process.argv.slice(2),
  env = process.env,
  deps = {},
) {
  const options = parseMemoryPgvectorBackfillArgs(argv);
  if (options.help) {
    console.log(usage());
    return { ok: true, mode: 'help' };
  }

  const mysqlUrl = env[options.mysqlUrlEnv];
  if (!mysqlUrl) {
    throw new Error(`Backfill requires ${options.mysqlUrlEnv} to be set`);
  }
  if (options.apply && !env[options.pgUrlEnv]) {
    throw new Error(`--apply requires ${options.pgUrlEnv} to be set`);
  }
  if (options.apply && !options.embeddingModule) {
    throw new Error('--apply requires --embedding-module');
  }

  const createMysqlPool = deps.createMysqlPool ?? defaultCreateMysqlPool;
  const createPgPool = deps.createPgPool ?? defaultCreatePgPool;
  const loadEmbedMemory = deps.loadEmbedMemory ?? defaultLoadEmbedMemory;

  const mysqlPool = await createMysqlPool(mysqlUrl);
  let pgPool = null;
  try {
    let embedMemory = null;
    if (options.apply) {
      pgPool = await createPgPool(env[options.pgUrlEnv]);
      embedMemory = await loadEmbedMemory(options.embeddingModule);
    }
    const result = await backfillLegacyMemoriesToPgvector({
      mysqlPool,
      pgPool,
      embedMemory,
      tableName: options.tableName,
      cursor: options.cursor,
      limit: options.limit,
      dryRun: !options.apply,
    });
    console.log(JSON.stringify(result, null, 2));
    return result;
  } finally {
    await mysqlPool?.end?.();
    await pgPool?.end?.();
  }
}

if (import.meta.url === `file://${process.argv[1]}`) {
  runMemoryPgvectorBackfillCli().catch((err) => {
    console.error(err instanceof Error ? err.message : err);
    process.exitCode = 1;
  });
}