import tempfile import os from typing import Dict, Any import pyvis.network as net from .graph import build_tree_structure def create_pyvis_network(graph: Dict[str, Any], root_id: int) -> str: nodes = graph.get("nodes", {}) nt = net.Network( height="600px", width="100%", bgcolor="#ffffff", font_color="black", directed=True ) nt.set_options(""" var options = { "physics": { "enabled": true, "hierarchicalRepulsion": { "centralGravity": 0.3, "springLength": 120, "springConstant": 0.01, "nodeDistance": 200, "damping": 0.09 }, "solver": "hierarchicalRepulsion" }, "layout": { "hierarchical": { "enabled": true, "direction": "UD", "sortMethod": "directed" } } } """) tree = build_tree_structure(graph, root_id) colors = ["#ff6b6b", "#4ecdc4", "#45b7d1", "#96ceb4", "#feca57", "#ff9ff3", "#54a0ff"] for node_id, node in nodes.items(): name = node.get("name", str(node_id)) year_str = f" ({node.get('year')})" if node.get('year') is not None else "" label = f"{name}{year_str}" depth = tree.get(node_id, {}).get("depth", 0) color = colors[depth % len(colors)] nt.add_node( node_id, label=label, title=f"Name: {name}\nYear: {node.get('year', 'N/A')}\nInstitution: {node.get('institution', 'N/A')}", color=color, size=30 if node_id == root_id else 20 ) for node_id, node in nodes.items(): for advisor_id in node.get("advisors", []): if advisor_id in nodes: nt.add_edge(advisor_id, node_id) temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".html") nt.save_graph(temp_file.name) with open(temp_file.name, 'r') as f: html_content = f.read() os.unlink(temp_file.name) return html_content