cognee_graph/
formatted.rs

1//! Formatted-graph-data helper — port of Python
2//! `cognee.modules.graph.methods.get_formatted_graph_data`.
3//!
4//! Reads all nodes and edges from a graph DB and reshapes them into the
5//! canonical wire payload:
6//!
7//! ```json
8//! {
9//!   "nodes": [{"id": "...", "label": "...", "type": "...", "properties": {...}}],
10//!   "edges": [{"source": "...", "target": "...", "label": "..."}]
11//! }
12//! ```
13//!
14//! Field ordering matches the Python helper exactly — the WS frame is part
15//! of the cross-SDK wire contract, see [`docs/http-server/websocket.md`].
16
17use uuid::Uuid;
18
19use crate::{GraphDBResult, GraphDBTrait};
20
21/// Fetch the formatted graph snapshot for a dataset.
22///
23/// # Python parity
24///
25/// The Rust port intentionally omits Python's `set_database_global_context_variables`
26/// branch — the Rust `GraphDBTrait` instance is already scoped to the caller's
27/// owner/tenant by construction (per `tenants.md §3`). The `dataset_id` and
28/// `user_id` parameters are accepted for API parity with the Python helper but
29/// are currently unused in the read path (matching how Python's helper also
30/// relies on the global context, not the parameters, to scope the query).
31///
32/// # Shape
33///
34/// Each node produces:
35/// - `id`     — string form of the node id
36/// - `label`  — `properties["name"]` if non-empty, else `"{type}_{id}"`
37/// - `type`   — `properties["type"]`
38/// - `properties` — every other property whose value is not null, with
39///   `id`, `type`, `name`, `created_at`, `updated_at` excluded.
40///
41/// Each edge produces:
42/// - `source`, `target`, `label` (the relationship name).
43pub async fn get_formatted_graph_data(
44    graph_db: &dyn GraphDBTrait,
45    dataset_id: Uuid,
46    user_id: Uuid,
47) -> GraphDBResult<serde_json::Value> {
48    let _ = (dataset_id, user_id);
49
50    let (nodes, edges) = graph_db.get_graph_data().await?;
51
52    let node_values: Vec<serde_json::Value> = nodes
53        .into_iter()
54        .map(|(node_id, props)| format_node(&node_id, &props))
55        .collect();
56
57    let edge_values: Vec<serde_json::Value> = edges
58        .into_iter()
59        .map(|(source, target, relationship_name, _props)| {
60            serde_json::json!({
61                "source": source,
62                "target": target,
63                "label": relationship_name,
64            })
65        })
66        .collect();
67
68    Ok(serde_json::json!({
69        "nodes": node_values,
70        "edges": edge_values,
71    }))
72}
73
74/// Build the per-node object matching Python's mapping shape.
75///
76/// `label = properties["name"]` if non-empty, else `"{type}_{id}"`.
77/// `properties` excludes id, type, name, created_at, updated_at and drops
78/// any value that is `null`.
79fn format_node(node_id: &str, props: &crate::NodeData) -> serde_json::Value {
80    let type_str = props
81        .get("type")
82        .and_then(|v| v.as_str())
83        .unwrap_or("")
84        .to_string();
85
86    let name = props.get("name").and_then(|v| v.as_str()).unwrap_or("");
87
88    let label = if !name.is_empty() {
89        name.to_string()
90    } else {
91        format!("{type_str}_{node_id}")
92    };
93
94    let mut properties_map = serde_json::Map::new();
95    for (key, value) in props.iter() {
96        let k = key.as_ref();
97        if matches!(k, "id" | "type" | "name" | "created_at" | "updated_at") {
98            continue;
99        }
100        if value.is_null() {
101            continue;
102        }
103        properties_map.insert(k.to_string(), value.clone());
104    }
105
106    serde_json::json!({
107        "id": node_id,
108        "label": label,
109        "type": type_str,
110        "properties": serde_json::Value::Object(properties_map),
111    })
112}
cognee_graph/formatted.rs

cognee_graph/
formatted.rs
