Update mcp/graph_metrics.py

mcp/graph_metrics.py

@@ -1,63 +1,97 @@
 # mcp/graph_metrics.py
 """
-Tiny NetworkX helpers with forgiving edge parsing.
-
-Why this exists
----------------
-Your Streamlit-agraph edges carry keys ``source/target`` while other
-toolchains (PyVis, cytoscape, etc.) use ``from/to``.  A naïve
-``e["target"]`` lookup therefore raises **KeyError**.  The functions below
-accept either spelling and silently skip malformed edges.
+Tiny NetworkX helpers for MedGenesis graphs.
+
+✨ 2025-06-25 REVAMP
+────────────────────
+• Accepts edge-dicts that use either
+      {'source': 'n1', 'target': 'n2'}   ← agraph / d3.js
+   or {'from'  : 'n1', 'to'    : 'n2'}   ← PyVis
+
+• Silently skips malformed edges (no more KeyError).
+
+• Works whether you pass plain dicts or streamlit-agraph Node/Edge objects.
 """
 
 from __future__ import annotations
-from typing import List, Dict, Tuple
 
 import networkx as nx
+from typing import List, Dict, Tuple, Union
 
 
-# ── private -----------------------------------------------------------------
-def _endpoints(edge: Dict) -> Tuple[str, str] | None:
-    """Return (src,dst) tuple or *None* if either end is missing."""
-    src = edge.get("source") or edge.get("from")
-    dst = edge.get("target") or edge.get("to")
-    return (src, dst) if src and dst else None
+# ── helpers -----------------------------------------------------------------
+def _edge_ends(e: Dict) -> Tuple[str, str] | None:
+    """
+    Normalise edge formats.
 
+    Returns
+    -------
+    (src, dst) tuple  – or None if either endpoint is missing.
+    """
+    src = e.get("source") or e.get("from")
+    dst = e.get("target") or e.get("to")
+    if src and dst:
+        return str(src), str(dst)
+    return None
 
-# ── public helpers ----------------------------------------------------------
-def build_nx(nodes: List[Dict], edges: List[Dict]) -> nx.Graph:
+
+def _node_id(n: Union[Dict, object]) -> str:
+    """
+    Accept either a dict *or* a streamlit_agraph.Node and return its id.
     """
-    Convert generic node/edge dictionaries into a NetworkX ``Graph``.
+    if isinstance(n, dict):
+        return str(n.get("id"))
+    # fallback for Node dataclass
+    return str(getattr(n, "id", ""))
 
-    •   Nodes are added with an attribute ``label`` (fallback: id).  
-    •   Edges are added only when both src & dst exist.  
-    •   **No KeyError** even if edge-schema is mixed.
+
+def _node_label(n: Union[Dict, object]) -> str:
+    """
+    Extract label safely from dict or Node.
+    """
+    if isinstance(n, dict):
+        return str(n.get("label", n.get("id")))
+    return str(getattr(n, "label", getattr(n, "id", "")))
+
+
+# ── public API --------------------------------------------------------------
+def build_nx(
+    nodes: List[Dict | object],
+    edges: List[Dict | object],
+) -> nx.Graph:
+    """
+    Convert generic node/edge payloads into a NetworkX graph.
     """
     G = nx.Graph()
 
-    # nodes
+    # add nodes
     for n in nodes:
-        try:
-            G.add_node(n["id"], label=n.get("label", n["id"]))
-        except KeyError:
-            # ignore malformed node dicts
+        nid = _node_id(n)
+        if not nid:
             continue
+        G.add_node(nid, label=_node_label(n))
 
-    # edges
+    # add edges
     for e in edges:
-        ends = _endpoints(e)
+        if not isinstance(e, dict):
+            e = e.__dict__  # Edge dataclass → dict
+        ends = _edge_ends(e)
         if ends:
             G.add_edge(*ends)
 
     return G
 
 
-def get_top_hubs(G: nx.Graph, k: int = 5):
-    """Return *k* nodes with highest degree-centrality."""
+def get_top_hubs(G: nx.Graph, k: int = 5) -> List[Tuple[str, float]]:
+    """
+    Return top-k nodes by **degree centrality**.
+    """
     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 ∈ [0,1]."""
+    """
+    Graph density in [0, 1].
+    """
     return nx.density(G)