sleep_ai_demo / db /src /sleep_db /visualize_graph.py
lokework's picture
Upload sleep full-agent demo
6025aa5 verified
Raw
History Blame Contribute Delete
22.3 kB
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"""<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1" />
<title>Sleep Knowledge Graph</title>
<style>
:root {{
color-scheme: dark;
--bg: #0b1020;
--panel: #111827;
--panel-2: #172033;
--line: #253249;
--text: #f8fafc;
--muted: #9ca3af;
--accent: #38bdf8;
}}
* {{ box-sizing: border-box; }}
body {{
margin: 0;
background: radial-gradient(circle at 20% 0%, #172554 0%, #0b1020 38%, #050816 100%);
color: var(--text);
font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
overflow: hidden;
}}
.shell {{
display: grid;
grid-template-columns: 320px 1fr 360px;
height: 100vh;
min-height: 680px;
}}
aside {{
background: rgba(17, 24, 39, 0.92);
border-right: 1px solid var(--line);
padding: 18px;
overflow: auto;
backdrop-filter: blur(14px);
}}
.right {{
border-left: 1px solid var(--line);
border-right: 0;
}}
h1, h2, h3 {{
margin: 0;
font-weight: 700;
letter-spacing: 0;
}}
h1 {{ font-size: 21px; line-height: 1.2; }}
h2 {{ font-size: 13px; color: var(--muted); margin-top: 4px; font-weight: 500; }}
h3 {{ font-size: 12px; color: #cbd5e1; text-transform: uppercase; margin: 18px 0 10px; }}
.canvas-wrap {{ position: relative; min-width: 0; }}
canvas {{ width: 100%; height: 100%; display: block; cursor: grab; }}
canvas:active {{ cursor: grabbing; }}
.toolbar {{
position: absolute;
top: 16px;
left: 16px;
display: flex;
gap: 8px;
flex-wrap: wrap;
max-width: calc(100% - 32px);
}}
button, input {{
background: rgba(15, 23, 42, 0.9);
border: 1px solid #334155;
color: var(--text);
border-radius: 6px;
padding: 9px 10px;
font: inherit;
}}
button {{ cursor: pointer; }}
button:hover {{ border-color: var(--accent); }}
input {{ width: 260px; }}
.stats {{
display: grid;
grid-template-columns: 1fr 1fr;
gap: 10px;
margin-top: 16px;
}}
.stat {{
background: var(--panel-2);
border: 1px solid var(--line);
border-radius: 8px;
padding: 12px;
}}
.stat strong {{ display: block; font-size: 22px; }}
.stat span {{ color: var(--muted); font-size: 12px; }}
.filter-row {{
display: grid;
grid-template-columns: 18px 16px 1fr auto;
align-items: center;
gap: 8px;
padding: 6px 0;
color: #d1d5db;
font-size: 13px;
}}
.swatch {{
width: 12px;
height: 12px;
border-radius: 50%;
}}
.pill {{
border: 1px solid #334155;
border-radius: 999px;
padding: 2px 7px;
color: var(--muted);
font-size: 11px;
}}
.details {{
background: var(--panel-2);
border: 1px solid var(--line);
border-radius: 8px;
padding: 12px;
white-space: pre-wrap;
overflow-wrap: anywhere;
font-size: 12px;
line-height: 1.45;
color: #dbeafe;
}}
.list {{
display: flex;
flex-direction: column;
gap: 8px;
}}
.mini {{
border: 1px solid var(--line);
border-radius: 8px;
padding: 9px;
background: rgba(15, 23, 42, 0.72);
font-size: 12px;
}}
.mini b {{ color: #e5e7eb; }}
.mini div {{ color: var(--muted); margin-top: 3px; }}
@media (max-width: 1100px) {{
.shell {{ grid-template-columns: 280px 1fr; }}
.right {{ display: none; }}
}}
</style>
</head>
<body>
<div class="shell">
<aside>
<h1>Sleep Knowledge Graph</h1>
<h2>Neo4j export generated at {graph.get("generated_at", "")}</h2>
<div class="stats">
<div class="stat"><strong id="nodeCount">0</strong><span>visible nodes</span></div>
<div class="stat"><strong id="edgeCount">0</strong><span>visible edges</span></div>
</div>
<h3>Node Labels</h3>
<div id="labelFilters"></div>
<h3>Edge Types</h3>
<div id="edgeFilters"></div>
</aside>
<main class="canvas-wrap">
<div class="toolbar">
<input id="search" placeholder="Search node id, name, rule, doc..." />
<button id="reset">Reset View</button>
<button id="labels">Edge Labels</button>
<button id="freeze">Pause</button>
<button id="png">Export PNG</button>
</div>
<canvas id="graph"></canvas>
</main>
<aside class="right">
<h1>Selection</h1>
<h2>Click a node or edge for details</h2>
<h3>Details</h3>
<div id="details" class="details">No selection yet.</div>
<h3>High Degree Nodes</h3>
<div id="topNodes" class="list"></div>
</aside>
</div>
<script>
const graph = {graph_json};
const colors = {label_color_json};
const canvas = document.getElementById("graph");
const ctx = canvas.getContext("2d");
const search = document.getElementById("search");
const details = document.getElementById("details");
const labelFilters = document.getElementById("labelFilters");
const edgeFilters = document.getElementById("edgeFilters");
let showEdgeLabels = false;
let paused = false;
let selected = null;
let hover = null;
let scale = 1;
let offsetX = 0;
let offsetY = 0;
let dragging = null;
let panning = false;
let lastX = 0;
let lastY = 0;
const activeLabels = new Set([...new Set(graph.nodes.map(n => n.label))]);
const activeEdges = new Set([...new Set(graph.edges.map(e => e.type))]);
const nodeById = new Map(graph.nodes.map(n => [n.element_id, n]));
const degree = new Map();
for (const n of graph.nodes) {{
n.x = Math.random() * 900 - 450;
n.y = Math.random() * 700 - 350;
n.vx = 0;
n.vy = 0;
degree.set(n.element_id, 0);
}}
for (const e of graph.edges) {{
degree.set(e.source_element_id, (degree.get(e.source_element_id) || 0) + 1);
degree.set(e.target_element_id, (degree.get(e.target_element_id) || 0) + 1);
}}
function resize() {{
const rect = canvas.getBoundingClientRect();
canvas.width = Math.max(800, Math.floor(rect.width * devicePixelRatio));
canvas.height = Math.max(600, Math.floor(rect.height * devicePixelRatio));
ctx.setTransform(devicePixelRatio, 0, 0, devicePixelRatio, 0, 0);
}}
function isVisibleNode(n) {{
const q = search.value.trim().toLowerCase();
const text = JSON.stringify([n.display_id, n.name, n.label, n.props]).toLowerCase();
return activeLabels.has(n.label) && (!q || text.includes(q));
}}
function isVisibleEdge(e) {{
const s = nodeById.get(e.source_element_id);
const t = nodeById.get(e.target_element_id);
return s && t && activeEdges.has(e.type) && isVisibleNode(s) && isVisibleNode(t);
}}
function worldToScreen(x, y) {{
return [x * scale + canvas.clientWidth / 2 + offsetX, y * scale + canvas.clientHeight / 2 + offsetY];
}}
function screenToWorld(x, y) {{
return [(x - canvas.clientWidth / 2 - offsetX) / scale, (y - canvas.clientHeight / 2 - offsetY) / scale];
}}
function tick() {{
const nodes = graph.nodes.filter(isVisibleNode);
const edges = graph.edges.filter(isVisibleEdge);
for (let i = 0; i < nodes.length; i++) {{
for (let j = i + 1; j < nodes.length; j++) {{
const a = nodes[i], b = nodes[j];
let dx = a.x - b.x, dy = a.y - b.y;
let d2 = dx * dx + dy * dy + 0.01;
let force = Math.min(1800 / d2, 2.8);
let d = Math.sqrt(d2);
a.vx += (dx / d) * force;
a.vy += (dy / d) * force;
b.vx -= (dx / d) * force;
b.vy -= (dy / d) * force;
}}
}}
for (const e of edges) {{
const a = nodeById.get(e.source_element_id), b = nodeById.get(e.target_element_id);
let dx = b.x - a.x, dy = b.y - a.y;
let d = Math.sqrt(dx * dx + dy * dy) || 1;
let target = e.type === "REFERS_TO" ? 180 : 125;
let force = (d - target) * 0.012;
a.vx += (dx / d) * force;
a.vy += (dy / d) * force;
b.vx -= (dx / d) * force;
b.vy -= (dy / d) * force;
}}
for (const n of nodes) {{
n.vx += -n.x * 0.002;
n.vy += -n.y * 0.002;
n.vx *= 0.82;
n.vy *= 0.82;
if (!n.fixed) {{
n.x += n.vx;
n.y += n.vy;
}}
}}
}}
function draw() {{
ctx.clearRect(0, 0, canvas.clientWidth, canvas.clientHeight);
const edges = graph.edges.filter(isVisibleEdge);
const nodes = graph.nodes.filter(isVisibleNode);
document.getElementById("nodeCount").textContent = nodes.length;
document.getElementById("edgeCount").textContent = edges.length;
ctx.lineCap = "round";
for (const e of edges) {{
const a = nodeById.get(e.source_element_id), b = nodeById.get(e.target_element_id);
const [x1, y1] = worldToScreen(a.x, a.y);
const [x2, y2] = worldToScreen(b.x, b.y);
const isSelected = selected && selected.element_id === e.element_id;
ctx.strokeStyle = isSelected ? "#f8fafc" : edgeColor(e.type);
ctx.globalAlpha = e.type === "REFERS_TO" ? 0.22 : 0.55;
ctx.lineWidth = isSelected ? 3 : e.type === "REFERS_TO" ? 1 : 1.5;
ctx.beginPath();
ctx.moveTo(x1, y1);
ctx.lineTo(x2, y2);
ctx.stroke();
ctx.globalAlpha = 1;
if (showEdgeLabels && scale > 0.5 && e.type !== "REFERS_TO") {{
ctx.fillStyle = "#cbd5e1";
ctx.font = "10px ui-sans-serif";
ctx.fillText(e.type, (x1 + x2) / 2, (y1 + y2) / 2);
}}
}}
for (const n of nodes) {{
const [x, y] = worldToScreen(n.x, n.y);
const r = nodeRadius(n);
const isSelected = selected && selected.element_id === n.element_id;
const isHover = hover && hover.element_id === n.element_id;
ctx.beginPath();
ctx.fillStyle = colors[n.label] || "#64748b";
ctx.shadowColor = ctx.fillStyle;
ctx.shadowBlur = isSelected || isHover ? 22 : 8;
ctx.arc(x, y, r, 0, Math.PI * 2);
ctx.fill();
ctx.shadowBlur = 0;
ctx.strokeStyle = isSelected ? "#ffffff" : "rgba(255,255,255,0.45)";
ctx.lineWidth = isSelected ? 3 : 1;
ctx.stroke();
if (scale > 0.55 || isSelected || isHover) {{
ctx.fillStyle = "#f8fafc";
ctx.font = "11px ui-sans-serif";
ctx.fillText(n.name || n.display_id, x + r + 4, y + 4);
}}
}}
}}
function nodeRadius(n) {{
return Math.min(8 + Math.sqrt((degree.get(n.element_id) || 1) * 2), 22);
}}
function edgeColor(type) {{
const map = {{
HAS_CONDITION: "#38bdf8",
IS_PART_OF: "#f59e0b",
TRIGGERS: "#fb923c",
CAUSES: "#fb7185",
SUGGESTS: "#34d399",
MITIGATES: "#22c55e",
INFLUENCES: "#06b6d4",
PROVIDES: "#84cc16",
DEFINES: "#c084fc",
REFERS_TO: "#94a3b8"
}};
return map[type] || "#64748b";
}}
function renderFilters() {{
const labels = [...new Set(graph.nodes.map(n => n.label))].sort();
labelFilters.innerHTML = labels.map(label => `
<label class="filter-row">
<input type="checkbox" data-label="${{label}}" checked />
<span class="swatch" style="background:${{colors[label] || "#64748b"}}"></span>
<span>${{label}}</span>
<span class="pill">${{graph.stats.label_counts[label] || 0}}</span>
</label>`).join("");
labelFilters.querySelectorAll("input").forEach(input => {{
input.addEventListener("change", () => {{
input.checked ? activeLabels.add(input.dataset.label) : activeLabels.delete(input.dataset.label);
draw();
}});
}});
const edgeTypes = [...new Set(graph.edges.map(e => e.type))].sort();
edgeFilters.innerHTML = edgeTypes.map(type => `
<label class="filter-row">
<input type="checkbox" data-edge="${{type}}" checked />
<span class="swatch" style="background:${{edgeColor(type)}}"></span>
<span>${{type}}</span>
<span class="pill">${{graph.stats.edge_type_counts[type] || 0}}</span>
</label>`).join("");
edgeFilters.querySelectorAll("input").forEach(input => {{
input.addEventListener("change", () => {{
input.checked ? activeEdges.add(input.dataset.edge) : activeEdges.delete(input.dataset.edge);
draw();
}});
}});
}}
function updateDetails(item) {{
if (!item) {{
details.textContent = "No selection yet.";
return;
}}
details.textContent = JSON.stringify(item, null, 2);
}}
function nearestNode(x, y) {{
const [wx, wy] = screenToWorld(x, y);
let best = null, bestDist = Infinity;
for (const n of graph.nodes.filter(isVisibleNode)) {{
const dx = n.x - wx, dy = n.y - wy;
const dist = Math.sqrt(dx * dx + dy * dy);
if (dist < bestDist && dist < nodeRadius(n) / scale + 8) {{
best = n;
bestDist = dist;
}}
}}
return best;
}}
function nearestEdge(x, y) {{
let best = null, bestDist = Infinity;
for (const e of graph.edges.filter(isVisibleEdge)) {{
const a = nodeById.get(e.source_element_id), b = nodeById.get(e.target_element_id);
const [x1, y1] = worldToScreen(a.x, a.y);
const [x2, y2] = worldToScreen(b.x, b.y);
const dist = pointLineDistance(x, y, x1, y1, x2, y2);
if (dist < bestDist && dist < 7) {{
best = e;
bestDist = dist;
}}
}}
return best;
}}
function pointLineDistance(px, py, x1, y1, x2, y2) {{
const dx = x2 - x1, dy = y2 - y1;
const len2 = dx * dx + dy * dy || 1;
const t = Math.max(0, Math.min(1, ((px - x1) * dx + (py - y1) * dy) / len2));
const x = x1 + t * dx, y = y1 + t * dy;
return Math.hypot(px - x, py - y);
}}
canvas.addEventListener("mousemove", e => {{
const rect = canvas.getBoundingClientRect();
const x = e.clientX - rect.left, y = e.clientY - rect.top;
if (dragging) {{
const [wx, wy] = screenToWorld(x, y);
dragging.x = wx;
dragging.y = wy;
dragging.vx = 0;
dragging.vy = 0;
draw();
return;
}}
if (panning) {{
offsetX += x - lastX;
offsetY += y - lastY;
lastX = x;
lastY = y;
draw();
return;
}}
hover = nearestNode(x, y);
draw();
}});
canvas.addEventListener("mousedown", e => {{
const rect = canvas.getBoundingClientRect();
const x = e.clientX - rect.left, y = e.clientY - rect.top;
dragging = nearestNode(x, y);
if (dragging) {{
dragging.fixed = true;
}} else {{
panning = true;
lastX = x;
lastY = y;
}}
}});
window.addEventListener("mouseup", () => {{
dragging = null;
panning = false;
}});
canvas.addEventListener("click", e => {{
const rect = canvas.getBoundingClientRect();
const x = e.clientX - rect.left, y = e.clientY - rect.top;
selected = nearestNode(x, y) || nearestEdge(x, y);
updateDetails(selected);
draw();
}});
canvas.addEventListener("wheel", e => {{
e.preventDefault();
const before = scale;
scale *= e.deltaY < 0 ? 1.1 : 0.9;
scale = Math.max(0.15, Math.min(scale, 4));
const rect = canvas.getBoundingClientRect();
const x = e.clientX - rect.left - canvas.clientWidth / 2;
const y = e.clientY - rect.top - canvas.clientHeight / 2;
offsetX = x - (x - offsetX) * (scale / before);
offsetY = y - (y - offsetY) * (scale / before);
draw();
}}, {{ passive: false }});
document.getElementById("reset").onclick = () => {{ scale = 1; offsetX = 0; offsetY = 0; selected = null; updateDetails(null); }};
document.getElementById("labels").onclick = () => {{ showEdgeLabels = !showEdgeLabels; draw(); }};
document.getElementById("freeze").onclick = e => {{ paused = !paused; e.target.textContent = paused ? "Resume" : "Pause"; }};
document.getElementById("png").onclick = () => {{
const link = document.createElement("a");
link.download = "neo4j_graph.png";
link.href = canvas.toDataURL("image/png");
link.click();
}};
search.addEventListener("input", draw);
function renderTopNodes() {{
const top = [...graph.nodes].sort((a, b) => (degree.get(b.element_id) || 0) - (degree.get(a.element_id) || 0)).slice(0, 12);
document.getElementById("topNodes").innerHTML = top.map(n => `
<div class="mini"><b>${{n.name || n.display_id}}</b><div>${{n.label}} · degree ${{degree.get(n.element_id) || 0}}</div></div>
`).join("");
}}
function loop() {{
if (!paused) tick();
draw();
requestAnimationFrame(loop);
}}
window.addEventListener("resize", resize);
resize();
renderFilters();
renderTopNodes();
loop();
</script>
</body>
</html>"""
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()