rmi-backend/app/streaming_rag.py

"""
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