john
79 lines
2.5 KiB
JavaScript
79 lines
2.5 KiB
JavaScript
const DEFAULT_TABLE = 'memory_embeddings';
|
|
const DEFAULT_DIMENSIONS = 1536;
|
|
|
|
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 resolveDimensions(value) {
|
|
const dimensions = Number(value ?? DEFAULT_DIMENSIONS);
|
|
if (!Number.isInteger(dimensions) || dimensions < 1 || dimensions > 16_384) {
|
|
throw new Error(`Invalid pgvector dimensions: ${value}`);
|
|
}
|
|
return dimensions;
|
|
}
|
|
|
|
export function buildPgvectorMemorySchemaSql({
|
|
tableName = DEFAULT_TABLE,
|
|
dimensions = DEFAULT_DIMENSIONS,
|
|
createExtension = false,
|
|
createVectorIndex = false,
|
|
} = {}) {
|
|
const table = quoteIdent(tableName);
|
|
const dim = resolveDimensions(dimensions);
|
|
const statements = [];
|
|
if (createExtension) {
|
|
statements.push('CREATE EXTENSION IF NOT EXISTS vector');
|
|
}
|
|
statements.push(`
|
|
CREATE TABLE IF NOT EXISTS ${table} (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
user_id TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
embedding vector(${dim}) NOT NULL,
|
|
type TEXT NOT NULL DEFAULT 'fact',
|
|
source_memory_id TEXT,
|
|
source_session_id TEXT,
|
|
source_message_id TEXT,
|
|
metadata JSONB NOT NULL DEFAULT '{}'::jsonb,
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
)`.trim());
|
|
statements.push(
|
|
`CREATE INDEX IF NOT EXISTS ${quoteIdent(`${tableName}_user_created_idx`)} ON ${table} (user_id, created_at DESC)`,
|
|
);
|
|
statements.push(
|
|
`CREATE UNIQUE INDEX IF NOT EXISTS ${quoteIdent(`${tableName}_source_memory_id_uidx`)} ON ${table} (source_memory_id) WHERE source_memory_id IS NOT NULL`,
|
|
);
|
|
if (createVectorIndex) {
|
|
statements.push(
|
|
`CREATE INDEX IF NOT EXISTS ${quoteIdent(`${tableName}_embedding_idx`)} ON ${table} USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100)`,
|
|
);
|
|
}
|
|
return statements;
|
|
}
|
|
|
|
export async function ensurePgvectorMemorySchema(pool, options = {}) {
|
|
if (!pool?.query) {
|
|
throw new Error('ensurePgvectorMemorySchema requires a PostgreSQL pool with query(sql)');
|
|
}
|
|
const statements = buildPgvectorMemorySchemaSql(options);
|
|
for (const statement of statements) {
|
|
await pool.query(statement);
|
|
}
|
|
return {
|
|
ok: true,
|
|
statements: statements.length,
|
|
tableName: options.tableName ?? DEFAULT_TABLE,
|
|
dimensions: resolveDimensions(options.dimensions),
|
|
};
|
|
}
|