cognee_models/
edge_metadata.rs

1//! EdgeMetadata - Metadata for relationships between DataPoints.
2//!
3//! Mirrors Python's `cognee/infrastructure/engine/models/Edge.py`.
4//! Represents edge properties like weight, relationship type, and edge text
5//! for relationships between DataPoints.
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9
10/// Edge metadata for relationships between DataPoints.
11///
12/// Matches the Python `Edge(BaseModel)` class. Supports single weight,
13/// multiple named weights, arbitrary properties, and an `edge_text` field
14/// that is auto-populated from `relationship_type` when not explicitly set
15/// (mirroring Python's `field_validator("edge_text")` behavior).
16///
17/// # Examples
18///
19/// ```
20/// use cognee_models::EdgeMetadata;
21///
22/// // Auto-populates edge_text from relationship_type
23/// let edge = EdgeMetadata::new(Some("contains".into()), None, None);
24/// assert_eq!(edge.edge_text.as_deref(), Some("contains"));
25///
26/// // Explicit edge_text takes priority
27/// let edge = EdgeMetadata::new(
28///     Some("contains".into()),
29///     Some(0.5),
30///     Some("relationship_name: contains; entity: Alice".into()),
31/// );
32/// assert_eq!(edge.edge_text.as_deref(), Some("relationship_name: contains; entity: Alice"));
33/// ```
34#[derive(Debug, Clone, Default, Serialize, PartialEq)]
35pub struct EdgeMetadata {
36    /// Relationship type name (e.g., "works_at", "contains").
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub relationship_type: Option<String>,
39
40    /// Single weight value (backward compatible with Python's `weight` field).
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub weight: Option<f64>,
43
44    /// Multiple named weights (e.g., `{"strength": 0.8, "confidence": 0.9}`).
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub weights: Option<HashMap<String, f64>>,
47
48    /// Arbitrary edge properties (flexible key-value storage).
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub properties: Option<HashMap<String, serde_json::Value>>,
51
52    /// Text representation for embedding. Auto-populated from `relationship_type`
53    /// if not explicitly set (matches Python's `field_validator` behavior).
54    #[serde(skip_serializing_if = "Option::is_none")]
55    pub edge_text: Option<String>,
56}
57
58impl EdgeMetadata {
59    /// Create a new EdgeMetadata, auto-populating `edge_text` from `relationship_type`
60    /// if `edge_text` is not provided (matches Python's `field_validator` behavior).
61    pub fn new(
62        relationship_type: Option<String>,
63        weight: Option<f64>,
64        edge_text: Option<String>,
65    ) -> Self {
66        let effective_edge_text = edge_text.or_else(|| relationship_type.clone());
67        Self {
68            relationship_type,
69            weight,
70            weights: None,
71            properties: None,
72            edge_text: effective_edge_text,
73        }
74    }
75
76    /// Create an EdgeMetadata with all fields specified.
77    pub fn with_all(
78        relationship_type: Option<String>,
79        weight: Option<f64>,
80        weights: Option<HashMap<String, f64>>,
81        properties: Option<HashMap<String, serde_json::Value>>,
82        edge_text: Option<String>,
83    ) -> Self {
84        let effective_edge_text = edge_text.or_else(|| relationship_type.clone());
85        Self {
86            relationship_type,
87            weight,
88            weights,
89            properties,
90            edge_text: effective_edge_text,
91        }
92    }
93
94    /// Auto-populate `edge_text` from `relationship_type` if `edge_text` is `None`.
95    ///
96    /// This mirrors Python's `field_validator("edge_text")` which sets
97    /// `edge_text = relationship_type` when `edge_text` is not provided.
98    pub fn ensure_edge_text(&mut self) {
99        if self.edge_text.is_none() {
100            self.edge_text.clone_from(&self.relationship_type);
101        }
102    }
103}
104
105/// Custom `Deserialize` implementation that auto-populates `edge_text` from
106/// `relationship_type` after deserialization, matching Python's `field_validator`
107/// behavior. This ensures that JSON like `{"relationship_type": "contains"}`
108/// will produce `edge_text == Some("contains")`.
109impl<'de> Deserialize<'de> for EdgeMetadata {
110    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
111    where
112        D: serde::Deserializer<'de>,
113    {
114        /// Helper struct with standard derive for deserialization.
115        #[derive(Deserialize)]
116        struct EdgeMetadataRaw {
117            relationship_type: Option<String>,
118            weight: Option<f64>,
119            weights: Option<HashMap<String, f64>>,
120            properties: Option<HashMap<String, serde_json::Value>>,
121            edge_text: Option<String>,
122        }
123
124        let raw = EdgeMetadataRaw::deserialize(deserializer)?;
125
126        // Auto-populate edge_text from relationship_type (Python field_validator)
127        let effective_edge_text = raw.edge_text.or_else(|| raw.relationship_type.clone());
128
129        Ok(EdgeMetadata {
130            relationship_type: raw.relationship_type,
131            weight: raw.weight,
132            weights: raw.weights,
133            properties: raw.properties,
134            edge_text: effective_edge_text,
135        })
136    }
137}
138
139#[cfg(test)]
140#[allow(
141    clippy::unwrap_used,
142    clippy::expect_used,
143    reason = "test code — panics are acceptable failures"
144)]
145mod tests {
146    use super::*;
147    use serde_json::json;
148
149    #[test]
150    fn test_new_auto_populates_edge_text() {
151        let edge = EdgeMetadata::new(Some("contains".into()), None, None);
152        assert_eq!(edge.relationship_type.as_deref(), Some("contains"));
153        assert_eq!(edge.edge_text.as_deref(), Some("contains"));
154    }
155
156    #[test]
157    fn test_new_explicit_edge_text_preserved() {
158        let edge = EdgeMetadata::new(
159            Some("contains".into()),
160            Some(0.5),
161            Some("custom text".into()),
162        );
163        assert_eq!(edge.relationship_type.as_deref(), Some("contains"));
164        assert_eq!(edge.weight, Some(0.5));
165        assert_eq!(edge.edge_text.as_deref(), Some("custom text"));
166    }
167
168    #[test]
169    fn test_new_no_relationship_type_no_edge_text() {
170        let edge = EdgeMetadata::new(None, Some(1.0), None);
171        assert_eq!(edge.relationship_type, None);
172        assert_eq!(edge.edge_text, None);
173        assert_eq!(edge.weight, Some(1.0));
174    }
175
176    #[test]
177    fn test_ensure_edge_text_populates_when_none() {
178        let mut edge = EdgeMetadata {
179            relationship_type: Some("works_at".into()),
180            edge_text: None,
181            ..Default::default()
182        };
183        edge.ensure_edge_text();
184        assert_eq!(edge.edge_text.as_deref(), Some("works_at"));
185    }
186
187    #[test]
188    fn test_ensure_edge_text_preserves_existing() {
189        let mut edge = EdgeMetadata {
190            relationship_type: Some("works_at".into()),
191            edge_text: Some("already set".into()),
192            ..Default::default()
193        };
194        edge.ensure_edge_text();
195        assert_eq!(edge.edge_text.as_deref(), Some("already set"));
196    }
197
198    #[test]
199    fn test_ensure_edge_text_both_none() {
200        let mut edge = EdgeMetadata::default();
201        edge.ensure_edge_text();
202        assert_eq!(edge.edge_text, None);
203    }
204
205    #[test]
206    fn test_default_all_none() {
207        let edge = EdgeMetadata::default();
208        assert_eq!(edge.relationship_type, None);
209        assert_eq!(edge.weight, None);
210        assert_eq!(edge.weights, None);
211        assert_eq!(edge.properties, None);
212        assert_eq!(edge.edge_text, None);
213    }
214
215    #[test]
216    fn test_with_all_fields() {
217        let mut weights = HashMap::new();
218        weights.insert("strength".into(), 0.8);
219        weights.insert("confidence".into(), 0.9);
220
221        let mut properties = HashMap::new();
222        properties.insert("source".into(), json!("manual"));
223
224        let edge = EdgeMetadata::with_all(
225            Some("contains".into()),
226            Some(0.5),
227            Some(weights.clone()),
228            Some(properties.clone()),
229            Some("custom text".into()),
230        );
231
232        assert_eq!(edge.relationship_type.as_deref(), Some("contains"));
233        assert_eq!(edge.weight, Some(0.5));
234        assert_eq!(edge.weights.as_ref().unwrap().get("strength"), Some(&0.8));
235        assert_eq!(edge.weights.as_ref().unwrap().get("confidence"), Some(&0.9));
236        assert_eq!(
237            edge.properties.as_ref().unwrap().get("source"),
238            Some(&json!("manual"))
239        );
240        assert_eq!(edge.edge_text.as_deref(), Some("custom text"));
241    }
242
243    #[test]
244    fn test_with_all_auto_populates_edge_text() {
245        let edge = EdgeMetadata::with_all(
246            Some("located_in".into()),
247            None,
248            None,
249            None,
250            None, // edge_text not provided
251        );
252        assert_eq!(edge.edge_text.as_deref(), Some("located_in"));
253    }
254
255    #[test]
256    fn test_serialization_roundtrip() {
257        let edge = EdgeMetadata::new(Some("works_at".into()), Some(0.75), None);
258        let json = serde_json::to_string(&edge).unwrap();
259        let deserialized: EdgeMetadata = serde_json::from_str(&json).unwrap();
260
261        assert_eq!(edge, deserialized);
262    }
263
264    #[test]
265    fn test_serialization_skips_none_fields() {
266        let edge = EdgeMetadata::new(Some("works_at".into()), None, None);
267        let json_value: serde_json::Value = serde_json::to_value(&edge).unwrap();
268        let obj = json_value.as_object().unwrap();
269
270        assert!(obj.contains_key("relationship_type"));
271        assert!(obj.contains_key("edge_text"));
272        assert!(!obj.contains_key("weight"));
273        assert!(!obj.contains_key("weights"));
274        assert!(!obj.contains_key("properties"));
275    }
276
277    #[test]
278    fn test_deserialize_auto_populates_edge_text() {
279        // Deserializing JSON without edge_text should auto-populate from relationship_type
280        let json = r#"{"relationship_type": "contains"}"#;
281        let edge: EdgeMetadata = serde_json::from_str(json).unwrap();
282
283        assert_eq!(edge.relationship_type.as_deref(), Some("contains"));
284        assert_eq!(edge.edge_text.as_deref(), Some("contains"));
285    }
286
287    #[test]
288    fn test_deserialize_explicit_edge_text_preserved() {
289        let json = r#"{"relationship_type": "contains", "edge_text": "custom"}"#;
290        let edge: EdgeMetadata = serde_json::from_str(json).unwrap();
291
292        assert_eq!(edge.relationship_type.as_deref(), Some("contains"));
293        assert_eq!(edge.edge_text.as_deref(), Some("custom"));
294    }
295
296    #[test]
297    fn test_deserialize_empty_json() {
298        let json = r#"{}"#;
299        let edge: EdgeMetadata = serde_json::from_str(json).unwrap();
300
301        assert_eq!(edge, EdgeMetadata::default());
302    }
303
304    #[test]
305    fn test_deserialize_with_weights() {
306        let json = r#"{
307            "relationship_type": "contains",
308            "weight": 0.5,
309            "weights": {"strength": 0.8, "confidence": 0.9}
310        }"#;
311        let edge: EdgeMetadata = serde_json::from_str(json).unwrap();
312
313        assert_eq!(edge.weight, Some(0.5));
314        let weights = edge.weights.as_ref().unwrap();
315        assert_eq!(weights.get("strength"), Some(&0.8));
316        assert_eq!(weights.get("confidence"), Some(&0.9));
317        // edge_text auto-populated from relationship_type
318        assert_eq!(edge.edge_text.as_deref(), Some("contains"));
319    }
320
321    #[test]
322    fn test_deserialize_with_properties() {
323        let json = r#"{
324            "relationship_type": "works_at",
325            "properties": {"since": "2020", "role": "engineer", "active": true}
326        }"#;
327        let edge: EdgeMetadata = serde_json::from_str(json).unwrap();
328
329        let props = edge.properties.as_ref().unwrap();
330        assert_eq!(props.get("since"), Some(&json!("2020")));
331        assert_eq!(props.get("role"), Some(&json!("engineer")));
332        assert_eq!(props.get("active"), Some(&json!(true)));
333    }
334
335    #[test]
336    fn test_clone() {
337        let edge = EdgeMetadata::new(Some("contains".into()), Some(0.5), None);
338        let cloned = edge.clone();
339        assert_eq!(edge, cloned);
340    }
341
342    #[test]
343    fn test_debug_format() {
344        let edge = EdgeMetadata::new(Some("contains".into()), None, None);
345        let debug = format!("{edge:?}");
346        assert!(debug.contains("EdgeMetadata"));
347        assert!(debug.contains("contains"));
348    }
349}
cognee_models/edge_metadata.rs

cognee_models/
edge_metadata.rs
