sqlite_graphrag/
entity_type.rs

1//! Canonical entity type taxonomy used across extraction, storage and CLI.
2//!
3//! `EntityType` is the single source of truth for the 13 graph entity kinds.
4//! It derives `clap::ValueEnum` so CLI flags can use it directly, and derives
5//! `serde::{Serialize, Deserialize}` with `rename_all = "lowercase"` so JSON
6//! round-trips remain backward-compatible with the pre-enum string format.
7
8use crate::errors::AppError;
9
10/// The 13 canonical graph entity classifications.
11///
12/// Values are serialized as lowercase strings (`"person"`, `"organization"`,
13/// etc.) matching the pre-enum wire format and the SQLite `type` column.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, clap::ValueEnum)]
15#[serde(rename_all = "snake_case")]
16#[clap(rename_all = "snake_case")]
17pub enum EntityType {
18    Concept,
19    Date,
20    Dashboard,
21    Decision,
22    File,
23    Incident,
24    IssueTracker,
25    Location,
26    Memory,
27    Organization,
28    Person,
29    Project,
30    Tool,
31}
32
33impl EntityType {
34    /// Returns the canonical lowercase string representation stored in SQLite.
35    pub fn as_str(self) -> &'static str {
36        match self {
37            EntityType::Concept => "concept",
38            EntityType::Date => "date",
39            EntityType::Dashboard => "dashboard",
40            EntityType::Decision => "decision",
41            EntityType::File => "file",
42            EntityType::Incident => "incident",
43            EntityType::IssueTracker => "issue_tracker",
44            EntityType::Location => "location",
45            EntityType::Memory => "memory",
46            EntityType::Organization => "organization",
47            EntityType::Person => "person",
48            EntityType::Project => "project",
49            EntityType::Tool => "tool",
50        }
51    }
52
53    /// Maps an arbitrary type label to the closest canonical [`EntityType`],
54    /// never failing (GAP-SG-47).
55    ///
56    /// LLM extraction routinely emits type labels outside the 13 canonical
57    /// kinds (`platform`, `language`, `feature`, `framework`, ...). The old
58    /// parse path discarded those entities with a `WARN`, silently losing
59    /// legitimate graph nodes. This function PRESERVES them by folding each
60    /// label onto the nearest canonical kind. Anything it cannot place falls
61    /// back to [`EntityType::Concept`], the most general kind — so a label is
62    /// never dropped.
63    ///
64    /// Matching is case-insensitive and treats hyphens as underscores, so
65    /// `"Issue-Tracker"` resolves to [`EntityType::IssueTracker`].
66    /// Exact match against the 13 canonical kinds (case/hyphen-insensitive).
67    fn parse_exact(key: &str) -> Option<EntityType> {
68        match key {
69            "concept" => Some(EntityType::Concept),
70            "date" => Some(EntityType::Date),
71            "dashboard" => Some(EntityType::Dashboard),
72            "decision" => Some(EntityType::Decision),
73            "file" => Some(EntityType::File),
74            "incident" => Some(EntityType::Incident),
75            "issue_tracker" => Some(EntityType::IssueTracker),
76            "location" => Some(EntityType::Location),
77            "memory" => Some(EntityType::Memory),
78            "organization" => Some(EntityType::Organization),
79            "person" => Some(EntityType::Person),
80            "project" => Some(EntityType::Project),
81            "tool" => Some(EntityType::Tool),
82            _ => None,
83        }
84    }
85
86    pub fn map_to_canonical(s: &str) -> EntityType {
87        let key = s.trim().to_lowercase().replace('-', "_");
88        // Exact canonical (and case/hyphen-insensitive) match first.
89        if let Some(et) = Self::parse_exact(&key) {
90            return et;
91        }
92        match key.as_str() {
93            // Concept-like: abstractions, technologies, capabilities, topics.
94            "platform" | "language" | "feature" | "framework" | "library" | "technology"
95            | "software" | "service" | "product" | "system" | "api" | "component" | "module"
96            | "package" | "dependency" | "protocol" | "standard" | "format" | "algorithm"
97            | "pattern" | "method" | "function" | "class" | "interface" | "command" | "flag"
98            | "option" | "config" | "setting" | "version" | "release" | "model" | "metric"
99            | "topic" | "skill" | "reference" | "note" | "feedback" | "url" | "link"
100            | "keyword" | "tag" | "category" => EntityType::Concept,
101            // File-like: documents, paths, code artifacts.
102            "document" | "doc" | "artifact" | "directory" | "folder" | "path" | "repository"
103            | "repo" | "codebase" | "script" => EntityType::File,
104            // Person-like roles.
105            "user" | "author" | "developer" | "maintainer" | "contributor" | "agent" | "owner"
106            | "assignee" => EntityType::Person,
107            // Organization-like collectives.
108            "company" | "org" | "vendor" | "group" | "team" | "department" | "institution" => {
109                EntityType::Organization
110            }
111            // Incident-like failures.
112            "bug" | "error" | "failure" | "outage" | "vulnerability" | "cve" | "regression"
113            | "defect" => EntityType::Incident,
114            // Decision-like records.
115            "adr" | "choice" | "policy" | "ruling" => EntityType::Decision,
116            // Date-like temporals.
117            "time" | "datetime" | "timestamp" | "day" | "month" | "year" | "deadline"
118            | "milestone" => EntityType::Date,
119            // Location-like places.
120            "city" | "country" | "region" | "place" | "address" | "site" => EntityType::Location,
121            // Issue-tracker-like.
122            "ticket" | "issue" | "jira" | "github_issue" | "pr" | "pull_request" => {
123                EntityType::IssueTracker
124            }
125            // Dashboard-like.
126            "panel" | "board" | "view" | "report" | "chart" => EntityType::Dashboard,
127            // Anything else: the most general canonical kind, never dropped.
128            _ => EntityType::Concept,
129        }
130    }
131}
132
133/// v1.1.8: manual `Deserialize` folds aliases via [`EntityType::map_to_canonical`]
134/// so EVERY JSON entry point (`--graph-stdin`, `--entities-file`, enrich)
135/// accepts extraction labels like `module`/`platform` without dropping nodes.
136/// Persistence always stores the 13 canonical kinds.
137impl<'de> serde::Deserialize<'de> for EntityType {
138    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
139    where
140        D: serde::Deserializer<'de>,
141    {
142        let s = String::deserialize(deserializer)?;
143        Ok(EntityType::map_to_canonical(&s))
144    }
145}
146
147impl std::fmt::Display for EntityType {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        f.write_str(self.as_str())
150    }
151}
152
153impl std::str::FromStr for EntityType {
154    type Err = AppError;
155
156    /// Folds non-canonical labels onto the nearest kind (never rejects).
157    /// Canonical kinds round-trip exactly; aliases like `module` → Concept.
158    fn from_str(s: &str) -> Result<Self, Self::Err> {
159        Ok(EntityType::map_to_canonical(s))
160    }
161}
162
163impl rusqlite::types::FromSql for EntityType {
164    fn column_result(value: rusqlite::types::ValueRef<'_>) -> rusqlite::types::FromSqlResult<Self> {
165        let s = String::column_result(value)?;
166        s.parse::<EntityType>().map_err(|e| {
167            rusqlite::types::FromSqlError::Other(Box::new(std::io::Error::other(e.to_string())))
168        })
169    }
170}
171
172impl rusqlite::types::ToSql for EntityType {
173    fn to_sql(&self) -> rusqlite::Result<rusqlite::types::ToSqlOutput<'_>> {
174        Ok(rusqlite::types::ToSqlOutput::from(self.as_str()))
175    }
176}
177
178#[cfg(test)]
179mod tests {
180    use super::*;
181
182    // v1.1.8: serde folds aliases via map_to_canonical (never rejects).
183    #[test]
184    fn deserialize_folds_aliases_and_unknown_to_canonical() {
185        assert_eq!(
186            serde_json::from_str::<EntityType>("\"reference\"").unwrap(),
187            EntityType::Concept
188        );
189        assert_eq!(
190            serde_json::from_str::<EntityType>("\"module\"").unwrap(),
191            EntityType::Concept
192        );
193        assert_eq!(
194            serde_json::from_str::<EntityType>("\"document\"").unwrap(),
195            EntityType::File
196        );
197        assert_eq!(
198            serde_json::from_str::<EntityType>("\"banana\"").unwrap(),
199            EntityType::Concept
200        );
201    }
202
203    #[test]
204    fn deserialize_valid_and_case_insensitive_entity_type() {
205        assert_eq!(
206            serde_json::from_str::<EntityType>("\"issue_tracker\"").unwrap(),
207            EntityType::IssueTracker
208        );
209        // FromStr lowercases, so serde now accepts mixed case like the CLI.
210        assert_eq!(
211            serde_json::from_str::<EntityType>("\"Tool\"").unwrap(),
212            EntityType::Tool
213        );
214    }
215
216    #[test]
217    fn serialize_stays_snake_case() {
218        assert_eq!(
219            serde_json::to_string(&EntityType::IssueTracker).unwrap(),
220            "\"issue_tracker\""
221        );
222    }
223
224    #[test]
225    fn from_str_lowercase_roundtrip() {
226        assert_eq!("person".parse::<EntityType>().unwrap(), EntityType::Person);
227        assert_eq!(
228            "organization".parse::<EntityType>().unwrap(),
229            EntityType::Organization
230        );
231        assert_eq!(
232            "issue_tracker".parse::<EntityType>().unwrap(),
233            EntityType::IssueTracker
234        );
235    }
236
237    #[test]
238    fn from_str_uppercase_is_case_insensitive() {
239        assert_eq!("PERSON".parse::<EntityType>().unwrap(), EntityType::Person);
240        assert_eq!(
241            "Organization".parse::<EntityType>().unwrap(),
242            EntityType::Organization
243        );
244    }
245
246    #[test]
247    fn from_str_unknown_folds_to_concept() {
248            let result = "not-a-real-type".parse::<EntityType>();
249            assert_eq!(result.unwrap(), EntityType::Concept);
250        }
251
252    #[test]
253    fn as_str_returns_canonical_lowercase() {
254        assert_eq!(EntityType::Person.as_str(), "person");
255        assert_eq!(EntityType::IssueTracker.as_str(), "issue_tracker");
256    }
257
258    #[test]
259    fn serde_json_serializes_as_lowercase_string() {
260        let json = serde_json::to_string(&EntityType::Person).unwrap();
261        assert_eq!(json, "\"person\"");
262        let json = serde_json::to_string(&EntityType::IssueTracker).unwrap();
263        assert_eq!(json, "\"issue_tracker\"");
264    }
265
266    #[test]
267    fn serde_json_deserializes_from_lowercase_string() {
268        let et: EntityType = serde_json::from_str("\"person\"").unwrap();
269        assert_eq!(et, EntityType::Person);
270    }
271
272    #[test]
273    fn map_to_canonical_preserves_canonical_types() {
274        assert_eq!(EntityType::map_to_canonical("person"), EntityType::Person);
275        assert_eq!(EntityType::map_to_canonical("concept"), EntityType::Concept);
276        assert_eq!(
277            EntityType::map_to_canonical("issue_tracker"),
278            EntityType::IssueTracker
279        );
280        // Hyphen + case variants normalize to the canonical kind.
281        assert_eq!(
282            EntityType::map_to_canonical("Issue-Tracker"),
283            EntityType::IssueTracker
284        );
285    }
286
287    #[test]
288    fn map_to_canonical_folds_non_canonical_instead_of_discarding() {
289        // GAP-SG-47: platform/language/feature were previously DISCARDED.
290        assert_eq!(
291            EntityType::map_to_canonical("platform"),
292            EntityType::Concept
293        );
294        assert_eq!(
295            EntityType::map_to_canonical("language"),
296            EntityType::Concept
297        );
298        assert_eq!(EntityType::map_to_canonical("feature"), EntityType::Concept);
299        // Role/collective folds.
300        assert_eq!(
301            EntityType::map_to_canonical("developer"),
302            EntityType::Person
303        );
304        assert_eq!(
305            EntityType::map_to_canonical("company"),
306            EntityType::Organization
307        );
308        assert_eq!(EntityType::map_to_canonical("document"), EntityType::File);
309    }
310
311    #[test]
312    fn map_to_canonical_unknown_falls_back_to_concept_never_dropped() {
313        assert_eq!(
314            EntityType::map_to_canonical("totally-made-up-kind"),
315            EntityType::Concept
316        );
317        assert_eq!(EntityType::map_to_canonical(""), EntityType::Concept);
318    }
319}
sqlite_graphrag/entity_type.rs

sqlite_graphrag/
entity_type.rs
