rmi-backend/app/graph_rag.py

"""
Graph RAG - Community detection + narrative summaries from Knowledge Graph.

Extends the existing Redis-backed Knowledge Graph (5,614 nodes) with:
  1. Label propagation community detection - find scammer rings, exchange clusters
  2. Community narrative synthesis - LLM-generated summaries per cluster
  3. Graph-augmented retrieval - expand queries with community context

Produces structured community reports that answer:
  "What scammer rings are active right now?"
  "Show me the cluster this address belongs to"
"""

import json
import logging
from collections import defaultdict
from typing import Any

logger = logging.getLogger(__name__)


async def detect_communities(
    max_communities: int = 20,
    min_community_size: int = 3,
) -> list[dict[str, Any]]:
    """
    Detect communities in the Knowledge Graph using label propagation.

    Algorithm:
      1. Each node starts as its own community
      2. Iteratively: each node adopts the most common label among neighbors
      3. Converge after ~5 iterations
      4. Return top communities by size

    Returns list of community dicts with nodes, edges, and metadata.
    """
    try:
        from app.knowledge_graph import get_knowledge_graph

        kg = await get_knowledge_graph()
    except ImportError:
        logger.debug("Knowledge Graph not available")
        return []
    except Exception as e:
        logger.warning(f"KG access failed: {e}")
        return []

    r = kg._redis_client if hasattr(kg, "_redis_client") else None
    if not r:
        logger.debug("No Redis client in KG")
        return []

    # 1. Get all nodes
    try:
        node_ids = list(await r.smembers("kg:nodes"))
        if not node_ids:
            logger.debug("KG has no nodes")
            return []
    except Exception as e:
        logger.warning(f"Failed to get KG nodes: {e}")
        return []

    # 2. Get edges (adjacency)
    adjacency: dict[str, set[str]] = defaultdict(set)
    for nid in node_ids:
        try:
            edges_raw = await r.smembers(f"kg:edges:{nid}")
            for edge in edges_raw:
                if isinstance(edge, bytes):
                    edge = edge.decode()
                parts = edge.split(":") if ":" in edge else [edge, ""]
                target = parts[1] if len(parts) > 1 else parts[0]
                if target in node_ids or target in {b.decode() if isinstance(b, bytes) else b for b in node_ids}:
                    adjacency[nid].add(target)
        except Exception:
            pass

    # 3. Label propagation
    N_ITER = 5
    labels = {nid: i for i, nid in enumerate(node_ids)}

    for _ in range(N_ITER):
        new_labels = {}
        for nid in node_ids:
            neighbors = adjacency.get(nid, set())
            if not neighbors:
                new_labels[nid] = labels.get(nid, 0)
                continue

            # Most common label among neighbors
            label_counts = defaultdict(int)
            for neighbor in neighbors:
                n_label = labels.get(neighbor)
                if n_label is not None:
                    label_counts[n_label] += 1

            if label_counts:
                new_labels[nid] = max(label_counts, key=label_counts.get)
            else:
                new_labels[nid] = labels.get(nid, 0)
        labels = new_labels

    # 4. Group into communities
    communities: dict[int, list[str]] = defaultdict(list)
    for nid, label in labels.items():
        communities[label].append(nid)

    # 5. Sort by size, filter minimum
    sorted_communities = sorted(
        communities.items(),
        key=lambda x: len(x[1]),
        reverse=True,
    )

    result = []
    for comm_id, (label, members) in enumerate(sorted_communities[:max_communities]):  # noqa: B007
        if len(members) < min_community_size:
            break

        # Get node types for community characterization
        node_types = defaultdict(int)
        for nid in members:
            try:
                node_data = await r.get(f"kg:node:{nid}")
                if node_data:
                    if isinstance(node_data, bytes):
                        node_data = node_data.decode()
                    data = json.loads(node_data) if node_data.startswith("{") else {}
                    ntype = data.get("type", "unknown")
                    node_types[ntype] += 1
            except Exception:
                node_types["unknown"] += 1

        # Get internal edges
        internal_edges = 0
        for nid in members:
            for neighbor in adjacency.get(nid, set()):
                if neighbor in members:
                    internal_edges += 1

        result.append(
            {
                "community_id": comm_id,
                "size": len(members),
                "members": members[:20],  # Top 20 members
                "node_types": dict(node_types),
                "internal_edges": internal_edges // 2,  # Undirected
                "density": round((internal_edges // 2) / max(len(members) * (len(members) - 1) / 2, 1), 4),
                "dominant_type": max(node_types, key=node_types.get) if node_types else "unknown",
            }
        )

    return result


async def build_community_narrative(
    community: dict[str, Any],
    max_chars: int = 2000,
) -> str:
    """
    Generate a narrative summary of a detected community.

    Uses LLM to synthesize: what this cluster is, key addresses, risk profile.
    Falls back to heuristic summary if LLM unavailable.
    """
    community.get("members", [])
    node_types = community.get("node_types", {})
    dominant = community.get("dominant_type", "unknown")

    heuristic = (
        f"Community {community['community_id']}: {community['size']} nodes, "
        f"primarily {dominant} addresses. "
        f"Density: {community.get('density', 0):.3f}. "
        f"Member types: {json.dumps(dict(node_types))}"
    )

    # Try LLM synthesis
    try:
        from app.investigation_narratives import AI_BASE, AI_MODEL, LLM_API_KEY

        if not LLM_API_KEY:
            return heuristic

        prompt = f"""Summarize this crypto wallet community:

Size: {community["size"]} addresses
Primary type: {dominant}
Density: {community.get("density", 0):.3f}
Member types: {json.dumps(dict(node_types))}

What kind of cluster is this? (scammer ring, exchange wallets, MEV bot network, etc.)
What is the risk profile? Limit to 3 sentences."""

        import httpx

        async with httpx.AsyncClient(timeout=30) as client:
            resp = await client.post(
                AI_BASE,
                headers={
                    "Authorization": f"Bearer {LLM_API_KEY}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": AI_MODEL,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": 200,
                },
            )
            if resp.status_code == 200:
                return resp.json()["choices"][0]["message"]["content"]
    except Exception:
        pass

    return heuristic


async def graph_rag_search(
    query: str,
    max_communities: int = 10,
) -> dict[str, Any]:
    """
    Full Graph RAG pipeline: detect communities → generate narratives → return.

    Returns communities with narrative summaries suitable for display.
    """
    communities = await detect_communities(max_communities=max_communities)

    # Generate narratives for top communities
    enriched = []
    for comm in communities[:5]:
        narrative = await build_community_narrative(comm)
        enriched.append({**comm, "narrative": narrative})

    return {
        "total_communities": len(communities),
        "communities": enriched,
        "total_nodes_in_communities": sum(c["size"] for c in communities),
        "summary": (f"Detected {len(communities)} communities across {sum(c['size'] for c in communities)} addresses."),
    }