Sleep Knowledge Graph

from __future__ import annotations import argparse import json from collections import Counter from datetime import date, datetime from pathlib import Path from typing import Any from neo4j import GraphDatabase from sleep_db.config import Settings, load_settings LABEL_COLORS = { "Metric": "#38bdf8", "ThresholdCondition": "#f59e0b", "ConditionGroup": "#f97316", "State": "#a78bfa", "Effect": "#fb7185", "Advice": "#34d399", "Activity": "#22c55e", "Nutrition": "#84cc16", "Environment": "#06b6d4", "Doc": "#e5e7eb", "RuleVersion": "#94a3b8", } def fetch_neo4j_graph(settings: Settings) -> dict[str, Any]: driver = GraphDatabase.driver( settings.neo4j_uri, auth=(settings.neo4j_user, settings.neo4j_password), ) try: with driver.session() as session: nodes = [ _node_from_record(record) for record in session.run( """ MATCH (n) RETURN elementId(n) AS element_id, labels(n) AS labels, properties(n) AS props ORDER BY labels(n)[0], coalesce(n.name, n.id, n.doc_id, n.rule_id) """ ) ] edges = [ _edge_from_record(record) for record in session.run( """ MATCH (a)-[r]->(b) RETURN elementId(r) AS element_id, elementId(a) AS source_element_id, elementId(b) AS target_element_id, type(r) AS type, properties(r) AS props ORDER BY type(r) """ ) ] finally: driver.close() label_counts = Counter(node["label"] for node in nodes) edge_counts = Counter(edge["type"] for edge in edges) return { "generated_at": datetime.now().isoformat(timespec="seconds"), "nodes": nodes, "edges": edges, "stats": { "node_count": len(nodes), "edge_count": len(edges), "label_counts": dict(label_counts), "edge_type_counts": dict(edge_counts), }, } def write_graph_visualization(graph: dict[str, Any], output_dir: Path) -> dict[str, Path]: output_dir.mkdir(parents=True, exist_ok=True) json_path = output_dir / "neo4j_graph.json" html_path = output_dir / "neo4j_graph.html" json_path.write_text(json.dumps(graph, indent=2, ensure_ascii=False), encoding="utf-8") html_path.write_text(render_graph_html(graph), encoding="utf-8") return {"json": json_path, "html": html_path} def render_graph_html(graph: dict[str, Any]) -> str: graph_json = json.dumps(graph, ensure_ascii=False) label_color_json = json.dumps(LABEL_COLORS) return f""" Sleep Knowledge Graph

""" def _node_from_record(record: Any) -> dict[str, Any]: labels = list(record["labels"]) props = _jsonable(dict(record["props"])) label = _primary_label(labels) display_id = props.get("id") or props.get("doc_id") or props.get("rule_id") or record["element_id"] return { "element_id": record["element_id"], "display_id": display_id, "label": label, "labels": labels, "name": props.get("name") or props.get("title") or display_id, "props": props, } def _edge_from_record(record: Any) -> dict[str, Any]: return { "element_id": record["element_id"], "source_element_id": record["source_element_id"], "target_element_id": record["target_element_id"], "type": record["type"], "props": _jsonable(dict(record["props"])), } def _primary_label(labels: list[str]) -> str: for label in LABEL_COLORS: if label in labels: return label return labels[0] if labels else "Node" def _jsonable(value: Any) -> Any: if isinstance(value, dict): return {str(key): _jsonable(item) for key, item in value.items()} if isinstance(value, list): return [_jsonable(item) for item in value] if isinstance(value, (datetime, date)): return value.isoformat() if isinstance(value, (str, int, float, bool)) or value is None: return value return str(value) def main() -> None: parser = argparse.ArgumentParser(description="Export Neo4j sleep graph to shareable HTML.") parser.add_argument("--output-dir", type=Path, default=Path("outputs/visualizations")) args = parser.parse_args() settings = load_settings() graph = fetch_neo4j_graph(settings) paths = write_graph_visualization(graph, args.output_dir) print({key: str(path) for key, path in paths.items()}) print(graph["stats"]) if __name__ == "__main__": main()