File size: 2,715 Bytes
3d1def9 447d959 3d1def9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
# mcp/knowledge_graph.py
from pyvis.network import Network
import tempfile
import jinja2
assert jinja2.Environment is not None
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
|