Skip to main content

xz_knowledge_graph/types/
entity.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use super::attribute::AttributeValue;
6
7/// Entity type with hierarchical classification.
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
9pub enum EntityType {
10    Person,
11    Location,
12    Item,
13    Organization,
14    Concept,
15    Resource,
16    Ability,
17    Event(String),
18    Custom { category: String, label: String },
19}
20
21impl EntityType {
22    pub fn category(&self) -> &str {
23        match self {
24            Self::Person => "Person",
25            Self::Location => "Location",
26            Self::Item => "Item",
27            Self::Organization => "Organization",
28            Self::Concept => "Concept",
29            Self::Resource => "Resource",
30            Self::Ability => "Ability",
31            Self::Event(_) => "Event",
32            Self::Custom { category, .. } => category,
33        }
34    }
35
36    pub fn as_str(&self) -> String {
37        match self {
38            Self::Person => "Person".into(),
39            Self::Location => "Location".into(),
40            Self::Item => "Item".into(),
41            Self::Organization => "Organization".into(),
42            Self::Concept => "Concept".into(),
43            Self::Resource => "Resource".into(),
44            Self::Ability => "Ability".into(),
45            Self::Event(s) => format!("Event:{}", s),
46            Self::Custom { category, label } => format!("{}:{}", category, label),
47        }
48    }
49
50    pub fn from_str(s: &str) -> Self {
51        match s {
52            "Person" => Self::Person,
53            "Location" => Self::Location,
54            "Item" => Self::Item,
55            "Organization" => Self::Organization,
56            "Concept" => Self::Concept,
57            "Resource" => Self::Resource,
58            "Ability" => Self::Ability,
59            other => {
60                if let Some(rest) = other.strip_prefix("Event:") {
61                    Self::Event(rest.to_string())
62                } else if let Some((cat, label)) = other.split_once(':') {
63                    Self::Custom {
64                        category: cat.to_string(),
65                        label: label.to_string(),
66                    }
67                } else {
68                    Self::Custom {
69                        category: "Other".into(),
70                        label: other.to_string(),
71                    }
72                }
73            }
74        }
75    }
76}
77
78/// Knowledge graph entity.
79#[derive(Debug, Clone, Serialize, Deserialize)]
80pub struct Entity {
81    pub id: String,
82    pub name: String,
83    pub aliases: Vec<String>,
84    pub entity_type: EntityType,
85    pub attributes: HashMap<String, AttributeValue>,
86    pub description: Option<String>,
87    pub created_at: u64,
88    pub updated_at: u64,
89    pub version: u64,
90    pub source: Option<String>,
91    pub tags: Vec<String>,
92}