mcp/graph_utils.py

"""
graph_metrics.py  ·  Lightweight NetworkX helpers for MedGenesis

Key features
────────────
• Accepts edge dictionaries in either Streamlit-agraph or PyVis style:
      {"source": "n1", "target": "n2"}   ← agraph
      {"from":   "n1", "to":     "n2"}   ← PyVis
• Silently skips malformed edges (no KeyError).
• Provides three public helpers:
      build_nx(nodes, edges)         → networkx.Graph
      get_top_hubs(G, k=5)           → List[(node_id, degree_centrality)]
      get_density(G)                 → float  (0–1)
"""

from __future__ import annotations
from typing import List, Dict, Tuple
import networkx as nx


# ────────────────────────────────────────────────────────────────────
# Internal helpers
# ────────────────────────────────────────────────────────────────────
def _edge_ends(e: Dict) -> Tuple[str, str] | None:
    """Return (src, dst) tuple if both ends exist; else None."""
    src = e.get("source") or e.get("from")
    dst = e.get("target") or e.get("to")
    if src and dst:
        return src, dst
    return None


# ────────────────────────────────────────────────────────────────────
# Public API
# ────────────────────────────────────────────────────────────────────
def build_nx(nodes: List[Dict], edges: List[Dict]) -> nx.Graph:
    """
    Convert agraph / PyVis node+edge dicts into a NetworkX Graph.

    Nodes: must contain "id"   (a unique string)
    Edges: accepted shapes → {"source":, "target":} or {"from":, "to":}
    """
    G = nx.Graph()

    # Add nodes with label attribute (used by Metrics tab)
    for n in nodes:
        G.add_node(n["id"], label=n.get("label", n["id"]))

    # Add edges (skip malformed)
    for e in edges:
        ends = _edge_ends(e)
        if ends:
            G.add_edge(*ends)

    return G


def get_top_hubs(G: nx.Graph, k: int = 5) -> List[Tuple[str, float]]:
    """
    Return top-k nodes by degree-centrality.
    Example output: [('TP53', 0.42), ('EGFR', 0.36), ...]
    """
    dc = nx.degree_centrality(G)
    return sorted(dc.items(), key=lambda x: x[1], reverse=True)[:k]


def get_density(G: nx.Graph) -> float:
    """Graph density in [0, 1]."""
    return nx.density(G)