| |
| """Fast RAG rebuild — run all ingestion sources + build FAISS indexes. |
| Single-pass, single client, proper error logging.""" |
|
|
| import asyncio |
| import csv |
| import hashlib |
| import json |
| import logging |
| import os |
| import sys |
|
|
| import httpx |
|
|
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") |
| logger = logging.getLogger("rag_rebuild") |
|
|
| DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data") |
| RAG_API = "http://localhost:8000/api/v1/rag/ingest" |
|
|
| SCAM_KEYWORDS = [ |
| "phish", |
| "hack", |
| "exploit", |
| "scam", |
| "drain", |
| "attack", |
| "heist", |
| "malicious", |
| "fraud", |
| "rug", |
| "compromis", |
| "fake", |
| ] |
|
|
|
|
| def doc_id(collection, key): |
| return hashlib.sha256(f"{collection}:{key.lower()}".encode()).hexdigest()[:16] |
|
|
|
|
| async def run(): |
| total_ingested = 0 |
| async with httpx.AsyncClient(timeout=30, limits=httpx.Limits(max_connections=20)) as client: |
| |
| csv_path = os.path.join(DATA_DIR, "etherscan_phish_hack.csv") |
| if os.path.exists(csv_path): |
| with open(csv_path, newline="", encoding="utf-8") as f: |
| rows = list(csv.DictReader(f)) |
| logger.info(f"1. Etherscan phish/hack: {len(rows)} addresses") |
| ok = err = 0 |
| for i, row in enumerate(rows): |
| addr = row.get("address", "").strip() |
| if not addr: |
| continue |
| payload = { |
| "collection": "known_scams", |
| "content": f"Etherscan {row.get('label_type', 'scam')} label: {addr} on {row.get('chain', 'Ethereum')}. Tag: {row.get('name_tag', '')}. Balance: {row.get('balance', '0')}. Txns: {row.get('txn_count', '0')}.", |
| "metadata": { |
| "address": addr.lower(), |
| "label_type": row.get("label_type", "scam"), |
| "name_tag": row.get("name_tag", ""), |
| "chain": (row.get("chain", "Ethereum") or "Ethereum").lower(), |
| "source": "etherscan_phish_hack", |
| "severity": "critical", |
| }, |
| } |
| try: |
| resp = await client.post(RAG_API, json=payload) |
| if resp.status_code == 200: |
| ok += 1 |
| else: |
| err += 1 |
| if err <= 3: |
| logger.warning(f" Ingest error {resp.status_code}: {resp.text[:200]}") |
| except Exception as e: |
| err += 1 |
| if err <= 3: |
| logger.warning(f" Connection error: {e}") |
| if (i + 1) % 25 == 0: |
| logger.info(f" Phish/hack: {i + 1}/{len(rows)} | ok={ok} err={err}") |
| await asyncio.sleep(0.1) |
| total_ingested += ok |
| logger.info(f" Done: {ok} ingested, {err} errors") |
|
|
| |
| csv_path = os.path.join(DATA_DIR, "etherscan_combined_labels.csv") |
| if os.path.exists(csv_path): |
| with open(csv_path, newline="", encoding="utf-8") as f: |
| rows = list(csv.DictReader(f)) |
| scam_rows = [r for r in rows if any(kw in (r.get("label", "") or "").lower() for kw in SCAM_KEYWORDS)] |
| wallet_rows = [r for r in rows if not any(kw in (r.get("label", "") or "").lower() for kw in SCAM_KEYWORDS)] |
| logger.info(f"2. Etherscan combined: {len(rows)} total → {len(scam_rows)} scam, {len(wallet_rows)} wallets") |
|
|
| ok_scam = ok_wallet = err = 0 |
| |
| for i, row in enumerate(scam_rows): |
| addr = row.get("address", "").strip() |
| if not addr: |
| continue |
| payload = { |
| "collection": "known_scams", |
| "content": f"Etherscan scam label: {addr} on {row.get('chain', 'Ethereum')}. Label: {row.get('label', '')}. Category: {row.get('category', '')}.", |
| "metadata": { |
| "address": addr.lower(), |
| "label": row.get("label", ""), |
| "category": row.get("category", ""), |
| "chain": (row.get("chain", "Ethereum") or "Ethereum").lower(), |
| "source": "etherscan_combined", |
| "severity": "high", |
| }, |
| } |
| try: |
| resp = await client.post(RAG_API, json=payload) |
| if resp.status_code == 200: |
| ok_scam += 1 |
| else: |
| err += 1 |
| except: |
| err += 1 |
| if (i + 1) % 100 == 0: |
| logger.info(f" Scam: {i + 1}/{len(scam_rows)} | ok={ok_scam}") |
| await asyncio.sleep(0.1) |
|
|
| |
| batch_size = 50 |
| for i in range(0, min(len(wallet_rows), 5000), batch_size): |
| batch = wallet_rows[i : i + batch_size] |
| content_parts = [] |
| for row in batch: |
| content_parts.append( |
| f"{row.get('address', '')}: {row.get('label', '')} on {row.get('chain', 'Ethereum')}" |
| ) |
| payload = { |
| "collection": "wallet_profiles", |
| "content": "Etherscan wallet labels batch: " + " | ".join(content_parts), |
| "metadata": { |
| "source": "etherscan_combined", |
| "chain": "multi", |
| "count": len(batch), |
| "label_type": "wallet_profile", |
| }, |
| } |
| try: |
| resp = await client.post(RAG_API, json=payload) |
| if resp.status_code == 200: |
| ok_wallet += len(batch) |
| except: |
| pass |
| if (i // batch_size + 1) % 20 == 0: |
| logger.info(f" Wallets: {i + len(batch)}/{min(len(wallet_rows), 5000)} | ok={ok_wallet}") |
| await asyncio.sleep(0.1) |
|
|
| total_ingested += ok_scam + ok_wallet |
| logger.info(f" Done: {ok_scam} scams, {ok_wallet} wallets ingested") |
|
|
| |
| try: |
| resp = await client.get( |
| "https://raw.githubusercontent.com/solana-labs/token-list/main/src/tokens/solana.tokenlist.json" |
| ) |
| data = resp.json() |
| flagged = [ |
| t |
| for t in data.get("tokens", []) |
| if any(kw in str(t.get("tags", [])).lower() for kw in ["scam", "spam", "fake"]) |
| ] |
| logger.info(f"3. Solana tokens: {len(flagged)} flagged as scam/spam") |
| ok = 0 |
| for i, token in enumerate(flagged): |
| payload = { |
| "collection": "known_scams", |
| "content": f"Solana scam token: {token.get('name', '')} ({token.get('symbol', '')}) at {token.get('address', '')}. Tags: {token.get('tags', [])}.", |
| "metadata": { |
| "address": token.get("address", ""), |
| "name": token.get("name", ""), |
| "symbol": token.get("symbol", ""), |
| "chain": "solana", |
| "source": "solana_token_registry", |
| "tags": token.get("tags", []), |
| "severity": "high", |
| }, |
| } |
| try: |
| resp = await client.post(RAG_API, json=payload) |
| if resp.status_code == 200: |
| ok += 1 |
| except: |
| pass |
| await asyncio.sleep(0.15) |
| total_ingested += ok |
| logger.info(f" Done: {ok} ingested") |
| except Exception as e: |
| logger.warning(f" Solana token fetch failed: {e}") |
|
|
| |
| rekt_path = os.path.join(DATA_DIR, "rekt_hacks.json") |
| if os.path.exists(rekt_path): |
| with open(rekt_path) as f: |
| hacks = json.load(f) |
| logger.info(f"4. Rekt hacks: {len(hacks)} incidents") |
| ok = 0 |
| for i, hack in enumerate(hacks): |
| payload = { |
| "collection": "forensic_reports", |
| "content": f"DeFi hack: {hack.get('name', '')} — {hack.get('description', '')}. Attackers: {hack.get('attackers', [])}. Exploited: {hack.get('exploited', [])}. Amount: {hack.get('amount_lost', '')}.", |
| "metadata": { |
| "source": "rekt_database", |
| "name": hack.get("name", ""), |
| "attackers": hack.get("attackers", []), |
| "exploited": hack.get("exploited", []), |
| "severity": "critical", |
| "chain": "multi", |
| }, |
| } |
| try: |
| resp = await client.post(RAG_API, json=payload) |
| if resp.status_code == 200: |
| ok += 1 |
| except: |
| pass |
| await asyncio.sleep(0.1) |
| total_ingested += ok |
| logger.info(f" Done: {ok} ingested") |
| else: |
| logger.warning("4. Rekt hacks: file not found, skipping") |
|
|
| logger.info(f"\n=== TOTAL INGESTED: {total_ingested} ===") |
|
|
| |
| logger.info("\n5. Building FAISS indexes...") |
| async with httpx.AsyncClient(timeout=120) as client: |
| try: |
| resp = await client.post("http://localhost:8000/api/v1/rag/build-index", json={}) |
| logger.info(f" Build index response: {resp.status_code} — {resp.text[:300]}") |
|
|
| |
| resp2 = await client.get("http://localhost:8000/api/v1/rag/stats") |
| stats = resp2.json() |
| logger.info( |
| f" RAG stats: {stats.get('total_docs', 0)} total docs across {len(stats.get('collections', {}))} collections" |
| ) |
| for name, count in stats.get("collections", {}).items(): |
| logger.info(f" {name}: {count} docs") |
| except Exception as e: |
| logger.error(f" FAISS build error: {e}") |
|
|
|
|
| asyncio.run(run()) |
|
|
