cognee_models/
entity.rs

1//! Entity - Storage-layer entity model.
2//!
3//! Mirrors Python's `cognee/modules/engine/models/Entity.py`
4//! Represents an entity extracted from text and stored in the graph database.
5
6use serde::{Deserialize, Serialize};
7use uuid::Uuid;
8
9use crate::DataPoint;
10use crate::has_datapoint::HasDataPoint;
11
12/// Storage-layer entity model.
13///
14/// Represents an entity (e.g., "TechCorp", "Alice", "London") extracted
15/// from text. Each entity has a name, description, and a reference to its
16/// EntityType (e.g., "Organization", "Person", "Location").
17#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
18pub struct Entity {
19    /// Base data point fields (id, timestamps, metadata, etc.)
20    #[serde(flatten)]
21    pub base: DataPoint,
22
23    /// Entity name (e.g., "TechCorp")
24    pub name: String,
25
26    /// Reference to EntityType UUID (e.g., UUID of "Organization" type)
27    pub is_a: Option<Uuid>,
28
29    /// Entity description from LLM extraction
30    pub description: String,
31}
32
33impl Entity {
34    /// Index fields to embed for vector search.
35    pub const INDEX_FIELDS: &'static [&'static str] = &["name"];
36
37    /// Deterministic, class-namespaced id for an Entity identity value.
38    ///
39    /// Mirrors Python's `Entity.id_for(value)` (`identity_fields=["name"]`):
40    /// `uuid5(NAMESPACE_OID, "Entity:<normalized_value>")`. This is the single
41    /// source of truth for "what id does the entity with this identity have",
42    /// used both when constructing entities and when looking them up from a raw
43    /// string before an instance exists.
44    pub fn id_for(value: &str) -> Uuid {
45        cognee_utils::data_point_id_for("Entity", &[value])
46    }
47
48    /// Create a new Entity.
49    ///
50    /// # Arguments
51    /// * `name` - Entity name
52    /// * `entity_type_id` - Optional reference to EntityType
53    /// * `description` - Entity description
54    /// * `dataset_id` - Dataset UUID
55    pub fn new(
56        name: impl Into<String>,
57        entity_type_id: Option<Uuid>,
58        description: impl Into<String>,
59        dataset_id: Option<Uuid>,
60    ) -> Self {
61        let name = name.into();
62        let mut metadata = std::collections::HashMap::new();
63        metadata.insert(
64            "index_fields".to_string(),
65            serde_json::json!(Self::INDEX_FIELDS),
66        );
67
68        // Deterministic, class-namespaced id derived from the identity value
69        // (`name`), mirroring Python's `identity_fields=["name"]` derivation in
70        // `DataPoint.__init__`. Prevents the random-uuid4 footgun that made the
71        // same entity duplicate across cognify runs.
72        let mut base = DataPoint::with_metadata("Entity", dataset_id, metadata);
73        base.id = Self::id_for(&name);
74
75        Self {
76            base,
77            name,
78            is_a: entity_type_id,
79            description: description.into(),
80        }
81    }
82
83    /// Create Entity from LLM-extracted Node.
84    ///
85    /// # Arguments
86    /// * `node_id` - Original node ID from LLM extraction
87    /// * `node_name` - Node name
88    /// * `node_description` - Node description
89    /// * `entity_type_id` - EntityType UUID
90    /// * `dataset_id` - Dataset UUID
91    pub fn from_node(
92        node_id: impl Into<String>,
93        node_name: impl Into<String>,
94        node_description: impl Into<String>,
95        entity_type_id: Uuid,
96        dataset_id: Option<Uuid>,
97    ) -> Self {
98        let node_id = node_id.into();
99        let mut entity = Self::new(
100            node_name,
101            Some(entity_type_id),
102            node_description,
103            dataset_id,
104        );
105
106        // Python `_create_entity_node` hashes the LLM-supplied node id (not the
107        // display name) into the id — `Entity(id=Entity.id_for(node_id), …)`
108        // (expand_with_nodes_and_edges.py:183,209). Override the name-derived id
109        // from `new` with the node-id-derived one for faithful parity.
110        entity.base.id = Self::id_for(&node_id);
111
112        entity
113            .base
114            .set_metadata("original_node_id", serde_json::json!(node_id));
115
116        entity
117    }
118
119    /// Get the entity name (for embedding).
120    pub fn get_embeddable_text(&self) -> String {
121        self.name.clone()
122    }
123
124    /// Update entity description.
125    pub fn set_description(&mut self, description: impl Into<String>) {
126        self.description = description.into();
127        self.base.touch();
128    }
129
130    /// Update entity type reference.
131    pub fn set_entity_type(&mut self, entity_type_id: Uuid) {
132        self.is_a = Some(entity_type_id);
133        self.base.touch();
134    }
135}
136
137impl HasDataPoint for Entity {
138    fn data_point(&self) -> &DataPoint {
139        &self.base
140    }
141    fn data_point_mut(&mut self) -> &mut DataPoint {
142        &mut self.base
143    }
144    // for_each_child_mut: default no-op — Entity references its EntityType
145    // by UUID (`is_a: Option<Uuid>`), not by ownership. If a future variant
146    // owns an `entity_type: Box<EntityType>` field, override here to recurse.
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    #[test]
154    fn test_entity_creation() {
155        let entity = Entity::new("TechCorp", None, "A technology company", None);
156
157        assert_eq!(entity.name, "TechCorp");
158        assert_eq!(entity.description, "A technology company");
159        assert_eq!(entity.base.data_type, "Entity");
160        assert!(entity.is_a.is_none());
161    }
162
163    #[test]
164    fn test_entity_with_type() {
165        let type_id = Uuid::new_v4();
166        let entity = Entity::new("TechCorp", Some(type_id), "A technology company", None);
167
168        assert_eq!(entity.is_a, Some(type_id));
169    }
170
171    #[test]
172    fn test_entity_from_node() {
173        let type_id = Uuid::new_v4();
174        let entity = Entity::from_node(
175            "techcorp_1",
176            "TechCorp",
177            "A technology company",
178            type_id,
179            None,
180        );
181
182        assert_eq!(entity.name, "TechCorp");
183        assert_eq!(entity.is_a, Some(type_id));
184        assert_eq!(
185            entity.base.get_metadata("original_node_id"),
186            Some(&serde_json::json!("techcorp_1"))
187        );
188    }
189
190    #[test]
191    fn test_id_for_matches_python() {
192        // Python: Entity.id_for("Alice") = uuid5(OID, "Entity:alice")
193        assert_eq!(
194            Entity::id_for("Alice"),
195            Uuid::new_v5(&Uuid::NAMESPACE_OID, b"Entity:alice"),
196        );
197    }
198
199    #[test]
200    fn test_new_id_is_deterministic_from_name() {
201        // Two entities with the same name resolve to the same id — this is what
202        // lets the same entity merge across cognify runs (regresses issue #57's
203        // inverse: random ids caused silent duplication).
204        let a = Entity::new("Acme Corp", None, "desc a", None);
205        let b = Entity::new(
206            "Acme Corp",
207            Some(Uuid::new_v4()),
208            "desc b",
209            Some(Uuid::new_v4()),
210        );
211        assert_eq!(a.base.id, b.base.id);
212        assert_eq!(a.base.id, Entity::id_for("Acme Corp"));
213    }
214
215    #[test]
216    fn test_from_node_hashes_node_id_not_name() {
217        // Python hashes the LLM node id, not the display name.
218        let e = Entity::from_node("node-42", "Alice", "desc", Uuid::new_v4(), None);
219        assert_eq!(e.base.id, Entity::id_for("node-42"));
220        assert_ne!(e.base.id, Entity::id_for("Alice"));
221    }
222
223    #[test]
224    fn test_entity_and_entity_type_ids_do_not_collide() {
225        // The class prefix keeps a Person "institution" and a type "institution"
226        // on distinct ids (topoteretes/cognee#2510/#2515).
227        assert_ne!(
228            Entity::id_for("institution"),
229            crate::EntityType::id_for("institution"),
230        );
231    }
232
233    #[test]
234    fn test_entity_index_fields() {
235        let entity = Entity::new("TechCorp", None, "A company", None);
236        let index_fields = entity.base.get_metadata("index_fields");
237
238        assert_eq!(index_fields, Some(&serde_json::json!(["name"])));
239    }
240
241    #[test]
242    fn test_entity_embeddable_text() {
243        let entity = Entity::new("TechCorp", None, "A company", None);
244        assert_eq!(entity.get_embeddable_text(), "TechCorp");
245    }
246
247    #[test]
248    fn test_entity_set_description() {
249        let mut entity = Entity::new("TechCorp", None, "Old desc", None);
250        let old_time = entity.base.updated_at;
251
252        std::thread::sleep(std::time::Duration::from_millis(10));
253        entity.set_description("New description");
254
255        assert_eq!(entity.description, "New description");
256        // updated_at is i64 (millis since epoch); touch() should advance it
257        assert!(entity.base.updated_at >= old_time);
258    }
259
260    #[test]
261    fn test_entity_set_type() {
262        let mut entity = Entity::new("TechCorp", None, "A company", None);
263        let type_id = Uuid::new_v4();
264
265        entity.set_entity_type(type_id);
266        assert_eq!(entity.is_a, Some(type_id));
267    }
268
269    #[test]
270    fn entity_implements_has_datapoint() {
271        let e = Entity::new("Foo", None, "desc", None);
272        let dp_id = e.base.id;
273        assert_eq!(e.data_point().id, dp_id);
274        let mut e2 = e;
275        assert_eq!(e2.data_point_mut().id, dp_id);
276    }
277}
cognee_models/entity.rs

cognee_models/
entity.rs
