Skip to main content

gaze_recognizers/
locale_aware.rs

1use std::ops::Range;
2
3use gaze_types::{LocaleTag, PiiClass};
4use thiserror::Error;
5
6/// Locale-aware model backend for Pass-2 NER and Pass-3 safety-net inference.
7pub trait LocaleAwareModel: Send + Sync {
8    /// Stable backend name used in diagnostics and audit trails.
9    fn name(&self) -> &str;
10
11    /// Locales this backend natively supports.
12    fn native_locales(&self) -> &[LocaleTag];
13
14    /// Run inference for the supplied text and locale hints.
15    fn infer(&self, input: ModelInput, hints: ModelHints) -> Result<Vec<ModelSpan>, ModelError>;
16}
17
18/// Input passed to a locale-aware model.
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct ModelInput {
21    pub text: String,
22    pub locale: LocaleTag,
23}
24
25/// Execution hints supplied by the caller.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct ModelHints {
28    pub stage: ModelStage,
29    pub max_spans: Option<u32>,
30}
31
32/// Pipeline stage requesting model inference.
33#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum ModelStage {
35    Pass2Ner,
36    Pass3SafetyNet,
37}
38
39/// Span emitted by a locale-aware model.
40#[derive(Debug, Clone, PartialEq)]
41pub struct ModelSpan {
42    pub text: String,
43    pub byte_range: Range<usize>,
44    pub class: PiiClass,
45    pub confidence: Option<f32>,
46    pub model_name: String,
47}
48
49/// Closed error set for locale-aware model routing and inference.
50#[derive(Debug, Clone, PartialEq, Eq, Error)]
51pub enum ModelError {
52    #[error("backend init failed: {0}")]
53    InitFailed(String),
54    #[error("inference failed: {0}")]
55    InferenceFailed(String),
56    #[error("locale not supported: {0:?}")]
57    LocaleNotSupported(LocaleTag),
58    #[error("model integrity mismatch")]
59    IntegrityMismatch,
60    #[error("no locale model coverage for {locale:?} at {stage:?}")]
61    NoLocaleModelCoverage {
62        locale: LocaleTag,
63        stage: ModelStage,
64    },
65    #[error("backend internal error: {0}")]
66    Internal(String),
67}
68
69/// Registry that resolves locale-aware model backends by locale fallback tier.
70#[derive(Default)]
71pub struct LocaleAwareModelRegistry {
72    backends: Vec<Box<dyn LocaleAwareModel>>,
73}
74
75impl LocaleAwareModelRegistry {
76    pub fn new() -> Self {
77        Self::default()
78    }
79
80    pub fn from_backends(backends: Vec<Box<dyn LocaleAwareModel>>) -> Self {
81        Self { backends }
82    }
83
84    pub fn register(&mut self, backend: impl LocaleAwareModel + 'static) {
85        self.backends.push(Box::new(backend));
86    }
87
88    pub fn len(&self) -> usize {
89        self.backends.len()
90    }
91
92    pub fn is_empty(&self) -> bool {
93        self.backends.is_empty()
94    }
95
96    /// Resolve models by four tiers:
97    /// exact locale, parent language, multilingual `Global`, then fail-closed.
98    pub fn resolve(
99        &self,
100        locale: &LocaleTag,
101        stage: ModelStage,
102    ) -> Result<Vec<&dyn LocaleAwareModel>, ModelError> {
103        if let Some(matches) = self.matches_for(locale) {
104            return Ok(matches);
105        }
106
107        if let Some(parent) = parent_language(locale) {
108            if let Some(matches) = self.matches_for(&parent) {
109                return Ok(matches);
110            }
111        }
112
113        if let Some(matches) = self.matches_for(&LocaleTag::Global) {
114            return Ok(matches);
115        }
116
117        Err(ModelError::NoLocaleModelCoverage {
118            locale: locale.clone(),
119            stage,
120        })
121    }
122
123    fn matches_for(&self, locale: &LocaleTag) -> Option<Vec<&dyn LocaleAwareModel>> {
124        let matches: Vec<&dyn LocaleAwareModel> = self
125            .backends
126            .iter()
127            .filter(|backend| backend.native_locales().contains(locale))
128            .map(|backend| backend.as_ref())
129            .collect();
130
131        (!matches.is_empty()).then_some(matches)
132    }
133}
134
135fn parent_language(locale: &LocaleTag) -> Option<LocaleTag> {
136    match locale {
137        LocaleTag::DeDe | LocaleTag::DeAt | LocaleTag::DeCh => {
138            Some(LocaleTag::Other("de".to_string()))
139        }
140        LocaleTag::EnUs | LocaleTag::EnGb | LocaleTag::EnIe | LocaleTag::EnAu | LocaleTag::EnCa => {
141            Some(LocaleTag::Other("en".to_string()))
142        }
143        LocaleTag::Other(raw) => raw
144            .split_once('-')
145            .map(|(language, _)| LocaleTag::Other(language.to_ascii_lowercase())),
146        LocaleTag::Global => None,
147        _ => None,
148    }
149}