Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -1,93 +1,62 @@
-"""
-Graph builder for Streamlit-agraph.
-• Safely skips bad nodes (RuntimeError / None)
-• Colour-codes    papers  = #3498db  (blue)
-                  drugs   = #e67e22  (orange)
-                  concepts= #2ecc71  (green)
-"""
-
-from __future__ import annotations
-import re
-from typing import List, Dict, Tuple
 from streamlit_agraph import Node, Edge, Config
-
-# ── colour palette ----------------------------------------------------
-CLR_PAPER   = "#3498db"
-CLR_DRUG    = "#e67e22"
-CLR_CONCEPT = "#2ecc71"
-
-def _add_node_safe(nodes: List[Node], nid: str, **kwargs) -> None:
-    if not any(n.id == nid for n in nodes):
-        nodes.append(Node(id=nid, **kwargs))
-
-def build_agraph(papers: List[Dict],
-                 umls  : List[Dict],
-                 safety: List[Dict]) -> Tuple[List[Node], List[Edge], Config]:
-    """
-    Map MedGenesis payload → (nodes, edges, config) for Streamlit-agraph.
-    """
-
-    nodes: List[Node] = []
-    edges: List[Edge] = []
-
-    # 1 UMLS CONCEPTS ---------------------------------------------------
-    for concept in umls:
-        try:
-            cui  = concept.get("cui")
-            name = concept.get("name", "")
-            if cui and name:
-                cid = f"CUI_{cui}"
-                _add_node_safe(nodes, cid, label=name, size=25, color=CLR_CONCEPT)
-        except Exception:
-            continue  # malformed concept
-
-    # 2 DRUGS (OpenFDA) -------------------------------------------------
-    drug_label_pairs: List[Tuple[str, str]] = []
-    for i, rec in enumerate(safety):
-        if not rec:       # may be {} if rate-limited
-            continue
-        drug_name = (rec.get("drug_name") or
-                     rec.get("patient", {}).get("drug") or
-                     rec.get("medicinalproduct") or
-                     f"unknown_{i}")
-        did = f"DRUG_{i}"
-        drug_label_pairs.append((did, drug_name))
-        _add_node_safe(nodes, did, label=drug_name, size=22, color=CLR_DRUG)
-
-    # 3 PAPERS ----------------------------------------------------------
-    for idx, p in enumerate(papers):
-        pid = f"PAPER_{idx}"
-        _add_node_safe(nodes, pid,
-                       label=f"P{idx+1}",
-                       tooltip=p.get("title", ""),
-                       size=15,
-                       color=CLR_PAPER)
-
-        plain = (p.get("title", "") + " " + p.get("summary", "")).lower()
-
-        # (i) concept edges
-        for concept in umls:
-            name = concept.get("name", "")
-            cui  = concept.get("cui")
-            if name and cui and name.lower() in plain:
-                edges.append(Edge(source=pid,
-                                  target=f"CUI_{cui}",
-                                  label="mentions"))
-
-        # (ii) drug edges
-        for did, dname in drug_label_pairs:
-            if dname.lower() in plain:
-                edges.append(Edge(source=pid,
-                                  target=did,
-                                  label="mentions"))
+import re, itertools
+from typing import List, Tuple
+
+_GREEN = "#00b894"
+_ORANGE = "#d35400"
+_BLUE   = "#0984e3"
+
+def _safe(node_like):
+    """Return empty dict if node_like is an Exception."""
+    return node_like if isinstance(node_like, dict) else {}
+
+def build_agraph(
+    papers: List[dict],
+    umls: List[dict],
+    drug_safety: List[dict],
+) -> Tuple[List[Node], List[Edge], Config]:
+    nodes, edges = [], []
+
+    # UMLS concept nodes
+    for c in filter(bool, map(_safe, umls)):
+        cui, name = c.get("cui"), c.get("name", "")
+        if cui and name:
+            nid = f"cui:{cui}"
+            nodes.append(Node(id=nid, label=name, color=_GREEN, size=25))
+
+    # Drug nodes
+    def _drug_name(d: dict) -> str:
+        return (
+            d.get("drug_name")
+            or d.get("patient", {}).get("drug", "")
+            or d.get("medicinalproduct", "")
+        )
+
+    for idx, rec in enumerate(itertools.chain.from_iterable(
+            [r if isinstance(r, list) else [r] for r in drug_safety])):
+        dn = _drug_name(rec) or f"drug_{idx}"
+        did = f"drug_{idx}"
+        nodes.append(Node(id=did, label=dn, color=_ORANGE, size=25))
+
+    # Paper nodes
+    for i, p in enumerate(papers, 1):
+        pid = f"paper_{i}"
+        nodes.append(Node(id=pid, label=f"P{i}", tooltip=p["title"], color=_BLUE, size=15))
+
+        txt = f"{p['title']} {p['summary']}".lower()
+        # link ↔ concepts
+        for c in filter(bool, map(_safe, umls)):
+            if (name := c.get("name", "")).lower() in txt and c.get("cui"):
+                edges.append(Edge(source=pid, target=f"cui:{c['cui']}", label="mentions"))
+        # link ↔ drugs
+        for n in nodes:
+            if n.id.startswith("drug_") and n.label.lower() in txt:
+                edges.append(Edge(source=pid, target=n.id, label="mentions"))
 
     cfg = Config(
-        width        = "100%",
-        height       = 600,
-        directed     = False,
-        nodeHighlightBehavior = True,
-        highlightColor = "#f1c40f",
-        collapsible  = True,
-        node         = {"labelProperty": "label"}
+        width="100%", height="600px", directed=False,
+        node={"labelProperty": "label"},
+        nodeHighlightBehavior=True, highlightColor="#f1c40f",
+        collapsible=True,
     )
     return nodes, edges, cfg