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

sqlite_graphrag/
entity_type.rs
