mcp/knowledge_graph.py

from streamlit_agraph import Node, Edge, Config
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="600px", directed=False,
        node={"labelProperty": "label"},
        nodeHighlightBehavior=True, highlightColor="#f1c40f",
        collapsible=True,
    )
    return nodes, edges, cfg