offline_intelligence/model_management/
storage.rs

1//! Model Storage Management
2//!
3//! Handles local storage of models in platform-appropriate locations:
4//! - Windows: %APPDATA%/Aud.io/models
5//! - Linux: ~/.local/share/aud.io/models
6//! - macOS: ~/Library/Application Support/Aud.io/models
7
8use anyhow::{Context, Result};
9use serde::{Deserialize, Serialize};
10use std::path::PathBuf;
11use tracing::{debug, info};
12
13/// Sanitize a model ID for use as a filesystem directory/file name.
14/// Replaces characters invalid on Windows (: / \ < > " | ? *) with underscores.
15/// Also trims trailing periods and spaces, handles reserved names, and limits length.
16fn sanitize_model_id(model_id: &str) -> String {
17    let mut sanitized = model_id
18        .replace(':', "_")
19        .replace('/', "_")
20        .replace('\\', "_")
21        .replace('<', "_")
22        .replace('>', "_")
23        .replace('"', "_")
24        .replace('|', "_")
25        .replace('?', "_")
26        .replace('*', "_");
27
28    // Trim trailing periods and spaces which are invalid on Windows
29    sanitized = sanitized.trim_end_matches('.').trim_end().to_string();
30
31    // Handle Windows reserved names by adding a suffix if needed
32    let upper_sanitized = sanitized.to_uppercase();
33    if is_windows_reserved_name(&upper_sanitized) {
34        sanitized.push_str("_model");
35    }
36
37    // Limit length to avoid Windows path length issues (keep under 200 chars to be safe)
38    if sanitized.len() > 200 {
39        // Hash the long name and append a portion of the original for readability
40        use std::collections::hash_map::DefaultHasher;
41        use std::hash::{Hash, Hasher};
42        
43        let mut hasher = DefaultHasher::new();
44        sanitized.hash(&mut hasher);
45        let hash = format!("{:x}", hasher.finish());
46        
47        // Take first 150 chars of sanitized + hash to ensure uniqueness while keeping readable
48        let prefix_len = std::cmp::min(150, sanitized.len());
49        sanitized = format!("{}_{}", &sanitized[..prefix_len], hash);
50    }
51
52    sanitized
53}
54
55/// Check if a name is a Windows reserved name
56fn is_windows_reserved_name(name: &str) -> bool {
57    matches!(name, 
58        "CON" | "PRN" | "AUX" | "NUL" |
59        "COM1" | "COM2" | "COM3" | "COM4" | "COM5" | "COM6" | "COM7" | "COM8" | "COM9" |
60        "LPT1" | "LPT2" | "LPT3" | "LPT4" | "LPT5" | "LPT6" | "LPT7" | "LPT8" | "LPT9"
61    )
62}
63
64/// Platform-specific storage location
65#[derive(Debug, Clone)]
66pub struct StorageLocation {
67    /// Base directory for all Aud.io data
68    pub app_data_dir: PathBuf,
69    /// Directory for storing downloaded models
70    pub models_dir: PathBuf,
71    /// Directory for model metadata and registry
72    pub registry_dir: PathBuf,
73}
74
75/// Model storage manager
76pub struct ModelStorage {
77    pub location: StorageLocation,
78}
79
80impl ModelStorage {
81    /// Create new storage manager with platform-appropriate paths
82    pub fn new() -> Result<Self> {
83        let location = Self::get_platform_storage_location()?;
84
85        // Ensure directories exist
86        Self::ensure_directories(&location)?;
87
88        info!(
89            "Model storage initialized at: {}",
90            location.models_dir.display()
91        );
92
93        Ok(Self { location })
94    }
95
96    /// Get platform-appropriate storage location
97    fn get_platform_storage_location() -> Result<StorageLocation> {
98        let app_data_dir = if cfg!(target_os = "windows") {
99            // Windows: %APPDATA%\Aud.io
100            dirs::data_dir()
101                .context("Failed to get APPDATA directory")?
102                .join("Aud.io")
103        } else if cfg!(target_os = "macos") {
104            // macOS: ~/Library/Application Support/Aud.io
105            dirs::data_dir()
106                .context("Failed to get Library directory")?
107                .join("Aud.io")
108        } else {
109            // Linux: ~/.local/share/aud.io
110            dirs::data_dir()
111                .context("Failed to get .local/share directory")?
112                .join("aud.io")
113        };
114
115        let models_dir = app_data_dir.join("models");
116        let registry_dir = app_data_dir.join("registry");
117
118        Ok(StorageLocation {
119            app_data_dir,
120            models_dir,
121            registry_dir,
122        })
123    }
124
125    /// Ensure all required directories exist
126    fn ensure_directories(location: &StorageLocation) -> Result<()> {
127        std::fs::create_dir_all(&location.app_data_dir)
128            .context("Failed to create app data directory")?;
129        std::fs::create_dir_all(&location.models_dir)
130            .context("Failed to create models directory")?;
131        std::fs::create_dir_all(&location.registry_dir)
132            .context("Failed to create registry directory")?;
133
134        debug!("Created storage directories successfully");
135        Ok(())
136    }
137
138    /// Get the full path for a model file
139    pub fn model_path(&self, model_id: &str, filename: &str) -> PathBuf {
140        self.location
141            .models_dir
142            .join(sanitize_model_id(model_id))
143            .join(filename)
144    }
145
146    /// Get the path for model metadata
147    pub fn metadata_path(&self, model_id: &str) -> PathBuf {
148        self.location
149            .registry_dir
150            .join(format!("{}.json", sanitize_model_id(model_id)))
151    }
152
153    /// List all available models in storage
154    pub fn list_models(&self) -> Result<Vec<String>> {
155        let mut models = Vec::new();
156
157        if self.location.models_dir.exists() {
158            for entry in std::fs::read_dir(&self.location.models_dir)? {
159                let entry = entry?;
160                if entry.file_type()?.is_dir() {
161                    if let Some(model_name) = entry.file_name().to_str() {
162                        models.push(model_name.to_string());
163                    }
164                }
165            }
166        }
167
168        Ok(models)
169    }
170
171    /// Check if a model exists locally
172    pub fn model_exists(&self, model_id: &str) -> bool {
173        self.location
174            .models_dir
175            .join(sanitize_model_id(model_id))
176            .exists()
177    }
178
179    /// Get total storage usage
180    pub fn get_storage_usage(&self) -> Result<u64> {
181        if !self.location.models_dir.exists() {
182            return Ok(0);
183        }
184
185        let mut total_size = 0u64;
186
187        for entry in walkdir::WalkDir::new(&self.location.models_dir) {
188            let entry = entry?;
189            if entry.file_type().is_file() {
190                total_size += entry.metadata()?.len();
191            }
192        }
193
194        Ok(total_size)
195    }
196
197    /// Remove a model from storage
198    pub fn remove_model(&self, model_id: &str) -> Result<()> {
199        let model_path = self.location.models_dir.join(sanitize_model_id(model_id));
200        let metadata_path = self.metadata_path(model_id);
201
202        if model_path.exists() {
203            std::fs::remove_dir_all(model_path)?;
204        }
205
206        if metadata_path.exists() {
207            std::fs::remove_file(metadata_path)?;
208        }
209
210        info!("Removed model: {}", model_id);
211        Ok(())
212    }
213
214    /// Create model directory structure
215    pub fn create_model_directory(&self, model_id: &str) -> Result<PathBuf> {
216        let model_dir = self.location.models_dir.join(sanitize_model_id(model_id));
217        std::fs::create_dir_all(&model_dir)?;
218        Ok(model_dir)
219    }
220
221    /// Get available disk space (cross-platform)
222    pub fn get_available_space(&self) -> Result<u64> {
223        let target_path = if self.location.models_dir.exists() {
224            self.location.models_dir.clone()
225        } else {
226            self.location.app_data_dir.clone()
227        };
228        let space = fs2::available_space(&target_path)?;
229        Ok(space)
230    }
231}
232
233/// Model metadata structure
234#[derive(Debug, Clone, Serialize, Deserialize)]
235pub struct ModelMetadata {
236    pub id: String,
237    pub name: String,
238    pub description: Option<String>,
239    pub author: Option<String>,
240    pub size_bytes: u64,
241    pub format: String,
242    pub download_source: String,
243    pub download_date: chrono::DateTime<chrono::Utc>,
244    pub last_used: Option<chrono::DateTime<chrono::Utc>>,
245    pub tags: Vec<String>,
246    pub hardware_requirements: HardwareRequirements,
247    pub compatibility_notes: Option<String>,
248    /// Platform-specific runtime binary paths
249    #[serde(default)]
250    pub runtime_binaries: std::collections::HashMap<String, std::path::PathBuf>,
251}
252
253/// Hardware requirements for a model
254#[derive(Debug, Clone, Serialize, Deserialize)]
255pub struct HardwareRequirements {
256    pub min_ram_gb: f32,
257    pub min_vram_gb: Option<f32>,
258    pub recommended_ram_gb: f32,
259    pub recommended_vram_gb: Option<f32>,
260    pub cpu_cores: Option<u32>,
261    pub gpu_required: bool,
262}
263
264impl Default for HardwareRequirements {
265    fn default() -> Self {
266        Self {
267            min_ram_gb: 4.0,
268            min_vram_gb: None,
269            recommended_ram_gb: 8.0,
270            recommended_vram_gb: None,
271            cpu_cores: None,
272            gpu_required: false,
273        }
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280    use tempfile::TempDir;
281
282    #[test]
283    fn test_storage_creation() -> Result<()> {
284        // Create a temporary directory for testing
285        let temp_dir = TempDir::new()?;
286        let test_base_dir = temp_dir.path().join("test_aud_io");
287
288        // Manually create the directory structure (simulating what ModelStorage::new() does)
289        let app_data_dir = test_base_dir.join("app_data");
290        let models_dir = app_data_dir.join("models");
291        let registry_dir = app_data_dir.join("registry");
292
293        std::fs::create_dir_all(&app_data_dir)?;
294        std::fs::create_dir_all(&models_dir)?;
295        std::fs::create_dir_all(&registry_dir)?;
296
297        // Create the storage struct with the pre-created directories
298        let storage = ModelStorage {
299            location: StorageLocation {
300                app_data_dir,
301                models_dir,
302                registry_dir,
303            },
304        };
305
306        // Test directory creation
307        assert!(storage.location.app_data_dir.exists());
308        assert!(storage.location.models_dir.exists());
309        assert!(storage.location.registry_dir.exists());
310
311        Ok(())
312    }
313
314    #[test]
315    fn test_model_path_generation() -> Result<()> {
316        let temp_dir = TempDir::new()?;
317        let storage = ModelStorage {
318            location: StorageLocation {
319                app_data_dir: temp_dir.path().to_path_buf(),
320                models_dir: temp_dir.path().join("models"),
321                registry_dir: temp_dir.path().join("registry"),
322            },
323        };
324
325        let path = storage.model_path("test-model", "model.gguf");
326        assert_eq!(
327            path,
328            temp_dir
329                .path()
330                .join("models")
331                .join("test-model")
332                .join("model.gguf")
333        );
334
335        Ok(())
336    }
337}
offline_intelligence/model_management/storage.rs

offline_intelligence/model_management/
storage.rs
