merge: chore/cleanup-remove-bloat-and-secrets into main
This commit is contained in:
commit
bde2f3a97d
1173 changed files with 437609 additions and 0 deletions
281
scripts/migrate_redis_to_pgvector.py
Normal file
281
scripts/migrate_redis_to_pgvector.py
Normal file
|
|
@ -0,0 +1,281 @@
|
|||
#!/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())
|
||||
Loading…
Add table
Add a link
Reference in a new issue