exclaim-recovery-database/drizzle/vector-init.ts

Shared database schemas for Exclaim Recovery platform

import 'dotenv/config';
import { drizzle } from 'drizzle-orm/node-postgres';
import { Pool } from 'pg';
import { sql } from 'drizzle-orm';
import * as schema from './schema.js';

async function main() {
  // Get connection string from environment
  const connectionString = process.env.VERCEL_POSTGRES_URL;

  if (!connectionString) {
    console.error('No VERCEL_POSTGRES_URL found in environment');
    process.exit(1);
  }

  console.log('Connecting to database...');
  const pool = new Pool({
    connectionString,
  });

  const db = drizzle(pool, { schema });

  try {
    // First, ensure the pgvector extension is installed
    console.log('Installing pgvector extension...');
    await db.execute(sql`CREATE EXTENSION IF NOT EXISTS vector`);

    // Create array_to_vector function for easy conversion from JavaScript arrays
    console.log('Creating array_to_vector function...');
    await db.execute(sql`
      CREATE OR REPLACE FUNCTION array_to_vector(float8[])
      RETURNS vector
      LANGUAGE SQL
      IMMUTABLE
      RETURNS NULL ON NULL INPUT
      PARALLEL SAFE
      AS 'SELECT $1::vector';
    `);

    // Create HNSW indexes for all vector columns
    console.log('Creating HNSW indexes for vector columns...');

    // Initialize vectors with HNSW index for cosine similarity search
    await db.execute(sql`
      CREATE INDEX IF NOT EXISTS smart_recovery_vector_idx ON "SMARTRecoveryDocument" USING hnsw ("vector" vector_cosine_ops);
      CREATE INDEX IF NOT EXISTS scientific_study_vector_idx ON "ScientificStudyDocument" USING hnsw ("vector" vector_cosine_ops);
      CREATE INDEX IF NOT EXISTS document_vector_idx ON "Document" USING hnsw ("vector" vector_cosine_ops);
      CREATE INDEX IF NOT EXISTS user_document_vector_idx ON "UserDocument" USING hnsw ("vector" vector_cosine_ops);
      CREATE INDEX IF NOT EXISTS ga_step_vector_idx ON "GAStepDocument" USING hnsw ("vector" vector_cosine_ops);
      CREATE INDEX IF NOT EXISTS ga_resource_vector_idx ON "GAResourceDocument" USING hnsw ("vector" vector_cosine_ops);
      CREATE INDEX IF NOT EXISTS ga_reflection_vector_idx ON "GAReflectionDocument" USING hnsw ("vector" vector_cosine_ops);
    `);

    console.log('Vector initialization complete!');
  } catch (error) {
    console.error('Error initializing vector database:', error);
    process.exit(1);
  } finally {
    await pool.end();
  }
}

main();