ai_model_vault/
formats.rs

1//! Model format detection and conversion registry
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5
6use crate::error::{Result, VaultError};
7
8/// Supported AI model formats
9#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
10pub enum ModelFormat {
11    // LLM-centric formats
12    /// Safetensors (.safetensors) - HuggingFace default for Transformers
13    Safetensors,
14    /// GGUF (.gguf) - Quantized LLM format (llama.cpp, LM Studio, Ollama)
15    GGUF,
16    /// PyTorch weights (.pt, .pth, .bin) - Classic state_dict files
17    PyTorch,
18    /// TensorRT Engine (.plan) - NVIDIA compiled engines
19    TensorRT,
20    /// ONNX (.onnx) - Interchange/serving format
21    ONNX,
22    /// MLX (.npz) - Apple Silicon optimized format
23    MLX,
24    /// Core ML (.mlmodel) - iOS/macOS on-device inference
25    CoreML,
26    /// TorchScript (.pt traced/scripted) - PyTorch serialization
27    TorchScript,
28    /// TensorFlow Lite (.tflite) - Mobile/edge deployment
29    TFLite,
30
31    // General DL formats
32    /// TensorFlow SavedModel - TensorFlow serving format
33    TensorFlow,
34    /// Keras (.h5, .keras) - Keras model format
35    Keras,
36    /// OpenVINO IR (.xml + .bin) - Intel optimization format
37    OpenVINO,
38    /// TVM compiled artifacts
39    TVM,
40    /// NCNN (.param + .bin) - Mobile-optimized format
41    NCNN,
42    /// MNN (.mnn) - Mobile Neural Network format
43    MNN,
44    /// RKNN (.rknn) - Rockchip NPU format
45    RKNN,
46
47    // Legacy formats
48    /// Caffe (.caffemodel) - Legacy Caffe format
49    Caffe,
50    /// MXNet (.params) - Apache MXNet format
51    MXNet,
52    /// Darknet (.weights) - YOLO/Darknet format
53    Darknet,
54
55    // Data formats
56    /// HDF5 (.h5, .hdf5) - Hierarchical data format
57    HDF5,
58    /// Pickle (.pkl) - Python pickle format
59    Pickle,
60    /// NumPy (.npy, .npz) - NumPy array format
61    NumPy,
62
63    /// Custom/Unknown format
64    Custom(String),
65}
66
67impl ModelFormat {
68    /// Detect format from file extension
69    pub fn from_extension(ext: &str) -> Self {
70        match ext.to_lowercase().as_str() {
71            "safetensors" => ModelFormat::Safetensors,
72            "gguf" => ModelFormat::GGUF,
73            "pt" | "pth" | "bin" => ModelFormat::PyTorch,
74            "plan" => ModelFormat::TensorRT,
75            "onnx" => ModelFormat::ONNX,
76            "mlx" => ModelFormat::MLX,
77            "mlmodel" | "mlmodelc" => ModelFormat::CoreML,
78            "tflite" => ModelFormat::TFLite,
79            "pb" => ModelFormat::TensorFlow,
80            "h5" | "keras" => ModelFormat::Keras,
81            "xml" => ModelFormat::OpenVINO,
82            "param" => ModelFormat::NCNN,
83            "mnn" => ModelFormat::MNN,
84            "rknn" => ModelFormat::RKNN,
85            "caffemodel" => ModelFormat::Caffe,
86            "params" => ModelFormat::MXNet,
87            "weights" => ModelFormat::Darknet,
88            "hdf5" => ModelFormat::HDF5,
89            "pkl" | "pickle" => ModelFormat::Pickle,
90            "npy" | "npz" => ModelFormat::NumPy,
91            _ => ModelFormat::Custom(ext.to_string()),
92        }
93    }
94
95    /// Get file extension for format
96    pub fn extension(&self) -> &str {
97        match self {
98            ModelFormat::Safetensors => "safetensors",
99            ModelFormat::GGUF => "gguf",
100            ModelFormat::PyTorch => "pt",
101            ModelFormat::TensorRT => "plan",
102            ModelFormat::ONNX => "onnx",
103            ModelFormat::MLX => "npz",
104            ModelFormat::CoreML => "mlmodel",
105            ModelFormat::TorchScript => "pt",
106            ModelFormat::TFLite => "tflite",
107            ModelFormat::TensorFlow => "pb",
108            ModelFormat::Keras => "h5",
109            ModelFormat::OpenVINO => "xml",
110            ModelFormat::TVM => "so",
111            ModelFormat::NCNN => "param",
112            ModelFormat::MNN => "mnn",
113            ModelFormat::RKNN => "rknn",
114            ModelFormat::Caffe => "caffemodel",
115            ModelFormat::MXNet => "params",
116            ModelFormat::Darknet => "weights",
117            ModelFormat::HDF5 => "h5",
118            ModelFormat::Pickle => "pkl",
119            ModelFormat::NumPy => "npy",
120            ModelFormat::Custom(ext) => ext,
121        }
122    }
123
124    /// Every non-`Custom` variant, for reverse lookups.
125    const ALL: [ModelFormat; 22] = [
126        ModelFormat::Safetensors,
127        ModelFormat::GGUF,
128        ModelFormat::PyTorch,
129        ModelFormat::TensorRT,
130        ModelFormat::ONNX,
131        ModelFormat::MLX,
132        ModelFormat::CoreML,
133        ModelFormat::TorchScript,
134        ModelFormat::TFLite,
135        ModelFormat::TensorFlow,
136        ModelFormat::Keras,
137        ModelFormat::OpenVINO,
138        ModelFormat::TVM,
139        ModelFormat::NCNN,
140        ModelFormat::MNN,
141        ModelFormat::RKNN,
142        ModelFormat::Caffe,
143        ModelFormat::MXNet,
144        ModelFormat::Darknet,
145        ModelFormat::HDF5,
146        ModelFormat::Pickle,
147        ModelFormat::NumPy,
148    ];
149
150    /// Parse a format from its display [`name`](Self::name), case- and
151    /// space-insensitively (`"Core ML"`, `"coreml"`, `"CORE ML"` all match).
152    #[must_use]
153    pub fn from_name(name: &str) -> Option<Self> {
154        let norm = |s: &str| s.to_lowercase().replace([' ', '-', '_'], "");
155        let target = norm(name);
156        Self::ALL.into_iter().find(|f| norm(f.name()) == target)
157    }
158
159    /// Parse a format string that may be either a display name or a file
160    /// extension.
161    ///
162    /// Version records persist `format.name()` (e.g. `"PyTorch"`), which
163    /// [`from_extension`](Self::from_extension) does not recognise — it would
164    /// silently yield `Custom("pytorch")` and break conversion-path lookup and
165    /// tensor-level diffing. Use this whenever the input came from storage.
166    #[must_use]
167    pub fn from_stored(value: &str) -> Self {
168        Self::from_name(value).unwrap_or_else(|| Self::from_extension(value))
169    }
170
171    /// Get format name
172    pub fn name(&self) -> &str {
173        match self {
174            ModelFormat::Safetensors => "Safetensors",
175            ModelFormat::GGUF => "GGUF",
176            ModelFormat::PyTorch => "PyTorch",
177            ModelFormat::TensorRT => "TensorRT",
178            ModelFormat::ONNX => "ONNX",
179            ModelFormat::MLX => "MLX",
180            ModelFormat::CoreML => "Core ML",
181            ModelFormat::TorchScript => "TorchScript",
182            ModelFormat::TFLite => "TensorFlow Lite",
183            ModelFormat::TensorFlow => "TensorFlow",
184            ModelFormat::Keras => "Keras",
185            ModelFormat::OpenVINO => "OpenVINO",
186            ModelFormat::TVM => "TVM",
187            ModelFormat::NCNN => "NCNN",
188            ModelFormat::MNN => "MNN",
189            ModelFormat::RKNN => "RKNN",
190            ModelFormat::Caffe => "Caffe",
191            ModelFormat::MXNet => "MXNet",
192            ModelFormat::Darknet => "Darknet",
193            ModelFormat::HDF5 => "HDF5",
194            ModelFormat::Pickle => "Pickle",
195            ModelFormat::NumPy => "NumPy",
196            ModelFormat::Custom(name) => name,
197        }
198    }
199
200    /// Format label safe to report in telemetry.
201    ///
202    /// [`name`](Self::name) returns the caller's own string for
203    /// [`ModelFormat::Custom`], which is free-form user input — a model name,
204    /// a path fragment, anything they typed. Reporting it would break the
205    /// guarantee that telemetry carries no user-supplied text, so this
206    /// collapses every `Custom` to the literal `"custom"` and is otherwise a
207    /// fixed set of `&'static str`.
208    #[must_use]
209    pub fn telemetry_name(&self) -> &'static str {
210        match self {
211            ModelFormat::Safetensors => "safetensors",
212            ModelFormat::GGUF => "gguf",
213            ModelFormat::PyTorch => "pytorch",
214            ModelFormat::TensorRT => "tensorrt",
215            ModelFormat::ONNX => "onnx",
216            ModelFormat::MLX => "mlx",
217            ModelFormat::CoreML => "coreml",
218            ModelFormat::TorchScript => "torchscript",
219            ModelFormat::TFLite => "tflite",
220            ModelFormat::TensorFlow => "tensorflow",
221            ModelFormat::Keras => "keras",
222            ModelFormat::OpenVINO => "openvino",
223            ModelFormat::TVM => "tvm",
224            ModelFormat::NCNN => "ncnn",
225            ModelFormat::MNN => "mnn",
226            ModelFormat::RKNN => "rknn",
227            ModelFormat::Caffe => "caffe",
228            ModelFormat::MXNet => "mxnet",
229            ModelFormat::Darknet => "darknet",
230            ModelFormat::HDF5 => "hdf5",
231            ModelFormat::Pickle => "pickle",
232            ModelFormat::NumPy => "numpy",
233            // Deliberately NOT `name`: that is whatever the user typed.
234            ModelFormat::Custom(_) => "custom",
235        }
236    }
237}
238
239impl std::fmt::Display for ModelFormat {
240    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
241        write!(f, "{}", self.name())
242    }
243}
244
245/// Model metadata
246#[derive(Debug, Clone, Serialize, Deserialize)]
247pub struct ModelMetadata {
248    pub name: String,
249    pub format: ModelFormat,
250    pub description: Option<String>,
251    pub framework: Option<String>,
252    pub task: Option<String>,
253    pub architecture: Option<String>,
254    pub parameters: Option<u64>,
255    pub custom_fields: HashMap<String, String>,
256}
257
258impl ModelMetadata {
259    /// Create new model metadata
260    pub fn new(name: String, format: ModelFormat) -> Self {
261        Self {
262            name,
263            format,
264            description: None,
265            framework: None,
266            task: None,
267            architecture: None,
268            parameters: None,
269            custom_fields: HashMap::new(),
270        }
271    }
272
273    /// Builder pattern for optional fields
274    pub fn with_description(mut self, desc: String) -> Self {
275        self.description = Some(desc);
276        self
277    }
278
279    /// Set the framework used to train the model.
280    pub fn with_framework(mut self, framework: String) -> Self {
281        self.framework = Some(framework);
282        self
283    }
284
285    /// Set the model's target task (e.g., classification, generation).
286    pub fn with_task(mut self, task: String) -> Self {
287        self.task = Some(task);
288        self
289    }
290
291    /// Set the model architecture (e.g., Transformer, CNN).
292    pub fn with_architecture(mut self, arch: String) -> Self {
293        self.architecture = Some(arch);
294        self
295    }
296
297    /// Set the model parameter count.
298    pub fn with_parameters(mut self, params: u64) -> Self {
299        self.parameters = Some(params);
300        self
301    }
302
303    /// Add a custom metadata field as a key-value pair.
304    pub fn add_custom_field(mut self, key: String, value: String) -> Self {
305        self.custom_fields.insert(key, value);
306        self
307    }
308}
309
310/// Type alias for format conversion function
311type ConversionFn = fn(&[u8]) -> Result<Vec<u8>>;
312
313/// Format converter registry
314pub struct FormatConverter {
315    converters: HashMap<(ModelFormat, ModelFormat), ConversionFn>,
316}
317
318impl FormatConverter {
319    /// Create new format converter registry
320    pub fn new() -> Self {
321        Self {
322            converters: HashMap::new(),
323        }
324    }
325
326    /// Register a format converter
327    pub fn register(
328        &mut self,
329        from: ModelFormat,
330        to: ModelFormat,
331        converter: fn(&[u8]) -> Result<Vec<u8>>,
332    ) {
333        self.converters.insert((from, to), converter);
334    }
335
336    /// Check if conversion is supported
337    pub fn can_convert(&self, from: ModelFormat, to: ModelFormat) -> bool {
338        self.converters.contains_key(&(from, to))
339    }
340
341    /// Convert between formats
342    pub fn convert(&self, data: &[u8], from: ModelFormat, to: ModelFormat) -> Result<Vec<u8>> {
343        if from == to {
344            return Ok(data.to_vec());
345        }
346
347        let converter = self
348            .converters
349            .get(&(from.clone(), to.clone()))
350            .ok_or_else(|| {
351                VaultError::ConversionError(format!(
352                    "No converter available for {} -> {}",
353                    from.name(),
354                    to.name()
355                ))
356            })?;
357
358        converter(data)
359    }
360}
361
362impl Default for FormatConverter {
363    fn default() -> Self {
364        Self::new()
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    #[test]
373    fn test_format_detection() {
374        assert_eq!(
375            ModelFormat::from_extension("safetensors"),
376            ModelFormat::Safetensors
377        );
378        assert_eq!(ModelFormat::from_extension("gguf"), ModelFormat::GGUF);
379        assert_eq!(ModelFormat::from_extension("pt"), ModelFormat::PyTorch);
380        assert_eq!(ModelFormat::from_extension("onnx"), ModelFormat::ONNX);
381        assert_eq!(ModelFormat::from_extension("tflite"), ModelFormat::TFLite);
382        assert_eq!(ModelFormat::from_extension("mlmodel"), ModelFormat::CoreML);
383        assert_eq!(ModelFormat::from_extension("plan"), ModelFormat::TensorRT);
384    }
385
386    #[test]
387    fn test_format_extension() {
388        assert_eq!(ModelFormat::Safetensors.extension(), "safetensors");
389        assert_eq!(ModelFormat::GGUF.extension(), "gguf");
390        assert_eq!(ModelFormat::PyTorch.extension(), "pt");
391        assert_eq!(ModelFormat::TensorRT.extension(), "plan");
392        assert_eq!(ModelFormat::ONNX.extension(), "onnx");
393        assert_eq!(ModelFormat::TFLite.extension(), "tflite");
394        assert_eq!(ModelFormat::CoreML.extension(), "mlmodel");
395    }
396
397    #[test]
398    fn test_format_names() {
399        assert_eq!(ModelFormat::Safetensors.name(), "Safetensors");
400        assert_eq!(ModelFormat::GGUF.name(), "GGUF");
401        assert_eq!(ModelFormat::TensorRT.name(), "TensorRT");
402        assert_eq!(ModelFormat::CoreML.name(), "Core ML");
403        assert_eq!(ModelFormat::TFLite.name(), "TensorFlow Lite");
404    }
405
406    #[test]
407    fn test_format_extension_all_variants() {
408        // Cover lines 111, 114, 116, 119 — TorchScript, TensorFlow, Keras, TVM
409        assert_eq!(ModelFormat::TorchScript.extension(), "pt");
410        assert_eq!(ModelFormat::TensorFlow.extension(), "pb");
411        assert_eq!(ModelFormat::Keras.extension(), "h5");
412        assert_eq!(ModelFormat::TVM.extension(), "so");
413        assert_eq!(ModelFormat::OpenVINO.extension(), "xml");
414        assert_eq!(ModelFormat::NCNN.extension(), "param");
415        assert_eq!(ModelFormat::MNN.extension(), "mnn");
416        assert_eq!(ModelFormat::RKNN.extension(), "rknn");
417        assert_eq!(ModelFormat::HDF5.extension(), "h5");
418        assert_eq!(ModelFormat::NumPy.extension(), "npy");
419        assert_eq!(ModelFormat::Custom("custom".into()).extension(), "custom");
420    }
421
422    #[test]
423    fn test_format_name_all_variants() {
424        // Cover line 146 — TorchScript name
425        assert_eq!(ModelFormat::TorchScript.name(), "TorchScript");
426        assert_eq!(ModelFormat::TensorFlow.name(), "TensorFlow");
427        assert_eq!(ModelFormat::Keras.name(), "Keras");
428        assert_eq!(ModelFormat::TVM.name(), "TVM");
429        assert_eq!(ModelFormat::Custom("custom".into()).name(), "custom");
430    }
431
432    #[test]
433    fn test_model_metadata_new() {
434        // Cover line 175 — ModelMetadata::new
435        let meta = ModelMetadata::new("model".to_string(), ModelFormat::PyTorch);
436        assert_eq!(meta.name, "model");
437        assert!(meta.description.is_none());
438    }
439
440    #[test]
441    fn test_format_converter_register_and_convert() {
442        // Cover line 258 — converter(data)
443        let mut fc = FormatConverter::new();
444        fc.register(ModelFormat::PyTorch, ModelFormat::ONNX, |data| {
445            Ok(data.to_vec())
446        });
447        assert!(fc.can_convert(ModelFormat::PyTorch, ModelFormat::ONNX));
448        let result = fc
449            .convert(b"test", ModelFormat::PyTorch, ModelFormat::ONNX)
450            .unwrap();
451        assert_eq!(result, b"test");
452    }
453
454    #[test]
455    fn test_format_converter_same_format() {
456        let fc = FormatConverter::new();
457        let result = fc
458            .convert(b"data", ModelFormat::PyTorch, ModelFormat::PyTorch)
459            .unwrap();
460        assert_eq!(result, b"data");
461    }
462
463    #[test]
464    fn test_format_display() {
465        let s = format!("{}", ModelFormat::Safetensors);
466        assert_eq!(s, "Safetensors");
467    }
468
469    #[test]
470    fn test_from_extension_all_variants() {
471        // Cover all from_extension branches
472        assert_eq!(ModelFormat::from_extension("mlx"), ModelFormat::MLX);
473        assert_eq!(ModelFormat::from_extension("pb"), ModelFormat::TensorFlow);
474        assert_eq!(ModelFormat::from_extension("h5"), ModelFormat::Keras);
475        assert_eq!(ModelFormat::from_extension("keras"), ModelFormat::Keras);
476        assert_eq!(ModelFormat::from_extension("xml"), ModelFormat::OpenVINO);
477        assert_eq!(ModelFormat::from_extension("param"), ModelFormat::NCNN);
478        assert_eq!(ModelFormat::from_extension("mnn"), ModelFormat::MNN);
479        assert_eq!(ModelFormat::from_extension("rknn"), ModelFormat::RKNN);
480        assert_eq!(
481            ModelFormat::from_extension("caffemodel"),
482            ModelFormat::Caffe
483        );
484        assert_eq!(ModelFormat::from_extension("params"), ModelFormat::MXNet);
485        assert_eq!(ModelFormat::from_extension("weights"), ModelFormat::Darknet);
486        assert_eq!(ModelFormat::from_extension("hdf5"), ModelFormat::HDF5);
487        assert_eq!(ModelFormat::from_extension("pkl"), ModelFormat::Pickle);
488        assert_eq!(ModelFormat::from_extension("pickle"), ModelFormat::Pickle);
489        assert_eq!(ModelFormat::from_extension("npy"), ModelFormat::NumPy);
490        assert_eq!(ModelFormat::from_extension("npz"), ModelFormat::NumPy);
491        assert_eq!(ModelFormat::from_extension("pth"), ModelFormat::PyTorch);
492        assert_eq!(ModelFormat::from_extension("bin"), ModelFormat::PyTorch);
493        assert_eq!(ModelFormat::from_extension("mlmodelc"), ModelFormat::CoreML);
494        assert_eq!(
495            ModelFormat::from_extension("xyz"),
496            ModelFormat::Custom("xyz".to_string())
497        );
498    }
499
500    #[test]
501    fn test_name_all_remaining_variants() {
502        assert_eq!(ModelFormat::NCNN.name(), "NCNN");
503        assert_eq!(ModelFormat::MNN.name(), "MNN");
504        assert_eq!(ModelFormat::RKNN.name(), "RKNN");
505        assert_eq!(ModelFormat::Caffe.name(), "Caffe");
506        assert_eq!(ModelFormat::MXNet.name(), "MXNet");
507        assert_eq!(ModelFormat::Darknet.name(), "Darknet");
508        assert_eq!(ModelFormat::HDF5.name(), "HDF5");
509        assert_eq!(ModelFormat::Pickle.name(), "Pickle");
510        assert_eq!(ModelFormat::NumPy.name(), "NumPy");
511        assert_eq!(ModelFormat::MLX.name(), "MLX");
512        assert_eq!(ModelFormat::OpenVINO.name(), "OpenVINO");
513    }
514
515    #[test]
516    fn test_model_metadata_builder_chain() {
517        let meta = ModelMetadata::new("m".to_string(), ModelFormat::ONNX)
518            .with_description("desc".to_string())
519            .with_framework("pytorch".to_string())
520            .with_task("classification".to_string())
521            .with_architecture("ResNet".to_string())
522            .with_parameters(1_000_000)
523            .add_custom_field("license".to_string(), "MIT".to_string());
524
525        assert_eq!(meta.description, Some("desc".to_string()));
526        assert_eq!(meta.framework, Some("pytorch".to_string()));
527        assert_eq!(meta.task, Some("classification".to_string()));
528        assert_eq!(meta.architecture, Some("ResNet".to_string()));
529        assert_eq!(meta.parameters, Some(1_000_000));
530        assert_eq!(meta.custom_fields.get("license"), Some(&"MIT".to_string()));
531    }
532
533    #[test]
534    fn test_format_converter_missing_conversion() {
535        let fc = FormatConverter::new();
536        let result = fc.convert(b"data", ModelFormat::PyTorch, ModelFormat::ONNX);
537        assert!(result.is_err());
538        assert!(!fc.can_convert(ModelFormat::PyTorch, ModelFormat::ONNX));
539    }
540
541    #[test]
542    fn test_format_converter_default() {
543        let fc = FormatConverter::default();
544        assert!(!fc.can_convert(ModelFormat::PyTorch, ModelFormat::ONNX));
545    }
546
547    #[test]
548    fn test_extension_all_remaining() {
549        assert_eq!(ModelFormat::Caffe.extension(), "caffemodel");
550        assert_eq!(ModelFormat::MXNet.extension(), "params");
551        assert_eq!(ModelFormat::Darknet.extension(), "weights");
552        assert_eq!(ModelFormat::Pickle.extension(), "pkl");
553        assert_eq!(ModelFormat::MLX.extension(), "npz");
554    }
555
556    #[test]
557    fn test_from_name_roundtrip_for_every_variant() {
558        // Version records persist `name()`, so every name must parse back to the
559        // same variant — otherwise stored formats silently degrade to Custom.
560        for fmt in ModelFormat::ALL {
561            assert_eq!(
562                ModelFormat::from_name(fmt.name()),
563                Some(fmt.clone()),
564                "name() -> from_name() must round-trip for {:?}",
565                fmt
566            );
567        }
568    }
569
570    #[test]
571    fn test_from_name_is_case_and_space_insensitive() {
572        assert_eq!(ModelFormat::from_name("core ml"), Some(ModelFormat::CoreML));
573        assert_eq!(ModelFormat::from_name("CoreML"), Some(ModelFormat::CoreML));
574        assert_eq!(ModelFormat::from_name("CORE-ML"), Some(ModelFormat::CoreML));
575        assert_eq!(
576            ModelFormat::from_name("tensorflow lite"),
577            Some(ModelFormat::TFLite)
578        );
579        assert_eq!(ModelFormat::from_name("not a format"), None);
580    }
581
582    #[test]
583    fn test_from_stored_accepts_names_and_extensions() {
584        // The shape actually written into version records.
585        assert_eq!(ModelFormat::from_stored("PyTorch"), ModelFormat::PyTorch);
586        assert_eq!(
587            ModelFormat::from_stored("Safetensors"),
588            ModelFormat::Safetensors
589        );
590        assert_eq!(ModelFormat::from_stored("ONNX"), ModelFormat::ONNX);
591        // Still accepts plain extensions.
592        assert_eq!(ModelFormat::from_stored("pt"), ModelFormat::PyTorch);
593        assert_eq!(ModelFormat::from_stored("onnx"), ModelFormat::ONNX);
594    }
595
596    #[test]
597    fn test_from_extension_alone_cannot_parse_stored_names() {
598        // Documents the trap `from_stored` exists to avoid: this is what the
599        // convert and diff paths used to do with a stored format string.
600        assert_ne!(ModelFormat::from_extension("PyTorch"), ModelFormat::PyTorch);
601    }
602
603    /// `name()` returns the caller's own string for `Custom`, so it must never
604    /// be what telemetry reports. This is the guard on that: whatever a user
605    /// types as `--format`, the telemetry label stays a fixed literal.
606    #[test]
607    fn test_telemetry_name_never_echoes_a_custom_format_string() {
608        for hostile in [
609            "/home/alice/models/customer-data.bin",
610            "s3://acme-private/secret-model",
611            "Bearer abc123",
612            "",
613        ] {
614            let fmt = ModelFormat::Custom(hostile.to_string());
615            assert_eq!(
616                fmt.telemetry_name(),
617                "custom",
618                "telemetry_name leaked a Custom payload"
619            );
620            assert_eq!(fmt.name(), hostile, "name() should still be verbatim");
621        }
622    }
623
624    /// Every non-Custom variant must map to a distinct lowercase literal, or
625    /// the collector groups unrelated formats together.
626    #[test]
627    fn test_telemetry_names_are_distinct_and_lowercase() {
628        let all = [
629            ModelFormat::Safetensors,
630            ModelFormat::GGUF,
631            ModelFormat::PyTorch,
632            ModelFormat::TensorRT,
633            ModelFormat::ONNX,
634            ModelFormat::MLX,
635            ModelFormat::CoreML,
636            ModelFormat::TorchScript,
637            ModelFormat::TFLite,
638            ModelFormat::TensorFlow,
639            ModelFormat::Keras,
640            ModelFormat::OpenVINO,
641            ModelFormat::TVM,
642            ModelFormat::NCNN,
643            ModelFormat::MNN,
644            ModelFormat::RKNN,
645            ModelFormat::Caffe,
646            ModelFormat::MXNet,
647            ModelFormat::Darknet,
648            ModelFormat::HDF5,
649            ModelFormat::Pickle,
650            ModelFormat::NumPy,
651        ];
652        let mut seen = std::collections::HashSet::new();
653        for f in &all {
654            let label = f.telemetry_name();
655            assert!(
656                label
657                    .chars()
658                    .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit()),
659                "{label:?} is not a plain lowercase token"
660            );
661            assert!(seen.insert(label), "duplicate telemetry label {label:?}");
662        }
663        assert_eq!(seen.len(), all.len());
664    }
665}
ai_model_vault/formats.rs

ai_model_vault/
formats.rs
