rmi-backend/scripts/migrate_redis_to_pgvector_v2.py

#!/usr/bin/env python3
"""
Resilient Redis → Supabase pgvector migration v2.
Fixes: smaller batches, retry on timeout, HNSW index dropped for speed.
"""

import asyncio
import json
import logging
import os
import time

import httpx
import redis.asyncio as aioredis
from dotenv import load_dotenv

load_dotenv("/app/.env", override=True)

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

TARGET_DIM = 640
BATCH_SIZE = 50  # Smaller batches to avoid Supabase statement timeout
MAX_RETRIES = 3

COLLECTIONS = [
    "wallet_profiles",  # ~50K already migrated
    "token_analysis",
    "known_scams",
    "scam_patterns",
    "forensic_reports",
    "market_intel",
    "contract_audits",
    "news_articles",
    "transaction_patterns",
    "general",
]


def _get_url():
    return os.environ.get("SUPABASE_URL", "")


def _get_key():
    return os.environ.get("SUPABASE_SERVICE_KEY", "") or os.environ.get(
        "SUPABASE_SERVICE_ROLE_KEY", ""
    )


def _get_headers():
    key = _get_key()
    return {
        "apikey": key,
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "Prefer": "resolution=merge-duplicates",
    }


def pad_vector(vec: list, target_dim: int) -> list:
    if len(vec) == target_dim:
        return vec
    if len(vec) > target_dim:
        return vec[:target_dim]
    return vec + [0.0] * (target_dim - len(vec))


async def get_existing_ids(client: httpx.AsyncClient) -> set:
    """Fetch existing doc IDs from rag_vectors for idempotency."""
    existing = set()
    offset = 0
    limit = 1000
    while True:
        url = f"{_get_url()}/rest/v1/rag_vectors"
        params = {"select": "id", "limit": str(limit), "offset": str(offset), "order": "id"}
        try:
            resp = await client.get(url, params=params, headers=_get_headers(), timeout=30)
            if resp.status_code != 200:
                logger.warning(f"Could not fetch existing IDs: {resp.status_code}")
                break
            rows = resp.json() if resp.text else []
            if not rows:
                break
            for row in rows:
                existing.add(row["id"])
            if len(rows) < limit:
                break
            offset += limit
        except Exception as e:
            logger.warning(f"Error fetching existing IDs: {e}")
            break
    logger.info(f"Found {len(existing)} existing docs in rag_vectors")
    return existing


async def insert_with_retry(client: httpx.AsyncClient, rows: list, batch_num: int) -> int:
    """Insert a batch with retry logic."""
    url = f"{_get_url()}/rest/v1/rag_vectors"
    params = {"on_conflict": "id"}

    for attempt in range(MAX_RETRIES):
        try:
            resp = await client.post(
                url, json=rows, params=params, headers=_get_headers(), timeout=60
            )
            if resp.status_code in (200, 201):
                return len(rows)
            elif resp.status_code == 500 and "timeout" in (resp.text or "").lower():
                # Statement timeout - split into smaller batches
                logger.warning(
                    f"Batch {batch_num}: statement timeout (attempt {attempt + 1}), splitting"
                )
                if len(rows) > 10:
                    # Split into two halves
                    mid = len(rows) // 2
                    count = 0
                    for half in [rows[:mid], rows[mid:]]:
                        r2 = await client.post(
                            url, json=half, params=params, headers=_get_headers(), timeout=90
                        )
                        if r2.status_code in (200, 201):
                            count += len(half)
                        else:
                            # Last resort: insert one by one
                            for row in half:
                                try:
                                    r3 = await client.post(
                                        url,
                                        json=[row],
                                        params=params,
                                        headers=_get_headers(),
                                        timeout=30,
                                    )
                                    if r3.status_code in (200, 201):
                                        count += 1
                                    await asyncio.sleep(0.02)
                                except Exception:
                                    logging.getLogger(__name__).warning(
                                        "swallowed exception", exc_info=True
                                    )
                    return count
            else:
                logger.error(
                    f"Batch {batch_num} insert failed: {resp.status_code} {resp.text[:300]}"
                )
                return 0
        except httpx.ReadTimeout:
            logger.warning(f"Batch {batch_num}: timeout (attempt {attempt + 1})")
            await asyncio.sleep(2**attempt)
        except Exception as e:
            logger.error(f"Batch {batch_num} exception: {e}")
            await asyncio.sleep(1)

    return 0


async def migrate_collection(
    r: aioredis.Redis, client: httpx.AsyncClient, collection: str, existing_ids: set
) -> int:
    """Migrate all docs from one Redis collection to Supabase."""
    idx_key = f"rag:idx:{collection}"
    doc_ids = await r.smembers(idx_key)
    if not doc_ids:
        logger.info(f"Collection {collection}: empty, skipping")
        return 0

    total_docs = len(doc_ids)
    to_migrate = [did for did in doc_ids if did not in existing_ids]
    skipped = total_docs - len(to_migrate)
    if skipped > 0:
        logger.info(
            f"Collection {collection}: {total_docs} total, skipping {skipped} already-migrated, {len(to_migrate)} to go"
        )

    migrated = 0
    errors = 0
    batch_rows = []
    batch_count = 0

    for _i, doc_id in enumerate(to_migrate):
        doc_key = f"rag:{collection}:{doc_id}"
        try:
            data = await r.get(doc_key)
            if not data:
                continue

            doc = json.loads(data)
            vector = doc.get("vector", [])
            content = doc.get("content", "") or ""
            metadata = doc.get("metadata", {}) or {}
            source = metadata.get("source", "") or doc.get("source", "") or ""
            severity = metadata.get("severity", "") or doc.get("severity", "") or "medium"
            chain = metadata.get("chain", "") or doc.get("chain", "") or ""

            padded = pad_vector(vector, TARGET_DIM) if vector else [0.0] * TARGET_DIM

            row = {
                "id": str(doc_id),
                "collection": collection,
                "content": content[:10000],
                "embedding": padded,
                "metadata": json.dumps(metadata) if isinstance(metadata, dict) else str(metadata),
                "source": source[:200] if source else "",
                "severity": severity[:50] if severity else "medium",
                "chain": chain[:50] if chain else "",
            }

            batch_rows.append(row)

            if len(batch_rows) >= BATCH_SIZE:
                count = await insert_with_retry(client, batch_rows, batch_count)
                migrated += count
                batch_count += 1
                if batch_count % 10 == 0:
                    logger.info(
                        f"  Progress: {migrated}/{len(to_migrate)} migrated ({batch_count} batches)"
                    )
                batch_rows = []
                await asyncio.sleep(0.1)

        except json.JSONDecodeError:
            errors += 1
        except Exception as e:
            errors += 1
            if errors <= 5:
                logger.warning(f"  Doc {doc_id}: {e}")

    # Flush remaining
    if batch_rows:
        count = await insert_with_retry(client, batch_rows, batch_count)
        migrated += count

    logger.info(f"Collection {collection}: DONE {migrated}/{len(to_migrate)} ({errors} errors)")
    return migrated


async def main():
    logger.info(f"Starting migration v2 (batch_size={BATCH_SIZE})")
    logger.info(f"Supabase: {_get_url()[:30]}...")
    logger.info(f"Target dim: {TARGET_DIM}")

    r = aioredis.Redis(
        host=os.environ.get("REDIS_HOST", "rmi-redis"),
        port=int(os.environ.get("REDIS_PORT", "6379")),
        password=os.environ.get("REDIS_PASSWORD", ""),
        decode_responses=True,
    )
    await r.ping()
    logger.info("Redis connected")

    async with httpx.AsyncClient(timeout=120) as client:
        existing_ids = await get_existing_ids(client)
        total = 0
        start = time.time()

        for coll in COLLECTIONS:
            try:
                n = await migrate_collection(r, client, coll, existing_ids)
                total += n
            except Exception as e:
                logger.error(f"Collection {coll} FAILED: {e}")

        elapsed = time.time() - start
        logger.info("=" * 60)
        logger.info(f"Migration v2 complete: {total} docs in {elapsed:.1f}s")
        logger.info(
            "Now rebuild HNSW index: CREATE INDEX idx_rag_vectors_hnsw ON rag_vectors USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);"
        )

    await r.aclose()


if __name__ == "__main__":
    asyncio.run(main())