cognee_visualization/
lib.rs

1//! Interactive HTML knowledge-graph visualization for Cognee-Rust.
2//!
3//! This crate ports the Python `cognee_network_visualization` module to Rust.
4//! It renders all nodes + edges of a `GraphDBTrait` into a single self-contained
5//! HTML file that uses d3.js v7 for force-directed layout and Canvas rendering.
6//!
7//! # Quick start
8//!
9//! ```no_run
10//! use cognee_graph::GraphDBTrait;
11//! use cognee_visualization::visualize;
12//! use std::path::Path;
13//!
14//! # async fn example(graph_db: &dyn GraphDBTrait) -> Result<(), Box<dyn std::error::Error>> {
15//! // Write the visualization to a caller-specified file.
16//! let _path = visualize(graph_db, Some(Path::new("/tmp/graph.html"))).await?;
17//!
18//! // Or write it to `~/graph_visualization.html` (matches Python behavior).
19//! let _path = visualize(graph_db, None).await?;
20//! # Ok(()) }
21//! ```
22
23mod colors;
24mod error;
25mod html;
26mod paths;
27mod serialize;
28
29pub use error::VisualizationError;
30
31use std::path::{Path, PathBuf};
32
33use cognee_graph::GraphDBTrait;
34use tokio::fs;
35use tokio::io::AsyncWriteExt;
36use tracing::info;
37
38/// Generate an interactive HTML knowledge-graph visualization of the supplied
39/// graph database.
40///
41/// * `graph_db` — graph database from which all nodes + edges are fetched via
42///   `get_graph_data()`.
43/// * `output_path` — optional destination path. When `None`, the file is
44///   written to `~/graph_visualization.html` (or `%USERPROFILE%` on Windows),
45///   matching Python's `visualize_graph()`.
46///
47/// Returns the absolute path the file was written to.
48pub async fn visualize(
49    graph_db: &dyn GraphDBTrait,
50    output_path: Option<&Path>,
51) -> Result<PathBuf, VisualizationError> {
52    let html = render(graph_db).await?;
53
54    let dest: PathBuf = match output_path {
55        Some(p) => p.to_path_buf(),
56        None => paths::default_output_path()?,
57    };
58
59    if let Some(parent) = dest.parent()
60        && !parent.as_os_str().is_empty()
61    {
62        fs::create_dir_all(parent).await?;
63    }
64    let mut file = fs::File::create(&dest).await?;
65    file.write_all(html.as_bytes()).await?;
66    file.flush().await?;
67
68    info!(path = %dest.display(), "Graph visualization saved");
69    Ok(dest)
70}
71
72/// Render the HTML visualization string for the supplied graph database,
73/// without writing it anywhere.
74///
75/// Useful when the caller wants to stream the HTML over HTTP, embed it into a
76/// larger page, or post-process it before persisting.
77pub async fn render(graph_db: &dyn GraphDBTrait) -> Result<String, VisualizationError> {
78    let (nodes, edges) = graph_db.get_graph_data().await?;
79    let serialized = serialize::serialize_graph(nodes, edges);
80    html::build_html(&serialized, None)
81}
82
83/// Render a combined HTML visualization aggregating multiple `(user_label, graph_db)`
84/// pairs into one output.
85///
86/// Each pair's nodes are tagged with a `source_user` attribute carrying the
87/// supplied human-readable label so the d3 template can color-code by user.
88/// Mirrors Python's `aggregate_multi_user_graphs()` in
89/// [`cognee/modules/visualization/cognee_network_visualization.py:115-157`](https://github.com/topoteretes/cognee/blob/main/cognee/modules/visualization/cognee_network_visualization.py#L115-L157):
90///
91/// - Nodes are deduplicated by `str(node_id)` with **first-write-wins**
92///   semantics so iteration order across the supplied pairs determines the
93///   surviving node entry.
94/// - Edges are deduplicated by the `(source, target, relation)` tuple.
95/// - The `source_user` field is only populated when the inbound node does not
96///   already carry one (Python's `if not node_info.get("source_user")`).
97///
98/// An empty input produces a valid-but-empty HTML document.
99///
100/// `pairs` is a slice of `(user_label, graph_db)` tuples; the label is taken
101/// as an arbitrary `&str` to keep this crate decoupled from
102/// `cognee_models::User`. Callers should resolve the underlying user record
103/// and pass `user.email` (or the stringified id as a fallback) so the
104/// `userColors` palette key matches Python.
105pub async fn render_multi_user(
106    pairs: &[(String, std::sync::Arc<dyn GraphDBTrait>)],
107) -> Result<String, VisualizationError> {
108    use std::borrow::Cow;
109    use std::collections::{HashMap, HashSet};
110
111    // First-write-wins by stringified node id (mirror Python L142).
112    let mut all_nodes: HashMap<String, cognee_graph::GraphNode> = HashMap::new();
113    let mut node_order: Vec<String> = Vec::new();
114
115    // Edge dedupe by (source, target, relation) (mirror Python L150-155).
116    let mut all_edges: Vec<cognee_graph::EdgeData> = Vec::new();
117    let mut seen_edges: HashSet<(String, String, String)> = HashSet::new();
118
119    for (user_label, gdb) in pairs {
120        let (nodes, edges) = gdb.get_graph_data().await?;
121        for (node_id, mut node_info) in nodes {
122            let key = node_id.to_string();
123            if all_nodes.contains_key(&key) {
124                continue;
125            }
126            // Mirror Python's `if not node_info.get("source_user"): node_info["source_user"] = user_label`
127            // — preserve any pre-existing `source_user` on the node so the
128            // owning user's value (if already tagged) survives.
129            let needs_label = match node_info.get("source_user") {
130                Some(serde_json::Value::String(s)) if !s.is_empty() => false,
131                Some(serde_json::Value::Null) | None => true,
132                _ => false,
133            };
134            if needs_label {
135                node_info.insert(
136                    Cow::Borrowed("source_user"),
137                    serde_json::Value::String(user_label.clone()),
138                );
139            }
140            node_order.push(key.clone());
141            all_nodes.insert(key, (node_id, node_info));
142        }
143        for edge in edges {
144            let edge_key = (edge.0.to_string(), edge.1.to_string(), edge.2.clone());
145            if seen_edges.insert(edge_key) {
146                all_edges.push(edge);
147            }
148        }
149    }
150
151    let ordered_nodes: Vec<cognee_graph::GraphNode> = node_order
152        .into_iter()
153        .filter_map(|k| all_nodes.remove(&k))
154        .collect();
155
156    let serialized = serialize::serialize_graph(ordered_nodes, all_edges);
157    html::build_html(&serialized, None)
158}
cognee_visualization/lib.rs

cognee_visualization/
lib.rs
