Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -3,56 +3,55 @@
 from streamlit_agraph import Node, Edge, Config
 import re
 
-# Colors for graph nodes
-EDGE_COLOR     = "#888"
-DRUG_COLOR     = "#f39c12"
-CONCEPT_COLOR  = "#00b894"
-PAPER_COLOR    = "#3498db"
-HL_COLOR       = "#f1c40f"
-DIM_COLOR      = "#d3d3d3"
+# Set colors for node types
+PAPER_COLOR = "#0984e3"
+UMLS_COLOR = "#00b894"
+DRUG_COLOR = "#d35400"
 
 def build_agraph(papers, umls, drug_safety):
     """
-    Build a Streamlit-agraph network:
-      - papers: list of PubMed/arXiv dicts
-      - umls: list of UMLSConcept dicts (may have None values)
-      - drug_safety: list of OpenFDA/other dicts
-    Returns (nodes, edges, config)
+    Build interactive agraph nodes and edges.
+    Defensive: handles unexpected types gracefully.
     """
     nodes, edges = [], []
 
-    # --- Add UMLS concept nodes ---
-    for c in (umls or []):
-        cui = c.get("cui") if c else None
-        name = c.get("name") if c else None
-        if not cui or not name:
+    # UMLS nodes
+    for c in umls or []:
+        if not isinstance(c, dict):
             continue
-        node_id = f"concept_{cui}"
-        nodes.append(Node(id=node_id, label=name, size=22, color=CONCEPT_COLOR))
+        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}"
+            ))
 
-    # --- Add Drug nodes ---
-    drug_ids = []
-    for i, drug_blob in enumerate(drug_safety or []):
-        # Support both list and dict style safety reports
-        if not drug_blob:
+    # Drug nodes
+    drug_names = []
+    for i, dr in enumerate(drug_safety or []):
+        if not dr:
             continue
-        reports = drug_blob if isinstance(drug_blob, list) else [drug_blob]
-        for j, rec in enumerate(reports):
-            label = (
-                rec.get("drug_name")
-                or rec.get("patient", {}).get("drug")
-                or rec.get("medicinalproduct")
-                or f"drug_{i}_{j}"
-            )
-            drug_id = f"drug_{i}_{j}"
-            drug_ids.append((drug_id, label))
-            nodes.append(Node(id=drug_id, label=label, size=25, color=DRUG_COLOR))
+        # 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}"))
 
-    # --- Add Paper nodes and connect to concepts/drugs ---
+    # Paper nodes and edges
     for k, p in enumerate(papers or []):
         pid = f"paper_{k}"
-        title = p.get("title", f"Paper {k+1}")
-        summary = p.get("summary", "")
+        title = str(p.get("title", f"Paper {k+1}"))
+        summary = str(p.get("summary", ""))
         label = f"P{k+1}"
         nodes.append(Node(
             id=pid,
@@ -62,31 +61,22 @@ def build_agraph(papers, umls, drug_safety):
             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"))
 
-        # Link to UMLS concepts if concept name in paper
-        for c in (umls or []):
-            cui = c.get("cui") if c else None
-            name = c.get("name") if c else None
-            if cui and name and isinstance(name, str) and name.lower() in txt:
-                edges.append(Edge(source=pid, target=f"concept_{cui}", color=EDGE_COLOR))
-
-        # Link to drug nodes if drug name appears in paper
-        for drug_id, drug_name in drug_ids:
-            if drug_name and isinstance(drug_name, str) and drug_name.lower() in txt:
-                edges.append(Edge(source=pid, target=drug_id, color=EDGE_COLOR))
-
-    # --- Graph config with physics enabled ---
-    cfg = Config(
-        width="100%",
-        height="520",
-        directed=False,
-        physics=True,
-        repulsion=True,
-        nodeHighlightBehavior=True,
-        highlightColor=HL_COLOR,
+    config = Config(
+        width="100%", height="600", directed=False,
+        nodeHighlightBehavior=True, highlightColor="#f1c40f",
         collapsible=True,
         node={"labelProperty": "label"},
-        edge={"color": EDGE_COLOR, "width": 1},
+        link={"labelProperty": "label"},
     )
-    return nodes, edges, cfg
-
+    return nodes, edges, config