Update mcp/knowledge_graph.py

mcp/knowledge_graph.py

@@ -1,61 +1,43 @@
 # mcp/knowledge_graph.py
 
-from pyvis.network import Network
-import tempfile
-import jinja2
-assert jinja2.Environment is not None
+from streamlit_agraph import agraph, Node, Edge, Config
 
-def build_knowledge_graph(papers, umls, drugs):
+def build_agraph(papers, umls, drugs):
     """
-    Build and return a pyvis Network object representing the knowledge graph.
-    Nodes: papers, UMLS concepts, drugs.
-    Edges: paper–concept, paper–drug, drug–concept.
+    Build nodes and edges for streamlit-agraph visualization.
     """
-    net = Network(height="550px", width="100%", bgcolor="#f3eaff", font_color="#222222", notebook=False, directed=False)
-    net.force_atlas_2based()  # More visually appealing layout
-
-    # Add UMLS concept nodes
-    concept_nodes = {}
-    for concept in umls:
-        if concept["cui"]:
-            node_id = f"concept_{concept['cui']}"
-            net.add_node(node_id, label=concept["name"], color="#00b894", title=f"UMLS: {concept['cui']}\n{concept['definition'] or ''}")
-            concept_nodes[concept['name']] = node_id
-
-    # Add drug nodes
-    drug_nodes = {}
-    for drug_report in drugs:
-        if isinstance(drug_report, list):
-            # OpenFDA can return a list of dicts
-            for d in drug_report:
-                if "safety_report_id" in d:
-                    node_id = f"drug_{d['safety_report_id']}"
-                    drug_name = d.get('reactions', ['drug'])[0] if d.get('reactions') else 'drug'
-                    net.add_node(node_id, label=drug_name, color="#d35400", title=str(d))
-                    drug_nodes[drug_name] = node_id
-
-    # Add paper nodes and edges to concepts and drugs
-    for i, paper in enumerate(papers, 1):
-        paper_id = f"paper_{i}"
-        net.add_node(paper_id, label=paper["title"][:40] + "...", color="#0984e3", title=f"{paper['title']}\n{paper['summary']}\n{paper['link']}")
-
-        # Link paper to concepts by presence in summary/title
-        for cname, concept_id in concept_nodes.items():
-            if cname.lower() in (paper["summary"] + paper["title"]).lower():
-                net.add_edge(paper_id, concept_id, color="#00b894")
-
-        # Link paper to drugs if drug name is present in summary/title
-        for dname, drug_id in drug_nodes.items():
-            if dname.lower() in (paper["summary"] + paper["title"]).lower():
-                net.add_edge(paper_id, drug_id, color="#d35400")
-
-    # Optionally: link concepts to drugs if name overlap
-    for cname, concept_id in concept_nodes.items():
-        for dname, drug_id in drug_nodes.items():
-            if cname.lower() in dname.lower() or dname.lower() in cname.lower():
-                net.add_edge(concept_id, drug_id, color="#636e72")
-
-    # Save and return path to HTML
-    temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".html").name
-    net.show(temp_path)
-    return temp_path
+    nodes, edges = [], []
+
+    # Map concepts
+    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
+        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"))
+    
+    config = Config(
+        width="100%", height="600", directed=False,
+        nodeHighlightBehavior=True, highlightColor="#f0a",
+        collapsible=True
+    )
+    return nodes, edges, config