Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -1,43 +1,61 @@
 # mcp/knowledge_graph.py
 
-from streamlit_agraph import agraph, Node, Edge, Config
+from streamlit_agraph import Node, Edge, Config
+import re
 
-def build_agraph(papers, umls, drugs):
+def build_agraph(papers, umls, drug_safety):
     """
-    Build nodes and edges for streamlit-agraph visualization.
+    Build interactive agraph nodes and edges.
+    Handles drug_safety entries that may be dict or list.
     """
+
     nodes, edges = [], []
 
-    # Map concepts
+    # Add UMLS concept nodes
     for c in umls:
-        if c.get("cui"):
-            cid = f"concept_{c['cui']}"
-            nodes.append(Node(id=cid, label=c["name"], size=25, color="#00b894"))
-    
-    # Map drugs
-    for i, drug_report in enumerate(drugs):
-        if drug_report:
-            did = f"drug_{i}"
-            dname = drug_report.get("drug_name", f"drug_{i}")
-            nodes.append(Node(id=did, label=dname, size=25, color="#d35400"))
-
-    # Map papers
-    for i, p in enumerate(papers):
-        pid = f"paper_{i}"
-        nodes.append(Node(id=pid, label=f"P{i+1}", tooltip=p["title"], size=15, color="#0984e3"))
-        # connect to concepts
+        cui = c.get("cui")
+        name = c.get("name", "")
+        if cui and name:
+            nid = f"concept_{cui}"
+            nodes.append(Node(id=nid, label=name, size=25, color="#00b894"))
+
+    # Add drug nodes, handling list or dict
+    drug_names = []
+    for i, dr in enumerate(drug_safety):
+        if not dr:
+            continue
+
+        # Normalize to single dict
+        recs = dr if isinstance(dr, list) else [dr]
+        for j, rec in enumerate(recs):
+            # Attempt to extract a drug name
+            dn = rec.get("drug_name") or rec.get("patient", {}).get("drug", "") or rec.get("medicinalproduct", "")
+            dn = 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="#d35400"))
+
+    # Add paper nodes and link to concepts & drugs
+    for pi, p in enumerate(papers):
+        pid = f"paper_{pi}"
+        nodes.append(Node(id=pid, label=f"P{pi+1}", tooltip=p["title"], size=15, color="#0984e3"))
+
+        text = f"{p.get('title','')} {p.get('summary','')}".lower()
+        # Link to concepts
         for c in umls:
-            if c["name"].lower() in (p["title"] + p["summary"]).lower():
-                edges.append(Edge(source=pid, target=f"concept_{c['cui']}", label="mentions"))
-        # connect to drugs
-        for j, drug_report in enumerate(drugs):
-            dname = drug_report.get("drug_name", "")
-            if dname and dname.lower() in (p["title"] + p["summary"]).lower():
-                edges.append(Edge(source=pid, target=f"drug_{j}", label="mentions"))
-    
+            cname = c.get("name", "")
+            cui = c.get("cui")
+            if cname and cui and cname.lower() in text:
+                edges.append(Edge(source=pid, target=f"concept_{cui}", label="mentions"))
+        # Link to drugs
+        for did, dn in drug_names:
+            if dn.lower() in text:
+                edges.append(Edge(source=pid, target=did, label="mentions"))
+
     config = Config(
         width="100%", height="600", directed=False,
-        nodeHighlightBehavior=True, highlightColor="#f0a",
-        collapsible=True
+        nodeHighlightBehavior=True, highlightColor="#f1c40f",
+        collapsible=True,
+        node={"labelProperty": "label"}
     )
     return nodes, edges, config