cognee_models/
data_point.rs

1//! DataPoint - Base model for all storage-layer entities.
2//!
3//! Mirrors Python's `cognee/infrastructure/engine/models/DataPoint.py`
4//! Provides common fields for UUID, timestamps, versioning, and metadata.
5
6use chrono::Utc;
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use uuid::Uuid;
10
11/// Default value for `feedback_weight` (used by serde).
12fn default_feedback_weight() -> f64 {
13    0.5
14}
15
16/// Default value for `importance_weight` (used by serde).
17///
18/// Mirrors Python's `DataPoint.importance_weight: float | None = 0.5`.
19/// Using a named serde default (rather than bare `#[serde(default)]`)
20/// ensures old payloads missing the key deserialize to `Some(0.5)` — the
21/// neutral rank — rather than `None`.
22fn default_importance_weight() -> Option<f64> {
23    Some(0.5)
24}
25
26/// Default value for `version` (used by serde).
27fn default_version() -> i32 {
28    1
29}
30
31/// Base model for all storage-layer entities.
32///
33/// Provides:
34/// - Unique identifier (UUID)
35/// - Timestamps (created_at, updated_at) as milliseconds since epoch
36/// - Ontology validation flag
37/// - Version tracking (integer)
38/// - Topological rank for graph traversal
39/// - Flexible metadata storage
40/// - Type discriminator
41/// - Dataset membership
42/// - Pipeline provenance fields
43/// - Feedback weight
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
45pub struct DataPoint {
46    /// Unique identifier
47    pub id: Uuid,
48
49    /// Creation timestamp (milliseconds since epoch, matching Python)
50    pub created_at: i64,
51
52    /// Last update timestamp (milliseconds since epoch, matching Python)
53    pub updated_at: i64,
54
55    /// Whether this entity has been validated against an ontology
56    pub ontology_valid: bool,
57
58    /// Version number (default 1, matching Python)
59    #[serde(default = "default_version")]
60    pub version: i32,
61
62    /// Topological rank for graph traversal optimization
63    pub topological_rank: Option<i32>,
64
65    /// Flexible metadata storage (e.g., index_fields, custom attributes)
66    pub metadata: HashMap<String, serde_json::Value>,
67
68    /// Type discriminator (e.g., "Entity", "EntityType", "EdgeType")
69    #[serde(rename = "type")]
70    pub data_type: String,
71
72    /// Dataset this data point belongs to (list of JSON values, matching Python)
73    pub belongs_to_set: Option<Vec<serde_json::Value>>,
74
75    /// Pipeline that created this data point
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub source_pipeline: Option<String>,
78
79    /// Task that created this data point
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub source_task: Option<String>,
82
83    /// Node set source
84    #[serde(default, skip_serializing_if = "Option::is_none")]
85    pub source_node_set: Option<String>,
86
87    /// User that triggered creation
88    #[serde(default, skip_serializing_if = "Option::is_none")]
89    pub source_user: Option<String>,
90
91    /// Content hash of the raw `Data` artefact that produced this DataPoint.
92    /// Propagates from upstream `Data.content_hash` through every task in
93    /// the cognify pipeline, enabling content-addressed lineage queries.
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub source_content_hash: Option<String>,
96
97    /// Feedback weight (default 0.5, matching Python)
98    #[serde(default = "default_feedback_weight")]
99    pub feedback_weight: f64,
100
101    /// Importance weight (default `Some(0.5)`, matching Python
102    /// `importance_weight: float | None = 0.5`). Propagated from the source
103    /// `Data` record through classify → chunk → summarize. Emitted
104    /// unconditionally (no `skip_serializing_if`) since it is effectively
105    /// always present.
106    #[serde(default = "default_importance_weight")]
107    pub importance_weight: Option<f64>,
108}
109
110impl DataPoint {
111    /// Create a new DataPoint with default values.
112    ///
113    /// # Arguments
114    /// * `data_type` - Type discriminator (e.g., "Entity", "EntityType")
115    /// * `dataset_id` - Optional dataset UUID
116    pub fn new(data_type: impl Into<String>, dataset_id: Option<Uuid>) -> Self {
117        let now = Utc::now().timestamp_millis();
118        Self {
119            id: Uuid::new_v4(),
120            created_at: now,
121            updated_at: now,
122            ontology_valid: false,
123            version: 1,
124            topological_rank: None,
125            metadata: HashMap::new(),
126            data_type: data_type.into(),
127            belongs_to_set: dataset_id.map(|id| vec![serde_json::json!(id.to_string())]),
128            source_pipeline: None,
129            source_task: None,
130            source_node_set: None,
131            source_user: None,
132            source_content_hash: None,
133            feedback_weight: 0.5,
134            importance_weight: Some(0.5),
135        }
136    }
137
138    /// Create a DataPoint with specific metadata.
139    pub fn with_metadata(
140        data_type: impl Into<String>,
141        dataset_id: Option<Uuid>,
142        metadata: HashMap<String, serde_json::Value>,
143    ) -> Self {
144        let now = Utc::now().timestamp_millis();
145        Self {
146            id: Uuid::new_v4(),
147            created_at: now,
148            updated_at: now,
149            ontology_valid: false,
150            version: 1,
151            topological_rank: None,
152            metadata,
153            data_type: data_type.into(),
154            belongs_to_set: dataset_id.map(|id| vec![serde_json::json!(id.to_string())]),
155            source_pipeline: None,
156            source_task: None,
157            source_node_set: None,
158            source_user: None,
159            source_content_hash: None,
160            feedback_weight: 0.5,
161            importance_weight: Some(0.5),
162        }
163    }
164
165    /// Get embeddable data as JSON string for vector indexing.
166    ///
167    /// Returns a JSON representation of this DataPoint.
168    pub fn get_embeddable_data(&self) -> String {
169        serde_json::to_string(self).unwrap_or_else(|_| "{}".to_string())
170    }
171
172    /// Convert to JSON value.
173    pub fn to_json(&self) -> serde_json::Value {
174        serde_json::to_value(self).unwrap_or(serde_json::Value::Null)
175    }
176
177    /// Canonical vector-store payload keys for this DataPoint.
178    ///
179    /// Mirrors Python's `DataPoint.model_dump()` payload shape: every
180    /// pydantic-equivalent field flows into the metadata map. Keys with
181    /// `None` values are omitted (consistent with the
182    /// `skip_serializing_if = "Option::is_none"` annotations on the
183    /// struct).
184    ///
185    /// Used by the cognify and memify pipelines when constructing
186    /// `VectorPoint` payloads to keep the Rust shape byte-comparable to
187    /// Python's for the cross-SDK parity tests. Note: the `data_type`
188    /// field carries `#[serde(rename = "type")]`, so the resulting map
189    /// uses the JSON key `"type"` (matching Python).
190    pub fn vector_metadata(&self) -> HashMap<String, serde_json::Value> {
191        match serde_json::to_value(self) {
192            Ok(serde_json::Value::Object(map)) => map.into_iter().collect(),
193            _ => HashMap::new(),
194        }
195    }
196
197    /// Update the timestamp to current time.
198    pub fn touch(&mut self) {
199        self.updated_at = Utc::now().timestamp_millis();
200    }
201
202    /// Set ontology validation status.
203    pub fn set_ontology_valid(&mut self, valid: bool) {
204        self.ontology_valid = valid;
205        self.touch();
206    }
207
208    /// Add or update metadata field.
209    pub fn set_metadata(&mut self, key: impl Into<String>, value: serde_json::Value) {
210        self.metadata.insert(key.into(), value);
211        self.touch();
212    }
213
214    /// Get metadata field.
215    pub fn get_metadata(&self, key: &str) -> Option<&serde_json::Value> {
216        self.metadata.get(key)
217    }
218}
219
220#[cfg(test)]
221#[allow(
222    clippy::unwrap_used,
223    clippy::expect_used,
224    reason = "test code — panics are acceptable failures"
225)]
226mod tests {
227    use super::*;
228    use serde_json::json;
229
230    #[test]
231    fn test_data_point_creation() {
232        let dp = DataPoint::new("TestType", None);
233        assert_eq!(dp.data_type, "TestType");
234        assert_eq!(dp.version, 1);
235        assert!(!dp.ontology_valid);
236        assert!(dp.metadata.is_empty());
237        assert!(dp.belongs_to_set.is_none());
238        assert!(dp.source_pipeline.is_none());
239        assert!(dp.source_task.is_none());
240        assert!(dp.source_node_set.is_none());
241        assert!(dp.source_user.is_none());
242        assert!(dp.source_content_hash.is_none());
243        assert!((dp.feedback_weight - 0.5).abs() < f64::EPSILON);
244        assert_eq!(dp.importance_weight, Some(0.5));
245        assert!(dp.created_at > 0);
246        assert!(dp.updated_at > 0);
247    }
248
249    #[test]
250    fn with_metadata_initializes_importance_weight_default() {
251        let dp = DataPoint::with_metadata("Entity", None, HashMap::new());
252        assert_eq!(dp.importance_weight, Some(0.5));
253    }
254
255    #[test]
256    fn test_data_point_importance_weight_default_on_missing_field() {
257        // Simulate an old stored payload lacking the `importance_weight` key.
258        let json = json!({
259            "id": Uuid::new_v4().to_string(),
260            "created_at": 1_i64,
261            "updated_at": 1_i64,
262            "ontology_valid": false,
263            "metadata": {},
264            "type": "Entity",
265            "belongs_to_set": null,
266        });
267        let dp: DataPoint = serde_json::from_value(json).unwrap();
268        assert_eq!(dp.importance_weight, Some(0.5));
269    }
270
271    #[test]
272    fn importance_weight_round_trips_with_expected_key() {
273        let mut dp = DataPoint::new("Entity", None);
274        dp.importance_weight = Some(0.9);
275        let json = serde_json::to_string(&dp).unwrap();
276        assert!(json.contains(r#""importance_weight":0.9"#));
277
278        let parsed: DataPoint = serde_json::from_str(&json).unwrap();
279        assert_eq!(parsed.importance_weight, Some(0.9));
280    }
281
282    #[test]
283    fn test_data_point_with_dataset() {
284        let dataset_id = Uuid::new_v4();
285        let dp = DataPoint::new("Entity", Some(dataset_id));
286        assert_eq!(
287            dp.belongs_to_set,
288            Some(vec![serde_json::json!(dataset_id.to_string())])
289        );
290    }
291
292    #[test]
293    fn test_metadata_operations() {
294        let mut dp = DataPoint::new("Entity", None);
295        dp.set_metadata("index_fields", serde_json::json!(["name"]));
296
297        assert_eq!(
298            dp.get_metadata("index_fields"),
299            Some(&serde_json::json!(["name"]))
300        );
301    }
302
303    #[test]
304    fn test_ontology_validation() {
305        let mut dp = DataPoint::new("Entity", None);
306        assert!(!dp.ontology_valid);
307
308        dp.set_ontology_valid(true);
309        assert!(dp.ontology_valid);
310    }
311
312    #[test]
313    fn test_get_embeddable_data() {
314        let dp = DataPoint::new("Entity", None);
315        let json_str = dp.get_embeddable_data();
316        assert!(json_str.contains("\"type\":\"Entity\""));
317    }
318
319    #[test]
320    fn source_content_hash_round_trips_when_set_and_omitted_when_none() {
321        let mut dp = DataPoint::new("Entity", None);
322        assert!(
323            !serde_json::to_string(&dp)
324                .unwrap()
325                .contains("source_content_hash"),
326            "absent field must be skipped by serde"
327        );
328
329        dp.source_content_hash = Some("md5:abcdef".to_string());
330        let json = serde_json::to_string(&dp).unwrap();
331        assert!(json.contains(r#""source_content_hash":"md5:abcdef""#));
332
333        let parsed: DataPoint = serde_json::from_str(&json).unwrap();
334        assert_eq!(parsed.source_content_hash.as_deref(), Some("md5:abcdef"));
335    }
336
337    #[test]
338    fn vector_metadata_includes_all_set_source_fields() {
339        let mut dp = DataPoint::new("Entity", None);
340        dp.source_pipeline = Some("cognify_pipeline".into());
341        dp.source_task = Some("classify_documents".into());
342        dp.source_user = Some("alice@example.com".into());
343        dp.source_node_set = Some("entity_nodes".into());
344        dp.source_content_hash = Some("md5:abcdef".into());
345
346        let m = dp.vector_metadata();
347        assert_eq!(
348            m.get("source_pipeline").unwrap(),
349            &json!("cognify_pipeline")
350        );
351        assert_eq!(m.get("source_task").unwrap(), &json!("classify_documents"));
352        assert_eq!(m.get("source_user").unwrap(), &json!("alice@example.com"));
353        assert_eq!(m.get("source_node_set").unwrap(), &json!("entity_nodes"));
354        assert_eq!(m.get("source_content_hash").unwrap(), &json!("md5:abcdef"));
355        // `data_type` round-trips as the JSON key `"type"` because of
356        // `#[serde(rename = "type")]` on the struct field.
357        assert_eq!(m.get("type").unwrap(), &json!("Entity"));
358        assert_eq!(m.get("version").unwrap(), &json!(1));
359        assert!(m.contains_key("created_at"));
360        assert!(m.contains_key("updated_at"));
361    }
362
363    #[test]
364    fn vector_metadata_omits_none_source_fields() {
365        let dp = DataPoint::new("Entity", None);
366        let m = dp.vector_metadata();
367        assert!(!m.contains_key("source_pipeline"));
368        assert!(!m.contains_key("source_task"));
369        assert!(!m.contains_key("source_user"));
370        assert!(!m.contains_key("source_node_set"));
371        assert!(!m.contains_key("source_content_hash"));
372    }
373
374    #[test]
375    fn test_touch_updates_timestamp() {
376        let mut dp = DataPoint::new("Entity", None);
377        let original_time = dp.updated_at;
378
379        std::thread::sleep(std::time::Duration::from_millis(10));
380        dp.touch();
381
382        assert!(dp.updated_at > original_time);
383    }
384}
cognee_models/data_point.rs

cognee_models/
data_point.rs
