Momento/memento-note/app/api/admin/embeddings/migrate/route.ts

import { NextRequest, NextResponse } from 'next/server'
import { prisma } from '@/lib/prisma'
import { auth } from '@/auth'
import { embeddingService } from '@/lib/ai/services/embedding.service'
import { getSystemConfig } from '@/lib/config'

export async function POST(req: NextRequest) {
  try {
    const session = await auth()
    if (!session?.user?.id || (session.user as { role?: string }).role !== 'ADMIN') {
      return NextResponse.json({ error: 'Unauthorized' }, { status: 401 })
    }

    const body = await req.json().catch(() => ({}))
    const targetDimension = body.targetDimension as number | undefined

    const config = await getSystemConfig()

    const testResult = await embeddingService.generateEmbedding('dimension test')
    const modelDimension = testResult.dimension

    const colDimResult: Array<{ dim: number | null }> = await prisma.$queryRawUnsafe(
      `SELECT a.atttypmod AS dim FROM pg_attribute a JOIN pg_class c ON a.attrelid = c.oid WHERE c.relname = 'NoteEmbedding' AND a.attname = 'embedding'`
    )
    const dbDimension = colDimResult[0]?.dim ?? null

    const newDimension = targetDimension || modelDimension

    if (dbDimension === newDimension) {
      return NextResponse.json({
        success: true,
        message: 'Dimensions match, no migration needed',
        dbDimension,
        modelDimension,
      })
    }

    const existingCount: Array<{ total: bigint }> = await prisma.$queryRawUnsafe(
      `SELECT COUNT(*)::bigint AS total FROM "NoteEmbedding"`
    )
    const count = Number(existingCount[0]?.total ?? 0)

    await prisma.$executeRawUnsafe(
      `DROP INDEX IF EXISTS "NoteEmbedding_embedding_hnsw"`
    )

    await prisma.$executeRawUnsafe(
      `TRUNCATE TABLE "NoteEmbedding"`
    )

    await prisma.$executeRawUnsafe(
      `ALTER TABLE "NoteEmbedding" ALTER COLUMN "embedding" TYPE vector(${newDimension}) USING NULL`
    )

    await prisma.$executeRawUnsafe(`
      CREATE INDEX "NoteEmbedding_embedding_hnsw" ON "NoteEmbedding"
      USING hnsw ("embedding" vector_cosine_ops)
      WITH (m = 16, ef_construction = 64)
    `)

    const docChunkExists: Array<{ exists: boolean }> = await prisma.$queryRawUnsafe(
      `SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = 'DocumentChunk' AND column_name = 'embedding') AS exists`
    )

    if (docChunkExists[0]?.exists) {
      await prisma.$executeRawUnsafe(
        `DROP INDEX IF EXISTS "DocumentChunk_embedding_hnsw_idx"`
      )
      await prisma.$executeRawUnsafe(
        `TRUNCATE TABLE "DocumentChunk"`
      )
      await prisma.$executeRawUnsafe(
        `ALTER TABLE "DocumentChunk" ALTER COLUMN "embedding" TYPE vector(${newDimension}) USING NULL`
      )
      await prisma.$executeRawUnsafe(`
        CREATE INDEX "DocumentChunk_embedding_hnsw_idx" ON "DocumentChunk"
        USING hnsw ("embedding" vector_cosine_ops)
        WITH (m = 16, ef_construction = 64)
      `)
    }

    await prisma.systemConfig.upsert({
      where: { key: 'AI_EMBEDDING_DIMENSION' },
      update: { value: String(newDimension) },
      create: { key: 'AI_EMBEDDING_DIMENSION', value: String(newDimension) },
    })

    return NextResponse.json({
      success: true,
      message: `Migration complete: vector(${dbDimension}) → vector(${newDimension}). ${count} old embeddings cleared. Re-indexing needed.`,
      previousDimension: dbDimension,
      newDimension,
      clearedEmbeddings: count,
    })
  } catch (error) {
    console.error('[EMBEDDING_MIGRATE] Error:', error)
    return NextResponse.json({ success: false, error: String(error) }, { status: 500 })
  }
}