cognee_models/
edge_type.rs

1//! EdgeType - Storage-layer edge type model for indexing.
2//!
3//! Mirrors Python's `cognee/modules/engine/models/EdgeType.py`
4//! Represents a type of relationship (e.g., "works_at", "located_in", "knows").
5
6use chrono::Utc;
7use serde::{Deserialize, Serialize};
8use uuid::Uuid;
9
10use crate::DataPoint;
11use crate::has_datapoint::HasDataPoint;
12
13/// Storage-layer edge type model.
14///
15/// Represents a type of relationship between entities (e.g., "works_at",
16/// "located_in", "knows"). Used for indexing and semantic search of
17/// relationship types.
18#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
19pub struct EdgeType {
20    /// Base data point fields (id, timestamps, metadata, etc.)
21    #[serde(flatten)]
22    pub base: DataPoint,
23
24    /// Relationship name (e.g., "works_at", "located_in")
25    pub relationship_name: String,
26
27    /// Number of edges of this type (for statistics)
28    pub number_of_edges: i32,
29}
30
31impl EdgeType {
32    /// Index fields to embed for vector search.
33    pub const INDEX_FIELDS: &'static [&'static str] = &["relationship_name"];
34
35    /// Compute a deterministic UUID for an EdgeType from its relationship name.
36    ///
37    /// Mirrors Python's `EdgeType.id_for(relationship_name)`
38    /// (`identity_fields=["relationship_name"]`):
39    /// `uuid5(NAMESPACE_OID, "EdgeType:<normalized relationship_name>")`.
40    ///
41    /// The `EdgeType:` class prefix was added upstream in cognee 1.2.0
42    /// (`namespace_edge_type_point_ids.py`) — the previous bare-name scheme
43    /// (`uuid5(OID, normalized)`) let a relationship and a same-named node
44    /// collide on one point id.
45    pub fn deterministic_id(relationship_name: &str) -> Uuid {
46        cognee_utils::data_point_id_for("EdgeType", &[relationship_name])
47    }
48
49    /// Resolve the retrieval text for an edge, mirroring Python's
50    /// `get_edge_retrieval_text(edge_text, relationship_name)`
51    /// (`prepare_edges_for_storage.py:26-28` via `index_graph_edges.py:33-53`):
52    /// prefer the nonblank `edge_text`, fall back to the nonblank
53    /// `relationship_name`, else return an empty string (caller drops empties).
54    ///
55    /// Takes primitive `&str` params (not a `GraphEdgePair`, which lives in
56    /// `cognee-cognify` and would introduce a cyclic dependency) so any lane
57    /// can recompute the exact text an `EdgeType`/`Triplet` vector id was
58    /// derived from.
59    pub fn retrieval_text(edge_text: Option<&str>, relationship_name: &str) -> String {
60        if let Some(text) = edge_text.map(str::trim).filter(|s| !s.is_empty()) {
61            return text.to_string();
62        }
63        relationship_name.trim().to_string()
64    }
65
66    /// Create a new EdgeType with a random UUID.
67    ///
68    /// # Arguments
69    /// * `relationship_name` - Relationship name (e.g., "works_at")
70    /// * `dataset_id` - Dataset UUID
71    pub fn new(relationship_name: impl Into<String>, dataset_id: Option<Uuid>) -> Self {
72        let mut metadata = std::collections::HashMap::new();
73        metadata.insert(
74            "index_fields".to_string(),
75            serde_json::json!(Self::INDEX_FIELDS),
76        );
77
78        Self {
79            base: DataPoint::with_metadata("EdgeType", dataset_id, metadata),
80            relationship_name: relationship_name.into(),
81            number_of_edges: 0,
82        }
83    }
84
85    /// Create a new EdgeType with a deterministic UUID derived from the
86    /// relationship name, matching Python's `generate_edge_id`.
87    ///
88    /// # Arguments
89    /// * `relationship_name` - Relationship name (e.g., "works_at")
90    /// * `dataset_id` - Dataset UUID
91    pub fn new_deterministic(
92        relationship_name: impl Into<String>,
93        dataset_id: Option<Uuid>,
94    ) -> Self {
95        let name = relationship_name.into();
96        let id = Self::deterministic_id(&name);
97        let now = Utc::now().timestamp_millis();
98
99        let mut metadata = std::collections::HashMap::new();
100        metadata.insert(
101            "index_fields".to_string(),
102            serde_json::json!(Self::INDEX_FIELDS),
103        );
104
105        Self {
106            base: DataPoint {
107                id,
108                created_at: now,
109                updated_at: now,
110                ontology_valid: false,
111                version: 1,
112                topological_rank: None,
113                metadata,
114                data_type: "EdgeType".to_string(),
115                belongs_to_set: dataset_id.map(|ds_id| vec![serde_json::json!(ds_id.to_string())]),
116                source_pipeline: None,
117                source_task: None,
118                source_node_set: None,
119                source_user: None,
120                source_content_hash: None,
121                feedback_weight: 0.5,
122                importance_weight: Some(0.5),
123            },
124            relationship_name: name,
125            number_of_edges: 0,
126        }
127    }
128
129    /// Get the relationship name (for embedding).
130    pub fn get_embeddable_text(&self) -> String {
131        self.relationship_name.clone()
132    }
133
134    /// Increment the edge count.
135    pub fn increment_count(&mut self) {
136        self.number_of_edges += 1;
137        self.base.touch();
138    }
139
140    /// Set the edge count.
141    pub fn set_count(&mut self, count: i32) {
142        self.number_of_edges = count;
143        self.base.touch();
144    }
145
146    /// Get the edge count.
147    pub fn count(&self) -> i32 {
148        self.number_of_edges
149    }
150}
151
152impl HasDataPoint for EdgeType {
153    fn data_point(&self) -> &DataPoint {
154        &self.base
155    }
156    fn data_point_mut(&mut self) -> &mut DataPoint {
157        &mut self.base
158    }
159    // for_each_child_mut: default no-op — EdgeType has no nested
160    // `HasDataPoint` children.
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn test_edge_type_creation() {
169        let et = EdgeType::new("works_at", None);
170
171        assert_eq!(et.relationship_name, "works_at");
172        assert_eq!(et.number_of_edges, 0);
173        assert_eq!(et.base.data_type, "EdgeType");
174    }
175
176    #[test]
177    fn test_edge_type_with_dataset() {
178        let dataset_id = Uuid::new_v4();
179        let et = EdgeType::new("works_at", Some(dataset_id));
180
181        assert_eq!(
182            et.base.belongs_to_set,
183            Some(vec![serde_json::json!(dataset_id.to_string())])
184        );
185    }
186
187    #[test]
188    fn test_edge_type_index_fields() {
189        let et = EdgeType::new("works_at", None);
190        let index_fields = et.base.get_metadata("index_fields");
191
192        assert_eq!(
193            index_fields,
194            Some(&serde_json::json!(["relationship_name"]))
195        );
196    }
197
198    #[test]
199    fn test_edge_type_embeddable_text() {
200        let et = EdgeType::new("works_at", None);
201        assert_eq!(et.get_embeddable_text(), "works_at");
202    }
203
204    #[test]
205    fn test_edge_type_increment_count() {
206        let mut et = EdgeType::new("works_at", None);
207        assert_eq!(et.count(), 0);
208
209        et.increment_count();
210        assert_eq!(et.count(), 1);
211
212        et.increment_count();
213        assert_eq!(et.count(), 2);
214    }
215
216    #[test]
217    fn test_edge_type_set_count() {
218        let mut et = EdgeType::new("works_at", None);
219        et.set_count(10);
220        assert_eq!(et.count(), 10);
221    }
222
223    #[test]
224    fn test_edge_type_increment_updates_timestamp() {
225        let mut et = EdgeType::new("works_at", None);
226        let old_time = et.base.updated_at;
227
228        std::thread::sleep(std::time::Duration::from_millis(10));
229        et.increment_count();
230
231        // updated_at is i64 (millis since epoch); touch() should advance it
232        assert!(et.base.updated_at >= old_time);
233    }
234
235    #[test]
236    fn test_deterministic_id_basic() {
237        let id1 = EdgeType::deterministic_id("works_at");
238        let id2 = EdgeType::deterministic_id("works_at");
239        assert_eq!(id1, id2, "same input must produce same UUID");
240    }
241
242    #[test]
243    fn test_deterministic_id_normalization() {
244        // Spaces become underscores, apostrophes removed, lowercased
245        let id1 = EdgeType::deterministic_id("Works At");
246        let id2 = EdgeType::deterministic_id("works_at");
247        assert_eq!(
248            id1, id2,
249            "normalization should make 'Works At' equal 'works_at'"
250        );
251
252        let id3 = EdgeType::deterministic_id("it's_related");
253        let id4 = EdgeType::deterministic_id("its_related");
254        assert_eq!(id3, id4, "apostrophe removal should match");
255    }
256
257    #[test]
258    fn test_deterministic_id_matches_python() {
259        // Python: EdgeType.id_for("works_at") = uuid5(OID, "EdgeType:works_at")
260        // (class-namespaced since cognee 1.2.0, namespace_edge_type_point_ids.py).
261        let id = EdgeType::deterministic_id("works_at");
262        assert_eq!(
263            id,
264            Uuid::new_v5(&Uuid::NAMESPACE_OID, b"EdgeType:works_at"),
265            "deterministic_id('works_at') should equal uuid5(OID, 'EdgeType:works_at')"
266        );
267    }
268
269    #[test]
270    fn test_new_deterministic_constructor() {
271        let et = EdgeType::new_deterministic("works_at", None);
272        assert_eq!(et.relationship_name, "works_at");
273        assert_eq!(et.base.data_type, "EdgeType");
274        assert_eq!(et.base.id, EdgeType::deterministic_id("works_at"));
275        assert_eq!(et.number_of_edges, 0);
276    }
277
278    #[test]
279    fn test_new_deterministic_with_dataset() {
280        let dataset_id = Uuid::new_v4();
281        let et = EdgeType::new_deterministic("located_in", Some(dataset_id));
282        assert_eq!(
283            et.base.belongs_to_set,
284            Some(vec![serde_json::json!(dataset_id.to_string())])
285        );
286        assert_eq!(et.base.id, EdgeType::deterministic_id("located_in"));
287    }
288
289    #[test]
290    fn test_deterministic_id_different_names_differ() {
291        let id1 = EdgeType::deterministic_id("works_at");
292        let id2 = EdgeType::deterministic_id("located_in");
293        assert_ne!(id1, id2, "different names must produce different UUIDs");
294    }
295
296    #[test]
297    fn retrieval_text_prefers_nonblank_edge_text() {
298        assert_eq!(
299            EdgeType::retrieval_text(Some("Alice knows Bob"), "knows"),
300            "Alice knows Bob"
301        );
302    }
303
304    #[test]
305    fn retrieval_text_falls_back_when_edge_text_none() {
306        assert_eq!(EdgeType::retrieval_text(None, "knows"), "knows");
307    }
308
309    #[test]
310    fn retrieval_text_falls_back_when_edge_text_blank() {
311        assert_eq!(EdgeType::retrieval_text(Some("   "), "knows"), "knows");
312    }
313
314    #[test]
315    fn retrieval_text_trims_both_sources() {
316        assert_eq!(
317            EdgeType::retrieval_text(Some("  spaced text  "), "knows"),
318            "spaced text"
319        );
320        assert_eq!(EdgeType::retrieval_text(None, "  knows  "), "knows");
321    }
322
323    #[test]
324    fn edge_type_implements_has_datapoint() {
325        let et = EdgeType::new("rel", None);
326        let dp_id = et.base.id;
327        assert_eq!(et.data_point().id, dp_id);
328        let mut et2 = et;
329        assert_eq!(et2.data_point_mut().id, dp_id);
330    }
331}
cognee_models/edge_type.rs

cognee_models/
edge_type.rs
