242 lines
7.8 KiB
Python
242 lines
7.8 KiB
Python
"""
|
|
Streaming RAG — Real-time token/address monitoring with push updates.
|
|
|
|
"Watch this token" → results stream as the scanner finds new evidence.
|
|
Uses Redis pub/sub for real-time delivery to WebSocket or polling clients.
|
|
|
|
Architecture:
|
|
1. Client POSTs /rag/watch with token address + filters
|
|
2. Server subscribes to scanner output channel
|
|
3. As new scans complete, results are pushed via:
|
|
- WebSocket (existing /ws/ endpoint)
|
|
- Redis pub/sub (for external clients)
|
|
4. Client polls /rag/watch/{id}/updates for new results
|
|
"""
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import os
|
|
import uuid
|
|
from datetime import UTC, datetime
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# ── In-memory watch registry ─────────────────────────────────────
|
|
# In production, this would use Redis for persistence across restarts
|
|
_watches: dict[str, dict[str, Any]] = {}
|
|
_updates: dict[str, list[dict[str, Any]]] = {}
|
|
|
|
|
|
async def create_watch(
|
|
token_address: str,
|
|
chain: str = "ethereum",
|
|
filters: list[str] | None = None,
|
|
ttl_seconds: int = 3600,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Create a new watch on a token/address.
|
|
|
|
Returns watch_id for polling or WebSocket subscription.
|
|
"""
|
|
watch_id = str(uuid.uuid4())[:12]
|
|
|
|
watch = {
|
|
"id": watch_id,
|
|
"token_address": token_address,
|
|
"chain": chain,
|
|
"filters": filters or ["scam_detection", "deployer_analysis", "liquidity_check"],
|
|
"created_at": datetime.now(UTC).isoformat(),
|
|
"expires_at": (datetime.now(UTC).timestamp() + ttl_seconds),
|
|
"status": "active",
|
|
"total_updates": 0,
|
|
}
|
|
|
|
_watches[watch_id] = watch
|
|
_updates[watch_id] = []
|
|
|
|
logger.info(f"Created watch {watch_id} for {token_address[:12]}... on {chain}")
|
|
|
|
# Kick off async scanning
|
|
asyncio.create_task(_run_watch_scans(watch_id))
|
|
|
|
return watch
|
|
|
|
|
|
async def get_watch_updates(watch_id: str, since_index: int = 0) -> dict[str, Any]:
|
|
"""Get new updates for a watch since the given index."""
|
|
watch = _watches.get(watch_id)
|
|
if not watch:
|
|
return {"error": "Watch not found", "watch_id": watch_id}
|
|
|
|
all_updates = _updates.get(watch_id, [])
|
|
new_updates = all_updates[since_index:]
|
|
|
|
return {
|
|
"watch_id": watch_id,
|
|
"status": watch["status"],
|
|
"new_updates": new_updates,
|
|
"total_updates": len(all_updates),
|
|
"next_index": len(all_updates),
|
|
}
|
|
|
|
|
|
async def _run_watch_scans(watch_id: str):
|
|
"""Run scanner checks for a watch and push results."""
|
|
watch = _watches.get(watch_id)
|
|
if not watch:
|
|
return
|
|
|
|
token_address = watch["token_address"]
|
|
chain = watch["chain"]
|
|
watch.get("filters", [])
|
|
|
|
try:
|
|
# 1. Run unified scanner if available
|
|
try:
|
|
from app.unified_scanner import scan_wallet
|
|
|
|
result = await scan_wallet(token_address, chain=chain)
|
|
_push_update(
|
|
watch_id,
|
|
{
|
|
"type": "scanner",
|
|
"source": "unified_scanner",
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
"risk_score": result.get("risk_score", 0),
|
|
"risk_level": result.get("risk_level", "unknown"),
|
|
"issues": result.get("issues", []),
|
|
"warnings": result.get("warnings", []),
|
|
},
|
|
)
|
|
except ImportError:
|
|
logger.debug("Unified scanner not available")
|
|
except Exception as e:
|
|
logger.debug(f"Scanner failed: {e}")
|
|
|
|
# 2. RAG search for known scam patterns
|
|
try:
|
|
from app.rag_service import three_pillar_search
|
|
|
|
rag_results = await three_pillar_search(
|
|
query=f"{token_address} scam detection",
|
|
collections=["known_scams", "scam_patterns", "forensic_reports"],
|
|
limit=5,
|
|
use_mmr=True,
|
|
)
|
|
if rag_results.get("results"):
|
|
_push_update(
|
|
watch_id,
|
|
{
|
|
"type": "rag",
|
|
"source": "three_pillar_search",
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
"matches": len(rag_results["results"]),
|
|
"top_match": rag_results["results"][0].get("content", "")[:200],
|
|
"confidence": rag_results.get("confidence"),
|
|
},
|
|
)
|
|
except ImportError:
|
|
logger.debug("RAG search not available for streaming")
|
|
except Exception as e:
|
|
logger.debug(f"RAG search failed in watch: {e}")
|
|
|
|
# 3. Deployer reputation
|
|
try:
|
|
from app.deployer_reputation import score_deployer_risk
|
|
|
|
reputation = score_deployer_risk(
|
|
address=token_address,
|
|
deployment_count=0,
|
|
)
|
|
_push_update(
|
|
watch_id,
|
|
{
|
|
"type": "reputation",
|
|
"source": "deployer_reputation",
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
"risk_score": reputation.get("risk_score", 0),
|
|
"risk_label": reputation.get("risk_label", "unknown"),
|
|
"summary": reputation.get("summary", ""),
|
|
},
|
|
)
|
|
except ImportError:
|
|
logger.debug("Deployer reputation not available")
|
|
except Exception as e:
|
|
logger.debug(f"Reputation failed: {e}")
|
|
|
|
finally:
|
|
watch["status"] = "completed"
|
|
watch["completed_at"] = datetime.now(UTC).isoformat()
|
|
|
|
# Publish to Redis for WebSocket clients
|
|
with contextlib.suppress(Exception):
|
|
await _publish_to_redis(watch_id, watch)
|
|
|
|
|
|
def _push_update(watch_id: str, update: dict[str, Any]):
|
|
"""Push an update to the in-memory queue."""
|
|
if watch_id not in _updates:
|
|
_updates[watch_id] = []
|
|
_updates[watch_id].append(update)
|
|
if watch_id in _watches:
|
|
_watches[watch_id]["total_updates"] = len(_updates[watch_id])
|
|
|
|
|
|
async def _publish_to_redis(watch_id: str, watch: dict[str, Any]):
|
|
"""Publish watch completion to Redis for WebSocket delivery."""
|
|
try:
|
|
import redis.asyncio as aioredis
|
|
|
|
redis_host = os.getenv("REDIS_HOST", "rmi-redis")
|
|
redis_pass = os.getenv("REDIS_PASSWORD", "")
|
|
r = await aioredis.from_url(
|
|
f"redis://{redis_host}:{os.getenv('REDIS_PORT', '6379')}",
|
|
password=redis_pass or None,
|
|
)
|
|
await r.publish(
|
|
"rmi:ws:alerts",
|
|
json.dumps(
|
|
{
|
|
"type": "watch_complete",
|
|
"watch_id": watch_id,
|
|
"token_address": watch.get("token_address", ""),
|
|
"total_updates": watch.get("total_updates", 0),
|
|
"timestamp": datetime.now(UTC).isoformat(),
|
|
}
|
|
),
|
|
)
|
|
await r.close()
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
async def list_active_watches() -> list[dict[str, Any]]:
|
|
"""List all active watches."""
|
|
now = datetime.now(UTC).timestamp()
|
|
active = []
|
|
expired = []
|
|
|
|
for wid, watch in _watches.items():
|
|
if watch.get("expires_at", 0) > now:
|
|
active.append(
|
|
{
|
|
"id": wid,
|
|
"token_address": watch["token_address"],
|
|
"chain": watch["chain"],
|
|
"status": watch["status"],
|
|
"updates": watch.get("total_updates", 0),
|
|
"created_at": watch["created_at"],
|
|
}
|
|
)
|
|
else:
|
|
expired.append(wid)
|
|
|
|
# Clean up expired
|
|
for wid in expired:
|
|
_watches.pop(wid, None)
|
|
_updates.pop(wid, None)
|
|
|
|
return active
|