john
124 lines
3.5 KiB
JavaScript
124 lines
3.5 KiB
JavaScript
#!/usr/bin/env node
|
|
import { runPgvectorMemorySmoke } from '../memory-v2-pgvector-smoke.mjs';
|
|
|
|
const DEFAULT_URL_ENV = 'MEMORY_PGVECTOR_DATABASE_URL';
|
|
|
|
function usage() {
|
|
return `
|
|
Usage:
|
|
node scripts/smoke-memory-v2-pgvector.mjs [options]
|
|
|
|
Runs a local pgvector read/write smoke test with synthetic rows.
|
|
|
|
Options:
|
|
--table <name> Table name. Default: memory_embeddings.
|
|
--dimensions <number> Embedding dimensions. Default: 3.
|
|
--create-schema Ensure schema before smoke. Default: false.
|
|
--create-extension Include CREATE EXTENSION IF NOT EXISTS vector when creating schema.
|
|
--keep-rows Do not delete synthetic smoke rows.
|
|
--url-env <name> Env var containing PostgreSQL URL. Default: ${DEFAULT_URL_ENV}.
|
|
-h, --help Show this help.
|
|
`.trim();
|
|
}
|
|
|
|
export function parseMemoryPgvectorSmokeArgs(argv = []) {
|
|
const options = {
|
|
tableName: 'memory_embeddings',
|
|
dimensions: 3,
|
|
createSchema: false,
|
|
createExtension: false,
|
|
cleanup: true,
|
|
urlEnv: DEFAULT_URL_ENV,
|
|
help: false,
|
|
};
|
|
|
|
for (let i = 0; i < argv.length; i += 1) {
|
|
const arg = argv[i];
|
|
switch (arg) {
|
|
case '--table':
|
|
i += 1;
|
|
options.tableName = argv[i];
|
|
break;
|
|
case '--dimensions':
|
|
i += 1;
|
|
options.dimensions = Number(argv[i]);
|
|
break;
|
|
case '--create-schema':
|
|
options.createSchema = true;
|
|
break;
|
|
case '--create-extension':
|
|
options.createExtension = true;
|
|
break;
|
|
case '--keep-rows':
|
|
options.cleanup = false;
|
|
break;
|
|
case '--url-env':
|
|
i += 1;
|
|
options.urlEnv = argv[i];
|
|
break;
|
|
case '-h':
|
|
case '--help':
|
|
options.help = true;
|
|
break;
|
|
default:
|
|
throw new Error(`Unknown argument: ${arg}`);
|
|
}
|
|
}
|
|
|
|
if (!options.tableName) throw new Error('--table requires a value');
|
|
if (!Number.isInteger(options.dimensions) || options.dimensions < 1) {
|
|
throw new Error('--dimensions requires a positive integer');
|
|
}
|
|
if (!options.urlEnv) throw new Error('--url-env requires a value');
|
|
return options;
|
|
}
|
|
|
|
export async function runMemoryPgvectorSmokeCli(
|
|
argv = process.argv.slice(2),
|
|
env = process.env,
|
|
{ importPg = (specifier) => import(specifier), stdout = process.stdout } = {},
|
|
) {
|
|
const options = parseMemoryPgvectorSmokeArgs(argv);
|
|
if (options.help) {
|
|
stdout.write(`${usage()}\n`);
|
|
return { ok: true, mode: 'help' };
|
|
}
|
|
|
|
const connectionString = env[options.urlEnv];
|
|
if (!connectionString) {
|
|
throw new Error(`pgvector smoke requires ${options.urlEnv} to be set`);
|
|
}
|
|
|
|
const imported = await importPg('pg');
|
|
const PgPool = imported?.Pool ?? imported?.default?.Pool;
|
|
if (typeof PgPool !== 'function') throw new Error('pg module does not export Pool');
|
|
|
|
const pool = new PgPool({ connectionString, max: 1 });
|
|
try {
|
|
const result = await runPgvectorMemorySmoke({
|
|
pool,
|
|
tableName: options.tableName,
|
|
dimensions: options.dimensions,
|
|
createSchema: options.createSchema,
|
|
createExtension: options.createExtension,
|
|
cleanup: options.cleanup,
|
|
});
|
|
stdout.write(`${JSON.stringify({
|
|
...result,
|
|
checkedAt: new Date().toISOString(),
|
|
}, null, 2)}\n`);
|
|
return result;
|
|
} finally {
|
|
await pool.end();
|
|
}
|
|
}
|
|
|
|
if (import.meta.url === `file://${process.argv[1]}`) {
|
|
runMemoryPgvectorSmokeCli().then((result) => {
|
|
process.exitCode = result.ok ? 0 : 1;
|
|
}).catch((err) => {
|
|
console.error(err instanceof Error ? err.message : err);
|
|
process.exitCode = 1;
|
|
});
|
|
}
|