Update mcp/graph_metrics.py

mcp/graph_metrics.py

@@ -1,31 +1,60 @@
-# mcp/graph_metrics.py
 """
-Basic graph-analytics helpers (pure CPU, no heavy maths):
-  • build_nx  – convert agraph nodes/edges → NetworkX graph
-  • get_top_hubs – return top-k nodes by degree-centrality
-  • get_density – overall graph density
+Basic NetworkX helpers for MedGenesis graphs.
+
+Key improvement:
+    build_nx() now accepts edge dictionaries in either of the two
+    common formats:
+
+        • {'source': 'n1', 'target': 'n2'}      (Streamlit-agraph)
+        • {'from':   'n1', 'to':     'n2'}      (PyVis)
+
+This prevents KeyError crashes when nodes / edges come from different
+UI toolchains.
 """
 
 from typing import List, Dict, Tuple
 import networkx as nx
 
+
 # ----------------------------------------------------------------------
+def _edge_endpoints(e: Dict) -> Tuple[str, str] | None:
+    """Return (src, dst) if both ends exist; otherwise 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
+
+
 def build_nx(nodes: List[Dict], edges: List[Dict]) -> nx.Graph:
+    """
+    Convert agraph/PyVis node+edge dicts into a NetworkX Graph.
+
+    * Skips malformed edges rather than raising KeyError.
+    * Node label stored as attribute 'label'.
+    """
     G = nx.Graph()
+
+    # add nodes
     for n in nodes:
         G.add_node(n["id"], label=n.get("label", n["id"]))
+
+    # add edges
     for e in edges:
-        G.add_edge(e["source"], e["target"])
+        endpoints = _edge_endpoints(e)
+        if endpoints:
+            G.add_edge(*endpoints)
+
     return G
 
 
+# ----------------------------------------------------------------------
 def get_top_hubs(G: nx.Graph, k: int = 5) -> List[Tuple[str, float]]:
-    """
-    Return [(node_id, centrality)] sorted desc.
-    """
+    """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:
+    """Return graph density in [0,1]."""
     return nx.density(G)