#!/usr/bin/env python3 """ Migrate RAG data from Redis to Supabase rag_vectors table. Reads all documents from 10 Redis collections and upserts into pgvector. Handles mixed embedding dimensions by padding to 640. Processes in batches of 100. Idempotent (skips already-migrated docs). """ import asyncio import json import logging import os import sys 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 dimension - max across all collections is 640 (token_analysis) TARGET_DIM = 640 COLLECTIONS = [ "wallet_profiles", "token_analysis", "known_scams", "scam_patterns", "forensic_reports", "market_intel", "contract_audits", "news_articles", "transaction_patterns", "general", ] BATCH_SIZE = 100 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 _get_redis_password(): return os.environ.get("REDIS_PASSWORD", "") def pad_vector(vec: list[float], target_dim: int) -> list[float]: """Pad or truncate embedding vector to target dimension.""" 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 all existing doc IDs from rag_vectors to skip already-migrated docs.""" 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} {resp.text[:300]}") 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_batch(client: httpx.AsyncClient, rows: list[dict]) -> int: """Insert a batch of rows into rag_vectors via Supabase REST API.""" url = f"{_get_url()}/rest/v1/rag_vectors" params = {"on_conflict": "id"} try: resp = await client.post(url, json=rows, params=params, headers=_get_headers(), timeout=120) if resp.status_code in (200, 201): return len(rows) else: logger.error(f"Batch insert failed: {resp.status_code} {resp.text[:500]}") # Fallback: try one by one success = 0 for row in rows: try: r2 = await client.post(url, json=[row], params=params, headers=_get_headers(), timeout=60) if r2.status_code in (200, 201): success += 1 else: if r2.status_code == 400 and "embedding" in (r2.text or ""): logger.warning(f"Skipping {row['id']} due to embedding error") else: logger.warning(f"Single insert failed for {row['id']}: {r2.status_code} {r2.text[:200]}") except Exception as e: logger.warning(f"Single insert error for {row['id']}: {e}") return success except Exception as e: logger.error(f"Batch insert exception: {e}") 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) logger.info(f"Collection {collection}: {total_docs} docs to process") # Filter out already-migrated docs 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" Skipping {skipped} already-migrated docs") migrated = 0 batch_rows = [] batch_count = 0 errors = 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: logger.debug(f" Doc {doc_id}: not found in Redis") continue doc = json.loads(data) # Extract fields 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 "" # Pad vector to target dimension if vector: padded = pad_vector(vector, TARGET_DIM) else: padded = [0.0] * TARGET_DIM logger.debug(f" Doc {doc_id}: no embedding, using zero vector") 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_batch(client, batch_rows) migrated += count batch_count += 1 logger.info( f" Batch {batch_count}: inserted {count}/{len(batch_rows)} ({migrated}/{len(to_migrate)} total)" ) batch_rows = [] await asyncio.sleep(0.05) except json.JSONDecodeError as e: errors += 1 if errors <= 5: logger.warning(f" Doc {doc_id}: JSON decode error: {e}") except Exception as e: errors += 1 if errors <= 5: logger.warning(f" Doc {doc_id}: error: {e}") # Progress every 500 docs if (i + 1) % 500 == 0: logger.info(f" Progress: processed {i + 1}/{len(to_migrate)}, migrated {migrated}, errors {errors}") # Flush remaining batch if batch_rows: count = await insert_batch(client, batch_rows) migrated += count batch_count += 1 logger.info(f" Final batch {batch_count}: inserted {count}/{len(batch_rows)}") logger.info( f"Collection {collection}: migrated {migrated}/{len(to_migrate)} docs ({skipped} skipped, {errors} errors)" ) return migrated async def main(): logger.info("Starting Redis -> Supabase rag_vectors migration") logger.info(f"Supabase URL: {_get_url()}") logger.info(f"Redis: {os.environ.get('REDIS_HOST', 'rmi-redis')}:{os.environ.get('REDIS_PORT', '6379')}") logger.info(f"Target vector dimension: {TARGET_DIM}") logger.info(f"Collections: {COLLECTIONS}") if not _get_url() or not _get_key(): logger.error("SUPABASE_URL and SUPABASE_SERVICE_KEY must be set") sys.exit(1) redis_password = _get_redis_password() redis_host = os.environ.get("REDIS_HOST", "rmi-redis") redis_port = int(os.environ.get("REDIS_PORT", "6379")) # Connect to Redis r = aioredis.Redis( host=redis_host, port=redis_port, password=redis_password, decode_responses=True, ) await r.ping() logger.info("Redis connected") # HTTP client for Supabase async with httpx.AsyncClient(timeout=120) as client: # Get existing IDs for idempotency existing_ids = await get_existing_ids(client) total_migrated = 0 start_time = time.time() for collection in COLLECTIONS: try: count = await migrate_collection(r, client, collection, existing_ids) total_migrated += count except Exception as e: logger.error(f"Collection {collection} failed: {e}") elapsed = time.time() - start_time logger.info("=" * 60) logger.info(f"Migration complete: {total_migrated} docs migrated in {elapsed:.1f}s") await r.aclose() if __name__ == "__main__": asyncio.run(main())