128 lines
3.7 KiB
JavaScript
128 lines
3.7 KiB
JavaScript
#!/usr/bin/env node
|
|
import {
|
|
buildPgvectorMemorySchemaSql,
|
|
ensurePgvectorMemorySchema,
|
|
} from '../memory-v2-pgvector-schema.mjs';
|
|
|
|
const DEFAULT_URL_ENV = 'MEMORY_PGVECTOR_DATABASE_URL';
|
|
|
|
function usage() {
|
|
return `
|
|
Usage:
|
|
node scripts/setup-memory-v2-pgvector-schema.mjs [options]
|
|
|
|
Default mode is dry-run: print SQL only.
|
|
|
|
Options:
|
|
--apply Execute SQL against PostgreSQL.
|
|
--table <name> Table name. Default: memory_embeddings.
|
|
--dimensions <number> Embedding dimensions. Default: 1536.
|
|
--create-extension Include CREATE EXTENSION IF NOT EXISTS vector.
|
|
--create-vector-index Include ivfflat vector index creation.
|
|
--url-env <name> Env var containing PostgreSQL URL. Default: ${DEFAULT_URL_ENV}.
|
|
-h, --help Show this help.
|
|
`.trim();
|
|
}
|
|
|
|
export function parseMemoryPgvectorSchemaArgs(argv = []) {
|
|
const options = {
|
|
apply: false,
|
|
tableName: 'memory_embeddings',
|
|
dimensions: 1536,
|
|
createExtension: false,
|
|
createVectorIndex: false,
|
|
urlEnv: DEFAULT_URL_ENV,
|
|
help: false,
|
|
};
|
|
|
|
for (let i = 0; i < argv.length; i += 1) {
|
|
const arg = argv[i];
|
|
switch (arg) {
|
|
case '--apply':
|
|
options.apply = true;
|
|
break;
|
|
case '--create-extension':
|
|
options.createExtension = true;
|
|
break;
|
|
case '--create-vector-index':
|
|
options.createVectorIndex = true;
|
|
break;
|
|
case '--table':
|
|
i += 1;
|
|
options.tableName = argv[i];
|
|
break;
|
|
case '--dimensions':
|
|
i += 1;
|
|
options.dimensions = Number(argv[i]);
|
|
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.isFinite(options.dimensions)) throw new Error('--dimensions requires a number');
|
|
if (!options.urlEnv) throw new Error('--url-env requires a value');
|
|
return options;
|
|
}
|
|
|
|
function printSql(statements) {
|
|
for (const statement of statements) {
|
|
console.log(`${statement};\n`);
|
|
}
|
|
}
|
|
|
|
export async function runMemoryPgvectorSchemaCli(argv = process.argv.slice(2), env = process.env) {
|
|
const options = parseMemoryPgvectorSchemaArgs(argv);
|
|
if (options.help) {
|
|
console.log(usage());
|
|
return { ok: true, mode: 'help' };
|
|
}
|
|
|
|
const schemaOptions = {
|
|
tableName: options.tableName,
|
|
dimensions: options.dimensions,
|
|
createExtension: options.createExtension,
|
|
createVectorIndex: options.createVectorIndex,
|
|
};
|
|
const statements = buildPgvectorMemorySchemaSql(schemaOptions);
|
|
|
|
if (!options.apply) {
|
|
console.log('-- Memory V2 pgvector schema dry-run');
|
|
console.log('-- No database changes were made. Re-run with --apply to execute.');
|
|
printSql(statements);
|
|
return { ok: true, mode: 'dry-run', statements: statements.length };
|
|
}
|
|
|
|
const connectionString = env[options.urlEnv];
|
|
if (!connectionString) {
|
|
throw new Error(`--apply requires ${options.urlEnv} to be set`);
|
|
}
|
|
const { default: pg } = await import('pg');
|
|
const pool = new pg.Pool({ connectionString, max: 1 });
|
|
try {
|
|
const result = await ensurePgvectorMemorySchema(pool, schemaOptions);
|
|
console.log(
|
|
`Memory V2 pgvector schema ensured: table=${result.tableName}, dimensions=${result.dimensions}, statements=${result.statements}`,
|
|
);
|
|
return { ok: true, mode: 'apply', ...result };
|
|
} finally {
|
|
await pool.end();
|
|
}
|
|
}
|
|
|
|
if (import.meta.url === `file://${process.argv[1]}`) {
|
|
runMemoryPgvectorSchemaCli().catch((err) => {
|
|
console.error(err instanceof Error ? err.message : err);
|
|
process.exitCode = 1;
|
|
});
|
|
}
|