Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -1,161 +1,182 @@
-#!/usr/bin/env python3
-"""MedGenesis – knowledge‑graph builder for Streamlit‑Agraph.
-
-This version recognises **all new enrichment layers** introduced in the
-latest orchestrator:
-    • UMLS concepts                → green nodes
-    • MyGene / NCBI gene hits      → purple nodes
-    • openFDA / DrugCentral drugs  → orange nodes
-    • ClinicalTrials.gov studies   → pink nodes
-    • Open Targets associations    → red drug–gene / gene–disease edges
-    • Literature papers            → blue nodes (tooltip = title)
-
-The entry‑point `build_agraph` now receives a richer payload and returns
-*(nodes, edges, config)* ready for `streamlit_agraph.agraph`.
+# mcp/knowledge_graph.py
 """
+Build an interactive Network Graph (Streamlit-Agraph) that links together
+the heterogeneous entities returned by the MedGenesis pipeline:
+
+• Papers (arXiv / PubMed)
+• UMLS concepts
+• DrugSafety records (OpenFDA)
+• Genes (optional – when available)
+
+Design goals
+------------
+1. **Resilience.**  Any malformed record (e.g. a RuntimeError placeholder from
+   an upstream failure) is skipped silently - the graph still renders.
+2. **Deduplication.**  Identical nodes (same CUI, drug-name, paper ID…) are
+   added only once; edges are de-duplicated as well.
+3. **Visual semantics.**
+      • Papers          → blue   (#3498db)
+      • UMLS concepts   → green  (#2ecc71)
+      • Drugs / safety  → orange (#e67e22)
+      • Genes           → purple (#9b59b6)
+4. **Config tuned for large graphs.**  Collapsible nodes, physics-based layout,
+   highlight on hover, full-width canvas.
+"""
+
 from __future__ import annotations
 
+from typing import Any, List, Tuple, Dict, Iterable
 import re
-from typing import List, Dict, Tuple
+from collections import defaultdict
 
 from streamlit_agraph import Node, Edge, Config
 
-# ---------------------------------------------------------------------
-# Colour palette (flat‑UI)
-# ---------------------------------------------------------------------
-C_PAPER   = "#0984e3"
-C_CONCEPT = "#00b894"
-C_GENE    = "#6c5ce7"
-C_DRUG    = "#d35400"
-C_TRIAL   = "#fd79a8"
-C_OT_EDGE = "#c0392b"
-
-
-# ---------------------------------------------------------------------
-# Helper builders
-# ---------------------------------------------------------------------
-
-def _add_node(nodes: List[Node], node_id: str, label: str, color: str, tooltip: str | None = None, size: int = 25):
-    """Append Node only if id not yet present (agraph duplicates crash)."""
-    if any(n.id == node_id for n in nodes):
-        return
-    nodes.append(Node(id=node_id, label=label, color=color, size=size, tooltip=tooltip))
-
-
-def _match(text: str, pattern: str) -> bool:
-    return bool(re.search(re.escape(pattern), text, flags=re.I))
-
-
-# ---------------------------------------------------------------------
-# Public API
-# ---------------------------------------------------------------------
-
+# --------------------------------------------------------------------------- #
+# utility helpers                                                             #
+# --------------------------------------------------------------------------- #
+RGB = {
+    "paper":   "#3498db",
+    "umls":    "#2ecc71",
+    "drug":    "#e67e22",
+    "gene":    "#9b59b6",
+}
+
+
+def _safe_iter(obj: Any) -> Iterable:
+    """Yield from *obj* if it is list-like, else yield the obj itself."""
+    if obj is None:
+        return ()
+    if isinstance(obj, (list, tuple, set)):
+        return obj
+    return (obj,)
+
+
+def _dedup(seq: Iterable[Tuple]) -> List[Tuple]:
+    """Remove duplicates while preserving order."""
+    seen = set()
+    out: List[Tuple] = []
+    for item in seq:
+        if item not in seen:
+            out.append(item)
+            seen.add(item)
+    return out
+
+
+# --------------------------------------------------------------------------- #
+# main builder                                                                #
+# --------------------------------------------------------------------------- #
 def build_agraph(
-    papers: List[Dict],
-    umls: List[Dict],
-    drug_safety: List[Dict],
-    genes: List[Dict] | None = None,
-    trials: List[Dict] | None = None,
-    ot_associations: List[Dict] | None = None,
-):
-    """Return (nodes, edges, config) for streamlit_agraph. Safe‑duplicates.
-
+    papers: List[Dict[str, Any]],
+    umls:   List[Dict[str, Any]],
+    drug_safety: List[Dict[str, Any]],
+    genes: List[Dict[str, Any]] | None = None,
+) -> Tuple[List[Node], List[Edge], Config]:
+    """
     Parameters
     ----------
-    papers          : PubMed / arXiv merged list (dicts with title & summary).
-    umls            : List of UMLS concept dicts `{cui, name}`.
-    drug_safety     : openFDA / DrugCentral outputs (mixed dict / list).
-    genes           : Optional list with MyGene/NCBI dicts (symbol, name,...).
-    trials          : Optional ClinicalTrials.gov v2 studies list.
-    ot_associations : Optional list from Open Targets.
+    papers : list[dict]
+        Output of `fetch_arxiv`/`fetch_pubmed` (must have *title* & *summary*).
+    umls : list[dict]
+        Items from `lookup_umls` (may be RuntimeError objects if failed).
+    drug_safety : list[dict]
+        Flattened OpenFDA adverse-event records.
+    genes : list[dict] | None
+        Optional – from gene resolver hub.
+
+    Returns
+    -------
+    nodes, edges, config
+        Objects suitable for `streamlit_agraph.agraph(...)`
     """
-
     nodes: List[Node] = []
     edges: List[Edge] = []
 
-    # 1️⃣ Concepts ----------------------------------------------------
-    for c in umls:
-        cui, name = c.get("cui"), c.get("name", "")
-        if cui and name:
-            cid = f"concept_{cui}"
-            _add_node(nodes, cid, name, C_CONCEPT)
-
-    # 2️⃣ Genes -------------------------------------------------------
-    genes = genes or []
-    for g in genes:
-        sym = g.get("symbol") or g.get("name")
-        gid = f"gene_{sym}"
-        tooltip = g.get("summary", "")
-        _add_node(nodes, gid, sym, C_GENE, tooltip=tooltip)
-
-    # 3️⃣ Drugs (normalize mixed structures) -------------------------
-    drug_tuples: List[Tuple[str, str]] = []  # (node_id, drug_name)
-    for i, dr in enumerate(drug_safety):
-        recs = dr if isinstance(dr, list) else [dr]
-        for j, rec in enumerate(recs):
-            name = (
-                rec.get("drug_name") or
-                rec.get("patient", {}).get("drug") or
-                rec.get("medicinalproduct") or
-                f"drug_{i}_{j}"
-            )
-            did = f"drug_{i}_{j}"
-            drug_tuples.append((did, name))
-            _add_node(nodes, did, name, C_DRUG)
-
-    # 4️⃣ Trials ------------------------------------------------------
-    trials = trials or []
-    for t in trials:
-        nct  = t.get("nctId") or t.get("nctid")
-        if not nct:
+    # ---- Papers ----------------------------------------------------------- #
+    for idx, p in enumerate(_safe_iter(papers)):
+        if not isinstance(p, dict):
             continue
-        tid  = f"trial_{nct}"
-        label = nct
-        tooltip = t.get("briefTitle") or "Clinical trial"
-        _add_node(nodes, tid, label, C_TRIAL, tooltip=tooltip, size=20)
-
-    # 5️⃣ Papers & mention edges -------------------------------------
-    for idx, p in enumerate(papers):
         pid = f"paper_{idx}"
-        _add_node(nodes, pid, f"P{idx+1}", C_PAPER, tooltip=p.get("title", ""), size=15)
-
-        text_blob = f"{p.get('title','')} {p.get('summary','')}".lower()
-
-        # concept links
-        for c in umls:
-            if c.get("name") and _match(text_blob, c["name"]):
-                edges.append(Edge(source=pid, target=f"concept_{c['cui']}", label="mentions"))
-        # gene links
-        for g in genes:
-            if g.get("symbol") and _match(text_blob, g["symbol"]):
-                edges.append(Edge(source=pid, target=f"gene_{g['symbol']}", label="mentions"))
-        # drug links
-        for did, dname in drug_tuples:
-            if _match(text_blob, dname):
-                edges.append(Edge(source=pid, target=did, label="mentions"))
-
-    # 6️⃣ Open Targets edges (drug–gene / gene–disease) --------------
-    if ot_associations:
-        for row in ot_associations:
-            gsym = row.get("target", {}).get("symbol")
-            dis  = row.get("disease", {}).get("name")
-            score = row.get("score", 0)
-            if gsym and dis:
-                gid = f"gene_{gsym}"
-                did = f"disease_{dis}"
-                _add_node(nodes, did, dis, C_CONCEPT, size=20)
-                edges.append(Edge(source=gid, target=did, color=C_OT_EDGE, label=f"OT {score:.2f}"))
-
-    # 7️⃣ Config ------------------------------------------------------
+        label = f"P{idx + 1}"
+        tooltip = p.get("title", "Paper")
+        nodes.append(Node(id=pid, label=label, tooltip=tooltip, size=15,
+                          color=RGB["paper"]))
+
+    # ---- UMLS concepts ---------------------------------------------------- #
+    for c in _safe_iter(umls):
+        if not isinstance(c, dict):
+            continue
+        cui = c.get("cui")
+        name = c.get("name")
+        if cui and name:
+            nodes.append(Node(id=f"umls_{cui}", label=name, size=22,
+                              color=RGB["umls"]))
+
+    # ---- Drug Safety ------------------------------------------------------ #
+    for i, rec in enumerate(_safe_iter(drug_safety)):
+        if not isinstance(rec, dict):
+            continue
+        dn = (
+            rec.get("drug_name")
+            or rec.get("medicinalproduct")
+            or "drug_{}".format(i)
+        )
+        nodes.append(Node(id=f"drug_{i}", label=dn, size=25,
+                          color=RGB["drug"]))
+
+    # ---- Genes (optional) -------------------------------------------------- #
+    if genes:
+        for g in _safe_iter(genes):
+            if not isinstance(g, dict):
+                continue
+            sym = g.get("symbol") or g.get("approvedSymbol") or g.get("name")
+            if sym:
+                gid = f"gene_{sym}"
+                nodes.append(Node(id=gid, label=sym, size=20,
+                                  color=RGB["gene"]))
+
+    # ---------------------------------------------------------------------- #
+    # Edges – naïve co-occurrence linking                                    #
+    # ---------------------------------------------------------------------- #
+    paper_texts = [
+        (n.id, f"{p.get('title','')} {p.get('summary','')}".lower())
+        for n, p in zip(nodes, papers)
+        if n.id.startswith("paper_")
+    ]
+
+    # connect paper ↔ umls / drugs / genes if mention appears in text
+    def _link(target_nodes: List[Node], pattern_getter):
+        for nid, blob in paper_texts:
+            for tn in target_nodes:
+                pat = pattern_getter(tn)
+                if pat and pat.search(blob):
+                    edges.append(Edge(source=nid, target=tn.id, label="mentions"))
+
+    umls_nodes = [n for n in nodes if n.id.startswith("umls_")]
+    drug_nodes = [n for n in nodes if n.id.startswith("drug_")]
+    gene_nodes = [n for n in nodes if n.id.startswith("gene_")]
+
+    _link(umls_nodes, lambda n: re.compile(re.escape(n.label.lower())))
+    _link(drug_nodes, lambda n: re.compile(rf"\b{re.escape(n.label.lower())}\b"))
+    _link(gene_nodes, lambda n: re.compile(rf"\b{re.escape(n.label.lower())}\b"))
+
+    # de-duplicate everything ------------------------------------------------ #
+    nodes = _dedup(nodes)
+    edges = _dedup(edges)
+
+    # ---------------------------------------------------------------------- #
+    # Graph config                                                            #
+    # ---------------------------------------------------------------------- #
     cfg = Config(
-        directed=False,
         width="100%",
-        height="600",
+        height="650px",
+        directed=False,
         nodeHighlightBehavior=True,
         highlightColor="#f1c40f",
         collapsible=True,
-        showLegend=False,
+        physics=True,
+        hierarchical=False,
         node={"labelProperty": "label"},
+        link={"labelProperty": "label", "renderLabel": False},
     )
 
     return nodes, edges, cfg