mcp/knowledge_graph.py

# mcp/knowledge_graph.py

from streamlit_agraph import Node, Edge, Config
import re

# Set colors for node types
PAPER_COLOR = "#0984e3"
UMLS_COLOR = "#00b894"
DRUG_COLOR = "#d35400"

def build_agraph(papers, umls, drug_safety):
    """
    Build interactive agraph nodes and edges.
    Defensive: handles unexpected types gracefully.
    """
    nodes, edges = [], []

    # UMLS nodes
    for c in umls or []:
        if not isinstance(c, dict):
            continue
        cui = str(c.get("cui", "") or "")
        name = str(c.get("name", "") or "")
        if cui and name:
            nid = f"concept_{cui}"
            nodes.append(Node(
                id=nid, label=name, size=25, color=UMLS_COLOR,
                tooltip=f"UMLS {cui}: {name}"
            ))

    # Drug nodes
    drug_names = []
    for i, dr in enumerate(drug_safety or []):
        if not dr:
            continue
        # Normalize to single dict
        recs = dr if isinstance(dr, list) else [dr]
        for j, rec in enumerate(recs):
            if not isinstance(rec, dict):
                continue
            dn = rec.get("drug_name") \
                or (rec.get("patient", {}) or {}).get("drug", "") \
                or rec.get("medicinalproduct", "")
            dn = str(dn or f"drug_{i}_{j}")
            did = f"drug_{i}_{j}"
            drug_names.append((did, dn))
            nodes.append(Node(id=did, label=dn, size=25, color=DRUG_COLOR,
                              tooltip=f"Drug: {dn}"))

    # Paper nodes and edges
    for k, p in enumerate(papers or []):
        pid = f"paper_{k}"
        title = str(p.get("title", f"Paper {k+1}"))
        summary = str(p.get("summary", ""))
        label = f"P{k+1}"
        nodes.append(Node(
            id=pid,
            label=label,
            tooltip=title,
            size=14,
            color=PAPER_COLOR,
        ))
        txt = (title + " " + summary).lower()
        # Link to concepts
        for c in umls or []:
            name = str(c.get("name", "") or "")
            cui = str(c.get("cui", "") or "")
            if name and name.lower() in txt and cui:
                edges.append(Edge(source=pid, target=f"concept_{cui}", label="mentions"))
        # Link to drugs
        for did, dn in drug_names:
            if dn and dn.lower() in txt:
                edges.append(Edge(source=pid, target=did, label="mentions"))

    config = Config(
        width="100%", height="600", directed=False,
        nodeHighlightBehavior=True, highlightColor="#f1c40f",
        collapsible=True,
        node={"labelProperty": "label"},
        link={"labelProperty": "label"},
    )
    return nodes, edges, config