mcp/knowledge_graph.py

# 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