entrenar/sovereign/registry/
types.rs

1//! Core types for the model registry.
2
3use serde::{Deserialize, Serialize};
4use std::collections::HashMap;
5use std::path::PathBuf;
6
7/// Model source type
8#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
9pub enum ModelSource {
10    /// HuggingFace Hub model
11    HuggingFace {
12        /// Repository ID (e.g., "bert-base-uncased")
13        repo_id: String,
14    },
15    /// Local file path
16    LocalFile {
17        /// Path to the model file
18        path: PathBuf,
19    },
20    /// Custom URL source
21    Custom {
22        /// Download URL
23        url: String,
24    },
25}
26
27impl ModelSource {
28    /// Create HuggingFace source
29    pub fn huggingface(repo_id: impl Into<String>) -> Self {
30        Self::HuggingFace { repo_id: repo_id.into() }
31    }
32
33    /// Create local file source
34    pub fn local(path: impl Into<PathBuf>) -> Self {
35        Self::LocalFile { path: path.into() }
36    }
37
38    /// Create custom URL source
39    pub fn custom(url: impl Into<String>) -> Self {
40        Self::Custom { url: url.into() }
41    }
42
43    /// Get a display string for the source
44    pub fn display_string(&self) -> String {
45        match self {
46            Self::HuggingFace { repo_id } => format!("hf://{repo_id}"),
47            Self::LocalFile { path } => format!("file://{}", path.display()),
48            Self::Custom { url } => url.clone(),
49        }
50    }
51}
52
53/// Model entry in the registry
54#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
55pub struct ModelEntry {
56    /// Model name (e.g., "bert-base-uncased")
57    pub name: String,
58    /// Model version
59    pub version: String,
60    /// SHA-256 checksum
61    pub sha256: String,
62    /// Model size in bytes
63    pub size_bytes: u64,
64    /// Model source
65    pub source: ModelSource,
66    /// Local path if mirrored
67    pub local_path: Option<PathBuf>,
68    /// Model format (gguf, safetensors, etc.)
69    pub format: Option<String>,
70    /// Model metadata
71    pub metadata: HashMap<String, String>,
72}
73
74impl ModelEntry {
75    /// Create a new model entry
76    pub fn new(
77        name: impl Into<String>,
78        version: impl Into<String>,
79        sha256: impl Into<String>,
80        size_bytes: u64,
81        source: ModelSource,
82    ) -> Self {
83        Self {
84            name: name.into(),
85            version: version.into(),
86            sha256: sha256.into(),
87            size_bytes,
88            source,
89            local_path: None,
90            format: None,
91            metadata: HashMap::new(),
92        }
93    }
94
95    /// Set local path
96    pub fn with_local_path(mut self, path: impl Into<PathBuf>) -> Self {
97        self.local_path = Some(path.into());
98        self
99    }
100
101    /// Set format
102    pub fn with_format(mut self, format: impl Into<String>) -> Self {
103        self.format = Some(format.into());
104        self
105    }
106
107    /// Add metadata
108    pub fn with_metadata(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
109        self.metadata.insert(key.into(), value.into());
110        self
111    }
112
113    /// Check if model is available locally
114    pub fn is_local(&self) -> bool {
115        self.local_path.as_ref().is_some_and(|p| p.exists())
116    }
117
118    /// Get size in megabytes
119    pub fn size_mb(&self) -> f64 {
120        self.size_bytes as f64 / (1024.0 * 1024.0)
121    }
122
123    /// Get size in gigabytes
124    pub fn size_gb(&self) -> f64 {
125        self.size_bytes as f64 / (1024.0 * 1024.0 * 1024.0)
126    }
127}
entrenar/sovereign/registry/types.rs

entrenar/sovereign/registry/
types.rs
