Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -1,82 +1,37 @@
 # 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.
-    """
+def build_agraph(res: Dict) -> (list, list, Config):
     nodes, edges = [], []
+    # add each paper as a node
+    for i,p in enumerate(res["papers"]):
+        nid = f"paper_{i}"
+        nodes.append(Node(id=nid, label=p["title"], size=20, color="#0984e3"))
+        # connect to AI summary?
+    # add UMLS concepts
+    for u in res["umls"]:
+        cid = f"cui_{u['cui']}"
+        label = f"{u['name']} ({u['cui']})"
+        nodes.append(Node(id=cid, label=label, size=25, color="#00b894"))
+        # connect concept → first paper
+        edges.append(Edge(source=cid, target="paper_0", label="mentioned_in"))
+    # genes
+    g = res.get("gene",{})
+    if g:
+        gid = "gene_node"
+        nodes.append(Node(id=gid, label=g.get("symbol",g.get("name","gene")), color="#d63031"))
+        edges.append(Edge(source=gid, target="cui_"+res["umls"][0]["cui"], label="related"))
+    # variants
+    for v in res["variants"]:
+        vid = f"var_{v['mutationId']}"
+        nodes.append(Node(id=vid, label=v["mutationId"], color="#fdcb6e", size=15))
+        edges.append(Edge(source=vid, target=gid, label="affects"))
+    # trials
+    for t in res["trials"]:
+        tid = t["NCTId"][0]
+        nodes.append(Node(id=tid, label=tid, color="#6c5ce7"))
+        edges.append(Edge(source=tid, target=gid, label="studies"))
 
-    # 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
+    cfg = Config(width="100%", height="600", directed=True,
+                 nodeHighlightBehavior=True, highlightColor="#fdcb6e")
+    return nodes, edges, cfg