cognee_models/
entity_type.rs

1//! EntityType - Storage-layer entity type model.
2//!
3//! Mirrors Python's `cognee/modules/engine/models/EntityType.py`
4//! Represents a category/type of entities (e.g., "Organization", "Person", "Location").
5
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9use crate::DataPoint;
10use crate::has_datapoint::HasDataPoint;
11
12/// Storage-layer entity type model.
13///
14/// Represents a category of entities (e.g., "Organization", "Person", "Location").
15/// Entity instances reference their EntityType via the `is_a` field.
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
17pub struct EntityType {
18    /// Base data point fields (id, timestamps, metadata, etc.)
19    #[serde(flatten)]
20    pub base: DataPoint,
21
22    /// Type name (e.g., "Organization", "Person", "Location")
23    pub name: String,
24
25    /// Type description
26    pub description: String,
27}
28
29impl EntityType {
30    /// Index fields to embed for vector search.
31    pub const INDEX_FIELDS: &'static [&'static str] = &["name"];
32
33    /// Deterministic, class-namespaced id for an EntityType identity value.
34    ///
35    /// Mirrors Python's `EntityType.id_for(value)` (`identity_fields=["name"]`):
36    /// `uuid5(NAMESPACE_OID, "EntityType:<normalized_value>")`. The distinct
37    /// class prefix is what prevents an `Entity` and an `EntityType` with the
38    /// same name from colliding on one id (topoteretes/cognee#2510/#2515).
39    pub fn id_for(value: &str) -> Uuid {
40        cognee_utils::data_point_id_for("EntityType", &[value])
41    }
42
43    /// Create a new EntityType.
44    ///
45    /// # Arguments
46    /// * `name` - Type name (e.g., "Organization")
47    /// * `description` - Type description
48    /// * `dataset_id` - Dataset UUID
49    pub fn new(
50        name: impl Into<String>,
51        description: impl Into<String>,
52        dataset_id: Option<Uuid>,
53    ) -> Self {
54        let mut metadata = std::collections::HashMap::new();
55        metadata.insert(
56            "index_fields".to_string(),
57            serde_json::json!(Self::INDEX_FIELDS),
58        );
59
60        let name_str = name.into();
61        let description_str = description.into();
62
63        // Deterministic, class-namespaced id derived from `name`, mirroring
64        // Python's `identity_fields=["name"]` derivation.
65        let mut base = DataPoint::with_metadata("EntityType", dataset_id, metadata);
66        base.id = Self::id_for(&name_str);
67
68        Self {
69            base,
70            name: name_str.clone(),
71            description: if description_str.is_empty() {
72                format!("Entity type: {name_str}")
73            } else {
74                description_str
75            },
76        }
77    }
78
79    /// Create EntityType from LLM-extracted node type string.
80    ///
81    /// # Arguments
82    /// * `type_name` - Node type from LLM (e.g., "Organization")
83    /// * `dataset_id` - Dataset UUID
84    pub fn from_node_type(type_name: impl Into<String>, dataset_id: Option<Uuid>) -> Self {
85        let type_str = type_name.into();
86        Self::new(
87            type_str.clone(),
88            format!("Entity type: {type_str}"),
89            dataset_id,
90        )
91    }
92
93    /// Get the type name (for embedding).
94    pub fn get_embeddable_text(&self) -> String {
95        self.name.clone()
96    }
97
98    /// Update type description.
99    pub fn set_description(&mut self, description: impl Into<String>) {
100        self.description = description.into();
101        self.base.touch();
102    }
103
104    /// Check if this type has been validated against an ontology.
105    pub fn is_ontology_valid(&self) -> bool {
106        self.base.ontology_valid
107    }
108
109    /// Mark as ontology-validated with canonical name.
110    ///
111    /// # Arguments
112    /// * `canonical_name` - Canonical name from ontology
113    pub fn mark_ontology_valid(&mut self, canonical_name: Option<String>) {
114        self.base.set_ontology_valid(true);
115
116        if let Some(canonical) = canonical_name
117            && canonical != self.name
118        {
119            self.base
120                .set_metadata("original_name", serde_json::json!(self.name.clone()));
121            self.name = canonical;
122        }
123    }
124}
125
126impl HasDataPoint for EntityType {
127    fn data_point(&self) -> &DataPoint {
128        &self.base
129    }
130    fn data_point_mut(&mut self) -> &mut DataPoint {
131        &mut self.base
132    }
133    // for_each_child_mut: default no-op — EntityType is a leaf in the
134    // model graph (no owned `HasDataPoint` children).
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    #[test]
142    fn test_entity_type_creation() {
143        let et = EntityType::new("Organization", "A company or institution", None);
144
145        assert_eq!(et.name, "Organization");
146        assert_eq!(et.description, "A company or institution");
147        assert_eq!(et.base.data_type, "EntityType");
148    }
149
150    #[test]
151    fn test_entity_type_empty_description() {
152        let et = EntityType::new("Person", "", None);
153
154        assert_eq!(et.name, "Person");
155        assert_eq!(et.description, "Entity type: Person");
156    }
157
158    #[test]
159    fn test_entity_type_from_node_type() {
160        let et = EntityType::from_node_type("Location", None);
161
162        assert_eq!(et.name, "Location");
163        assert_eq!(et.description, "Entity type: Location");
164    }
165
166    #[test]
167    fn test_id_for_matches_python() {
168        // Python: EntityType.id_for("Organization") = uuid5(OID, "EntityType:organization")
169        assert_eq!(
170            EntityType::id_for("Organization"),
171            Uuid::new_v5(&Uuid::NAMESPACE_OID, b"EntityType:organization"),
172        );
173    }
174
175    #[test]
176    fn test_new_id_is_deterministic_from_name() {
177        let a = EntityType::from_node_type("Organization", None);
178        let b = EntityType::new(
179            "Organization",
180            "different description",
181            Some(Uuid::new_v4()),
182        );
183        assert_eq!(a.base.id, b.base.id);
184        assert_eq!(a.base.id, EntityType::id_for("Organization"));
185    }
186
187    #[test]
188    fn test_entity_type_index_fields() {
189        let et = EntityType::new("Organization", "A company", None);
190        let index_fields = et.base.get_metadata("index_fields");
191
192        assert_eq!(index_fields, Some(&serde_json::json!(["name"])));
193    }
194
195    #[test]
196    fn test_entity_type_embeddable_text() {
197        let et = EntityType::new("Organization", "A company", None);
198        assert_eq!(et.get_embeddable_text(), "Organization");
199    }
200
201    #[test]
202    fn test_entity_type_set_description() {
203        let mut et = EntityType::new("Organization", "Old desc", None);
204        et.set_description("New description");
205        assert_eq!(et.description, "New description");
206    }
207
208    #[test]
209    fn test_ontology_validation() {
210        let mut et = EntityType::new("Mathematician", "", None);
211        assert!(!et.is_ontology_valid());
212
213        // Mark as valid with canonical name
214        et.mark_ontology_valid(Some("Person".to_string()));
215
216        assert!(et.is_ontology_valid());
217        assert_eq!(et.name, "Person");
218        assert_eq!(
219            et.base.get_metadata("original_name"),
220            Some(&serde_json::json!("Mathematician"))
221        );
222    }
223
224    #[test]
225    fn test_ontology_validation_same_name() {
226        let mut et = EntityType::new("Person", "", None);
227        et.mark_ontology_valid(Some("Person".to_string()));
228
229        assert!(et.is_ontology_valid());
230        assert_eq!(et.name, "Person");
231        assert_eq!(et.base.get_metadata("original_name"), None);
232    }
233
234    #[test]
235    fn test_ontology_validation_no_canonical() {
236        let mut et = EntityType::new("Person", "", None);
237        et.mark_ontology_valid(None);
238
239        assert!(et.is_ontology_valid());
240        assert_eq!(et.name, "Person");
241    }
242
243    #[test]
244    fn entity_type_implements_has_datapoint() {
245        let et = EntityType::new("Org", "desc", None);
246        let dp_id = et.base.id;
247        assert_eq!(et.data_point().id, dp_id);
248        let mut et2 = et;
249        assert_eq!(et2.data_point_mut().id, dp_id);
250    }
251}
cognee_models/entity_type.rs

cognee_models/
entity_type.rs
