scripts/backfill_768.py

"""
GovBridge India — Phase 2: Backfill 768-dim Embeddings
Sprint 17: Zero-Downtime Expand-Contract Migration

PREREQUISITES:
  1. Migration 007 has been executed (embedding_v2 column exists)
  2. Nomic model is available (pip install sentence-transformers einops)

USAGE:
  cd /workspaces/govbridge
  python3 gov_backend/scripts/backfill_768.py

BEHAVIOR:
  - Reads ALL rows from document_chunks where embedding_v2 IS NULL
  - Generates 768-dim embeddings using nomic-embed-text-v1
  - UPSERTs in batches of 50 rows using keyset pagination
  - Also backfills query_cache.query_embedding_v2
  - Safe to re-run (idempotent via IS NULL filter)
"""
import os
import sys
import time
import torch

# Ensure gov_backend is on the Python path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from config import settings

from supabase import create_client
from sentence_transformers import SentenceTransformer

# ── Configuration ────────────────────────────────────────────
BATCH_SIZE = 50          # Rows per batch (free-tier Supabase safe)
SLEEP_BETWEEN = 0.5      # Seconds between batches (prevent rate limit)
MAX_ROWS = 5000          # Safety ceiling per run

def main():
    supabase_url = settings.SUPABASE_URL
    supabase_key = settings.SUPABASE_KEY

    if not supabase_url or not supabase_key:
        print("❌ FATAL: Missing SUPABASE_URL or SUPABASE_KEY in environment")
        sys.exit(1)

    print("⏳ Loading Nomic Embedding Model (768-dim)...")
    model = SentenceTransformer(
        'nomic-ai/nomic-embed-text-v1',
        trust_remote_code=True
    )
    print("✅ Nomic model loaded")

    supabase = create_client(supabase_url, supabase_key)

    # ── Phase 2A: Backfill document_chunks.embedding_v2 ──────
    print("\n═══ PHASE 2A: Backfilling document_chunks.embedding_v2 ═══")
    
    result = supabase.table('document_chunks') \
        .select('id, chunk_text, scheme_title') \
        .is_('embedding_v2', 'null') \
        .limit(MAX_ROWS) \
        .execute()

    chunks = result.data or []
    total = len(chunks)
    
    if total == 0:
        print("🎉 document_chunks: All rows already have embedding_v2")
    else:
        print(f"Found {total} chunks to backfill")
        
        for i in range(0, total, BATCH_SIZE):
            batch = chunks[i:i + BATCH_SIZE]
            batch_num = (i // BATCH_SIZE) + 1
            total_batches = (total + BATCH_SIZE - 1) // BATCH_SIZE

            # Nomic requires "search_document:" prefix for document embeddings
            texts = [
                f"search_document: {c.get('chunk_text', '') or c.get('scheme_title', '')}"
                for c in batch
            ]

            print(f"🔄 Batch {batch_num}/{total_batches} ({len(batch)} rows)...")
            
            with torch.inference_mode():
                embeddings = model.encode(
                    texts,
                    normalize_embeddings=True,
                    show_progress_bar=False,
                    batch_size=32
                ).tolist()

            # UPSERT each row's embedding_v2
            for chunk, embedding in zip(batch, embeddings):
                supabase.table('document_chunks') \
                    .update({'embedding_v2': embedding}) \
                    .eq('id', chunk['id']) \
                    .execute()

            print(f"✅ Batch {batch_num}/{total_batches} complete")
            time.sleep(SLEEP_BETWEEN)

        print(f"\n🎉 document_chunks backfill complete: {total} rows updated")

    # ── Phase 2B: Backfill query_cache.query_embedding_v2 ────
    print("\n═══ PHASE 2B: Backfilling query_cache.query_embedding_v2 ═══")
    
    cache_result = supabase.table('query_cache') \
        .select('id, query_text') \
        .is_('query_embedding_v2', 'null') \
        .limit(MAX_ROWS) \
        .execute()

    cache_rows = cache_result.data or []
    cache_total = len(cache_rows)

    if cache_total == 0:
        print("🎉 query_cache: All rows already have query_embedding_v2")
    else:
        print(f"Found {cache_total} cache entries to backfill")
        
        for i in range(0, cache_total, BATCH_SIZE):
            batch = cache_rows[i:i + BATCH_SIZE]
            batch_num = (i // BATCH_SIZE) + 1

            # Nomic requires "search_query:" prefix for query embeddings
            texts = [
                f"search_query: {c.get('query_text', '')}"
                for c in batch
            ]

            with torch.inference_mode():
                embeddings = model.encode(
                    texts,
                    normalize_embeddings=True,
                    show_progress_bar=False,
                    batch_size=32
                ).tolist()

            for row, embedding in zip(batch, embeddings):
                supabase.table('query_cache') \
                    .update({'query_embedding_v2': embedding}) \
                    .eq('id', row['id']) \
                    .execute()

            time.sleep(SLEEP_BETWEEN)

        print(f"\n🎉 query_cache backfill complete: {cache_total} rows updated")

    # ── Summary ──────────────────────────────────────────────
    print("\n" + "═" * 60)
    print("MIGRATION PHASE 2 COMPLETE")
    print(f"  document_chunks: {total} rows backfilled")
    print(f"  query_cache: {cache_total} rows backfilled")
    print("═" * 60)
    print("\nNEXT STEPS:")
    print("  1. Run migration 008 (concurrent index creation)")
    print("  2. Verify indexes are VALID")
    print("  3. Run migration 009 (contract — column swap)")


if __name__ == '__main__':
    main()