Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -1,37 +1,52 @@
 # mcp/knowledge_graph.py
 from streamlit_agraph import Node, Edge, Config
+import re
 
-def build_agraph(res: Dict) -> (list, list, Config):
+def build_agraph(papers, umls, drug_safety, umls_relations=None):
     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"))
 
-    cfg = Config(width="100%", height="600", directed=True,
-                 nodeHighlightBehavior=True, highlightColor="#fdcb6e")
+    # 1) concepts
+    for c in umls:
+        cui, name = c.get("cui"), c.get("name")
+        types = c.get("types", [])
+        if not cui or not name: 
+            continue
+        label = f"{name}\n({'/'.join(types)})"
+        size  = 20 + int(c.get("score",0)*30)
+        nodes.append(Node(
+            id=f"cui_{cui}",
+            label=label,
+            size=size,
+            color="#00b894",
+            tooltip=c.get("definition","")
+        ))
+
+    # 2) concept→concept edges
+    if umls_relations:
+        for rels in umls_relations:
+            for r in rels:
+                s, t = f"cui_{r['cui']}", f"cui_{r['cui2']}"
+                if s != t:
+                    edges.append(Edge(source=s, target=t, label=r["label"]))
+
+    # 3) papers
+    for i,p in enumerate(papers):
+        pid = f"paper_{i}"
+        nodes.append(Node(
+            id=pid, label=f"P{i+1}", size=15, color="#0984e3",
+            tooltip=p.get("title","")
+        ))
+        txt = (p.get("title","") + " " + p.get("summary","")).lower()
+        for c in umls:
+            if c.get("name","").lower() in txt:
+                edges.append(Edge(source=pid, target=f"cui_{c['cui']}", label="mentions"))
+
+    # 4) drugs (as before)…
+    # nodes.extend(...), edges.extend(...)
+
+    cfg = Config(
+        width="100%", height="600px", directed=True,
+        nodeHighlightBehavior=True, highlightColor="#f1c40f",
+        collapsible=True, node={"labelProperty":"label"}
+    )
     return nodes, edges, cfg