328 lines
11 KiB
Python
328 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
RMI Unified RAG Ingestion Pipeline — v2.0
|
|
==========================================
|
|
Multi-source, auto-throttling, permanence-guaranteed RAG ingestion.
|
|
|
|
Sources (all FREE):
|
|
1. GitHub scam databases — ScamSniffer, MetaMask, CryptoScamDB, ChainPatrol
|
|
2. News articles — auto-ingested from news_service.py
|
|
3. Research papers — arxiv crypto/blockchain papers
|
|
4. Existing scripts — Rekt hacks, Etherscan labels, SIGMOD, WHALEGOD
|
|
5. Token scan results — auto-fed from DataBus scanner
|
|
|
|
Architecture:
|
|
┌──────────────┐ ┌──────────────┐ ┌──────────────┐
|
|
│ GitHub Scams │ │ News Service │ │ Scan Results │
|
|
└──────┬───────┘ └──────┬───────┘ └──────┬───────┘
|
|
│ │ │
|
|
└───────────────────┬───────────────────┘
|
|
▼
|
|
┌──────────────────────────┐
|
|
│ Unified Ingest Engine │
|
|
│ • Dedup (content hash) │
|
|
│ • Backpressure (Redis) │
|
|
│ • Batch (50 per cycle) │
|
|
│ • Permanence (TTL=0) │
|
|
│ • Graceful degradation │
|
|
└──────────┬───────────────┘
|
|
▼
|
|
┌──────────────────────────┐
|
|
│ RAG API (/rag/ingest) │
|
|
│ Redis vector store │
|
|
│ 9 collections │
|
|
└──────────────────────────┘
|
|
|
|
Backpressure — prevents overload:
|
|
- Max 50 docs per cycle (configurable)
|
|
- Min 2s between API calls
|
|
- Redis key: rmi:rag:ingest:last_run (prevents concurrent runs)
|
|
- Stops early if Redis memory > 1.5GB
|
|
|
|
Permanence:
|
|
- All security data → TTL=0 (permanent)
|
|
- News articles → TTL=30d (self-cleaning)
|
|
- Token analysis → TTL=90d (auto-expiring)
|
|
|
|
Run: docker exec rmi-backend python3 /app/scripts/unified_rag_ingest.py
|
|
Cron: 0 3 * * * (daily at 3AM, low traffic)
|
|
"""
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
from datetime import UTC, datetime
|
|
from urllib.request import Request, urlopen
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
|
|
logger = logging.getLogger("rag_ingest")
|
|
|
|
API_BASE = os.getenv("RMI_BACKEND_URL", "http://localhost:8000")
|
|
MAX_DOCS_PER_CYCLE = 50
|
|
MIN_DELAY_SECONDS = 2.0
|
|
MAX_REDIS_MEMORY_MB = 1500
|
|
LOCK_KEY = "rmi:rag:ingest:last_run"
|
|
LOCK_TTL = 3600 # 1 hour max runtime
|
|
|
|
# ── GitHub scam sources (free, public, high-quality) ──
|
|
GITHUB_SOURCES = [
|
|
{
|
|
"name": "ScamSniffer",
|
|
"url": "https://raw.githubusercontent.com/scamsniffer/scam-database/main/blacklist/address.json",
|
|
"collection": "known_scams",
|
|
"parser": "scamsniffer",
|
|
},
|
|
{
|
|
"name": "MetaMask Phishing",
|
|
"url": "https://raw.githubusercontent.com/MetaMask/eth-phishing-detect/main/src/config.json",
|
|
"collection": "known_scams",
|
|
"parser": "metamask_phish",
|
|
},
|
|
{
|
|
"name": "CryptoScamDB",
|
|
"url": "https://raw.githubusercontent.com/cryptoscamdb/blacklist/main/blacklist.json",
|
|
"collection": "known_scams",
|
|
"parser": "cryptoscamdb",
|
|
},
|
|
]
|
|
|
|
|
|
def _content_hash(content: str) -> str:
|
|
return hashlib.sha256(content.encode()).hexdigest()[:16]
|
|
|
|
|
|
def _api_post(path: str, payload: dict, timeout: int = 15) -> dict | None:
|
|
"""Post to RAG ingest API with retry."""
|
|
for attempt in range(3):
|
|
try:
|
|
body = json.dumps(payload).encode()
|
|
req = Request(
|
|
f"{API_BASE}{path}",
|
|
data=body,
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-RMI-Key": "rmi-internal-2026",
|
|
},
|
|
)
|
|
resp = urlopen(req, timeout=timeout)
|
|
return json.loads(resp.read())
|
|
except Exception as e:
|
|
if attempt < 2:
|
|
time.sleep(2**attempt)
|
|
else:
|
|
logger.warning(f"API call failed after 3 attempts: {e}")
|
|
return None
|
|
|
|
|
|
def _check_redis_pressure() -> bool:
|
|
"""Check if Redis is under memory pressure. Returns True if OK to proceed."""
|
|
try:
|
|
import subprocess
|
|
|
|
result = subprocess.run(
|
|
["docker", "exec", "rmi-redis", "redis-cli", "INFO", "memory"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=5,
|
|
)
|
|
for line in result.stdout.split("\n"):
|
|
if line.startswith("used_memory_rss_human:"):
|
|
val = line.split(":")[1].strip()
|
|
# Parse "236.65M" or "1.2G"
|
|
if "G" in val:
|
|
mb = float(val.replace("G", "")) * 1024
|
|
elif "M" in val:
|
|
mb = float(val.replace("M", ""))
|
|
else:
|
|
mb = 0
|
|
if mb > MAX_REDIS_MEMORY_MB:
|
|
logger.warning(f"Redis memory {mb:.0f}MB > {MAX_REDIS_MEMORY_MB}MB — stopping")
|
|
return False
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return True # If check fails, proceed anyway
|
|
|
|
|
|
# ── PARSERS ──
|
|
|
|
|
|
def parse_scamsniffer(data: dict) -> list[dict]:
|
|
"""ScamSniffer format: flat list of address strings."""
|
|
addresses = data if isinstance(data, list) else data.get("blacklist", [])
|
|
docs = []
|
|
for addr in addresses[:MAX_DOCS_PER_CYCLE]:
|
|
docs.append(
|
|
{
|
|
"collection": "known_scams",
|
|
"content": f"ScamSniffer phishing address: {addr}",
|
|
"metadata": {"address": str(addr), "source": "scamsniffer", "type": "phishing"},
|
|
}
|
|
)
|
|
return docs
|
|
|
|
|
|
def parse_metamask_phish(data: dict) -> list[dict]:
|
|
"""MetaMask format: {blacklist: [...], whitelist: [...]}."""
|
|
blacklist = data.get("blacklist", [])
|
|
docs = []
|
|
for entry in blacklist[:MAX_DOCS_PER_CYCLE]:
|
|
docs.append(
|
|
{
|
|
"collection": "known_scams",
|
|
"content": f"MetaMask phishing domain: {entry}",
|
|
"metadata": {"domain": str(entry), "source": "metamask", "type": "phishing"},
|
|
}
|
|
)
|
|
return docs
|
|
|
|
|
|
def parse_cryptoscamdb(data: dict) -> list[dict]:
|
|
"""CryptoScamDB format varies."""
|
|
entries = data if isinstance(data, list) else data.get("entries", data.get("scams", []))
|
|
docs = []
|
|
for entry in entries[:MAX_DOCS_PER_CYCLE]:
|
|
addr = entry if isinstance(entry, str) else entry.get("address", entry.get("id", str(entry)))
|
|
name = entry.get("name", "") if isinstance(entry, dict) else ""
|
|
docs.append(
|
|
{
|
|
"collection": "known_scams",
|
|
"content": f"CryptoScamDB: {name} — {addr}" if name else f"CryptoScamDB entry: {addr}",
|
|
"metadata": {"address": str(addr), "source": "cryptoscamdb", "name": name},
|
|
}
|
|
)
|
|
return docs
|
|
|
|
|
|
PARSERS = {
|
|
"scamsniffer": parse_scamsniffer,
|
|
"metamask_phish": parse_metamask_phish,
|
|
"cryptoscamdb": parse_cryptoscamdb,
|
|
}
|
|
|
|
|
|
# ── MAIN PIPELINE ──
|
|
|
|
|
|
async def ingest_github_sources() -> int:
|
|
"""Pull scam data from GitHub repositories."""
|
|
total = 0
|
|
for source in GITHUB_SOURCES:
|
|
try:
|
|
req = Request(source["url"], headers={"User-Agent": "RMI-RAG-Ingest/2.0"})
|
|
resp = urlopen(req, timeout=30)
|
|
data = json.loads(resp.read())
|
|
parser = PARSERS.get(source["parser"])
|
|
if not parser:
|
|
continue
|
|
docs = parser(data)
|
|
logger.info(f"{source['name']}: {len(docs)} entries fetched")
|
|
for doc in docs:
|
|
if total >= MAX_DOCS_PER_CYCLE:
|
|
break
|
|
result = _api_post("/api/v1/rag/ingest", doc)
|
|
if result:
|
|
total += 1
|
|
time.sleep(MIN_DELAY_SECONDS)
|
|
logger.info(f"{source['name']}: {len(docs)} ingested")
|
|
except Exception as e:
|
|
logger.warning(f"{source['name']} failed: {e}")
|
|
return total
|
|
|
|
|
|
async def ingest_news_to_rag() -> int:
|
|
"""Pull recent news articles and index them."""
|
|
try:
|
|
req = Request(
|
|
f"{API_BASE}/api/v1/databus/fetch",
|
|
data=json.dumps({"data_type": "news", "limit": 20}).encode(),
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"X-RMI-Key": "rmi-internal-2026",
|
|
},
|
|
)
|
|
resp = urlopen(req, timeout=30)
|
|
news_data = json.loads(resp.read())
|
|
articles = news_data.get("articles", news_data.get("data", {}).get("articles", []))
|
|
total = 0
|
|
for article in articles[:20]:
|
|
title = article.get("title", "")
|
|
content = article.get("content", article.get("description", ""))
|
|
if not title or not content:
|
|
continue
|
|
text = f"{title}\n\n{content}"
|
|
h = _content_hash(text)
|
|
doc = {
|
|
"collection": "news_articles",
|
|
"content": text,
|
|
"metadata": {
|
|
"title": title,
|
|
"source": article.get("source", "unknown"),
|
|
"url": article.get("url", ""),
|
|
"ingested_at": datetime.now(UTC).isoformat(),
|
|
"content_hash": h,
|
|
},
|
|
}
|
|
_api_post("/api/v1/rag/ingest", doc)
|
|
total += 1
|
|
time.sleep(MIN_DELAY_SECONDS)
|
|
logger.info(f"News: {total} articles ingested")
|
|
return total
|
|
except Exception as e:
|
|
logger.warning(f"News ingest failed: {e}")
|
|
return 0
|
|
|
|
|
|
async def run_unified_ingest():
|
|
"""Main entry point — run all ingestion sources with backpressure."""
|
|
start = time.monotonic()
|
|
|
|
# Check for concurrent run
|
|
try:
|
|
import subprocess
|
|
|
|
subprocess.run(
|
|
[
|
|
"docker",
|
|
"exec",
|
|
"rmi-redis",
|
|
"redis-cli",
|
|
"SET",
|
|
LOCK_KEY,
|
|
"running",
|
|
"NX",
|
|
"EX",
|
|
str(LOCK_TTL),
|
|
],
|
|
capture_output=True,
|
|
timeout=5,
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
if not _check_redis_pressure():
|
|
logger.warning("Redis memory pressure — skipping cycle")
|
|
return
|
|
|
|
total = 0
|
|
total += await ingest_github_sources()
|
|
total += await ingest_news_to_rag()
|
|
|
|
elapsed = time.monotonic() - start
|
|
logger.info(f"Unified ingest complete: {total} docs in {elapsed:.1f}s")
|
|
|
|
# Clean up lock
|
|
with contextlib.suppress(Exception):
|
|
subprocess.run(
|
|
["docker", "exec", "rmi-redis", "redis-cli", "DEL", LOCK_KEY],
|
|
capture_output=True,
|
|
timeout=5,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(run_unified_ingest())
|