laminar_db/ai/
registry.rs

1//! Model registry: maps SQL-referenced model names to backends and tasks.
2//! Built once at startup; AI functions resolve against it at plan time.
3
4use std::collections::HashMap;
5use std::fmt;
6use std::str::FromStr;
7
8use thiserror::Error;
9
10/// Task contract for an AI SQL function.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12pub enum Task {
13    /// `ai_classify`
14    Classify,
15    /// `ai_sentiment`
16    Sentiment,
17    /// `ai_embed`
18    Embed,
19    /// `ai_extract`
20    Extract,
21    /// `ai_complete`
22    Complete,
23    /// `ai_summarize`
24    Summarize,
25    /// `ai_translate`
26    Translate,
27    /// `ai_gen`
28    Gen,
29}
30
31impl Task {
32    /// Canonical name matching the `task = "…"` config spelling.
33    #[must_use]
34    pub fn as_str(self) -> &'static str {
35        match self {
36            Task::Classify => "classify",
37            Task::Sentiment => "sentiment",
38            Task::Embed => "embed",
39            Task::Extract => "extract",
40            Task::Complete => "complete",
41            Task::Summarize => "summarize",
42            Task::Translate => "translate",
43            Task::Gen => "gen",
44        }
45    }
46}
47
48impl fmt::Display for Task {
49    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
50        f.write_str(self.as_str())
51    }
52}
53
54impl FromStr for Task {
55    type Err = RegistryError;
56
57    fn from_str(s: &str) -> Result<Self, Self::Err> {
58        match s {
59            "classify" => Ok(Task::Classify),
60            "sentiment" => Ok(Task::Sentiment),
61            "embed" => Ok(Task::Embed),
62            "extract" => Ok(Task::Extract),
63            "complete" => Ok(Task::Complete),
64            "summarize" => Ok(Task::Summarize),
65            "translate" => Ok(Task::Translate),
66            "gen" => Ok(Task::Gen),
67            other => Err(RegistryError::UnknownTask(other.to_string())),
68        }
69    }
70}
71
72/// Which runtime executes a model.
73#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
74pub enum BackendKind {
75    /// In-process ONNX Runtime.
76    Local,
77    /// Configured remote provider.
78    Remote,
79}
80
81/// Backend-specific wiring for a registered model.
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub enum ModelBackend {
84    /// In-process ONNX model.
85    Local {
86        /// `hf:org/repo`, `file://<path>`, or a bare path.
87        source: String,
88        /// `id2label` from config.json; `None` until derived at load time.
89        labels: Option<Vec<String>>,
90    },
91    /// Remote provider model.
92    Remote {
93        /// Key into the configured providers map.
94        provider: String,
95        /// Provider-specific model id.
96        model: String,
97    },
98}
99
100/// One registered model.
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct ModelEntry {
103    /// Name referenced as `model => '<id>'` in SQL.
104    pub id: String,
105    /// Tasks this model can perform.
106    pub tasks: Vec<Task>,
107    /// Backend that runs it.
108    pub backend: ModelBackend,
109}
110
111impl ModelEntry {
112    /// Backend kind for this entry.
113    #[must_use]
114    pub fn kind(&self) -> BackendKind {
115        match self.backend {
116            ModelBackend::Local { .. } => BackendKind::Local,
117            ModelBackend::Remote { .. } => BackendKind::Remote,
118        }
119    }
120
121    /// Whether this model supports `task`.
122    #[must_use]
123    pub fn supports(&self, task: Task) -> bool {
124        self.tasks.contains(&task)
125    }
126
127    /// Whether results are deterministic (local only) and permanently cacheable.
128    #[must_use]
129    pub fn is_deterministic(&self) -> bool {
130        matches!(self.kind(), BackendKind::Local)
131    }
132
133    /// Whether calls incur metered cost (remote providers).
134    #[must_use]
135    pub fn is_costed(&self) -> bool {
136        matches!(self.kind(), BackendKind::Remote)
137    }
138
139    /// Intrinsic labels of a local classifier, if known.
140    #[must_use]
141    pub fn labels(&self) -> Option<&[String]> {
142        match &self.backend {
143            ModelBackend::Local { labels, .. } => labels.as_deref(),
144            ModelBackend::Remote { .. } => None,
145        }
146    }
147}
148
149/// Map of model name to entry, with optional per-task defaults.
150#[derive(Debug, Default)]
151pub struct ModelRegistry {
152    models: HashMap<String, ModelEntry>,
153    defaults: HashMap<Task, String>,
154}
155
156impl ModelRegistry {
157    /// Create an empty registry.
158    #[must_use]
159    pub fn new() -> Self {
160        Self::default()
161    }
162
163    /// Register a model.
164    ///
165    /// # Errors
166    ///
167    /// Returns [`RegistryError::DuplicateModel`] if the id is already registered.
168    pub fn register(&mut self, entry: ModelEntry) -> Result<(), RegistryError> {
169        if self.models.contains_key(&entry.id) {
170            return Err(RegistryError::DuplicateModel(entry.id.clone()));
171        }
172        self.models.insert(entry.id.clone(), entry);
173        Ok(())
174    }
175
176    /// Set the default model for a task (`[ai.defaults]` config block).
177    pub fn set_default(&mut self, task: Task, model: impl Into<String>) {
178        self.defaults.insert(task, model.into());
179    }
180
181    /// Default model for a task, if configured.
182    #[must_use]
183    pub fn default_for(&self, task: Task) -> Option<&str> {
184        self.defaults.get(&task).map(String::as_str)
185    }
186
187    /// Resolve a model by name.
188    ///
189    /// # Errors
190    ///
191    /// Returns [`RegistryError::UnknownModel`] if not registered.
192    pub fn resolve(&self, name: &str) -> Result<&ModelEntry, RegistryError> {
193        self.models
194            .get(name)
195            .ok_or_else(|| RegistryError::UnknownModel(name.to_string()))
196    }
197
198    /// Resolve a model and confirm it supports `task`.
199    ///
200    /// # Errors
201    ///
202    /// Returns [`RegistryError::UnknownModel`] or [`RegistryError::TaskUnsupported`].
203    pub fn validate(&self, name: &str, task: Task) -> Result<&ModelEntry, RegistryError> {
204        let entry = self.resolve(name)?;
205        if entry.supports(task) {
206            Ok(entry)
207        } else {
208            Err(RegistryError::TaskUnsupported {
209                model: name.to_string(),
210                task,
211                supported: entry.tasks.clone(),
212            })
213        }
214    }
215
216    /// Number of registered models.
217    #[must_use]
218    pub fn len(&self) -> usize {
219        self.models.len()
220    }
221
222    /// Whether the registry is empty.
223    #[must_use]
224    pub fn is_empty(&self) -> bool {
225        self.models.is_empty()
226    }
227
228    /// Iterate entries in unspecified order. Backs `laminar.models`.
229    pub fn iter(&self) -> impl Iterator<Item = &ModelEntry> {
230        self.models.values()
231    }
232}
233
234/// Errors from resolving or validating a model.
235#[derive(Debug, Error, PartialEq, Eq)]
236pub enum RegistryError {
237    /// No model with the given name is registered.
238    #[error("unknown model '{0}'")]
239    UnknownModel(String),
240
241    /// The model exists but does not support the requested task.
242    #[error("model '{model}' does not support task '{task}' (supports: {supported:?})")]
243    TaskUnsupported {
244        /// The model name.
245        model: String,
246        /// The requested task.
247        task: Task,
248        /// Tasks the model actually supports.
249        supported: Vec<Task>,
250    },
251
252    /// A model with this id is already registered.
253    #[error("model '{0}' is already registered")]
254    DuplicateModel(String),
255
256    /// A `task = "…"` config string named an unknown task.
257    #[error("unknown task '{0}'")]
258    UnknownTask(String),
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    fn local_classifier() -> ModelEntry {
266        ModelEntry {
267            id: "finbert".to_string(),
268            tasks: vec![Task::Classify, Task::Sentiment],
269            backend: ModelBackend::Local {
270                source: "hf:onnx-community/finbert".to_string(),
271                labels: Some(vec![
272                    "positive".to_string(),
273                    "negative".to_string(),
274                    "neutral".to_string(),
275                ]),
276            },
277        }
278    }
279
280    fn remote_llm() -> ModelEntry {
281        ModelEntry {
282            id: "haiku".to_string(),
283            tasks: vec![Task::Classify, Task::Extract, Task::Complete],
284            backend: ModelBackend::Remote {
285                provider: "anthropic".to_string(),
286                model: "claude-haiku-4-5-20251001".to_string(),
287            },
288        }
289    }
290
291    #[test]
292    fn resolve_and_validate() {
293        let mut reg = ModelRegistry::new();
294        assert!(reg.is_empty());
295        reg.register(local_classifier()).unwrap();
296        reg.register(remote_llm()).unwrap();
297        reg.set_default(Task::Classify, "finbert");
298        assert_eq!(reg.len(), 2);
299
300        // Unknown model.
301        assert_eq!(
302            reg.resolve("missing").unwrap_err(),
303            RegistryError::UnknownModel("missing".to_string())
304        );
305
306        // Supported task resolves.
307        assert_eq!(
308            reg.validate("finbert", Task::Sentiment).unwrap().id,
309            "finbert"
310        );
311
312        // Unsupported task is rejected with the supported set.
313        match reg.validate("finbert", Task::Complete).unwrap_err() {
314            RegistryError::TaskUnsupported {
315                model,
316                task,
317                supported,
318            } => {
319                assert_eq!(model, "finbert");
320                assert_eq!(task, Task::Complete);
321                assert_eq!(supported, vec![Task::Classify, Task::Sentiment]);
322            }
323            other => panic!("unexpected error: {other}"),
324        }
325
326        assert_eq!(reg.default_for(Task::Classify), Some("finbert"));
327        assert_eq!(reg.default_for(Task::Embed), None);
328    }
329
330    #[test]
331    fn duplicate_registration_rejected() {
332        let mut reg = ModelRegistry::new();
333        reg.register(local_classifier()).unwrap();
334        assert_eq!(
335            reg.register(local_classifier()).unwrap_err(),
336            RegistryError::DuplicateModel("finbert".to_string())
337        );
338    }
339}
laminar_db/ai/registry.rs

laminar_db/ai/
registry.rs
