Create knowledge_graph.py

mcp/knowledge_graph.py

@@ -0,0 +1,59 @@
+# mcp/knowledge_graph.py
+
+from pyvis.network import Network
+import tempfile
+
+def build_knowledge_graph(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.
+    """
+    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