use anyhow::Result; use futures_util::StreamExt; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::fs; use std::io::Write; use std::path::PathBuf; use std::sync::Mutex; use tauri::{App, AppHandle, Emitter, Manager}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelInfo { pub id: String, pub name: String, pub description: String, pub filename: String, pub url: Option, pub size_mb: u64, pub is_downloaded: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DownloadProgress { pub model_id: String, pub downloaded: u64, pub total: u64, pub percentage: f64, } pub struct ModelManager { app_handle: AppHandle, models_dir: PathBuf, available_models: Mutex>, } impl ModelManager { pub fn new(app: &App) -> Result { let app_handle = app.app_handle().clone(); // Create models directory in app data let models_dir = app .path() .app_data_dir() .map_err(|e| anyhow::anyhow!("Failed to get app data dir: {}", e))? .join("models"); if !models_dir.exists() { fs::create_dir_all(&models_dir)?; } let mut available_models = HashMap::new(); available_models.insert( "small".to_string(), ModelInfo { id: "small".to_string(), name: "Whisper Small".to_string(), description: "Fast and efficient, great for most use cases".to_string(), filename: "ggml-small.bin".to_string(), url: Some("https://blob.handy.computer/ggml-small.bin".to_string()), size_mb: 244, is_downloaded: false, }, ); // Add downloadable models available_models.insert( "medium".to_string(), ModelInfo { id: "medium".to_string(), name: "Whisper Medium".to_string(), description: "Good accuracy, medium speed".to_string(), filename: "ggml-medium-q5_0.bin".to_string(), url: Some("https://blob.handy.computer/ggml-medium-q5_0.bin".to_string()), size_mb: 539, // Approximate size is_downloaded: false, }, ); available_models.insert( "turbo".to_string(), ModelInfo { id: "turbo".to_string(), name: "Whisper Turbo".to_string(), description: "Good accuracy, medium speed".to_string(), filename: "ggml-large-v3-turbo-q5_0.bin".to_string(), url: Some("https://blob.handy.computer/ggml-large-v3-turbo-q5_0.bin".to_string()), size_mb: 574, // Approximate size is_downloaded: false, }, ); available_models.insert( "large".to_string(), ModelInfo { id: "large".to_string(), name: "Whisper Large".to_string(), description: "Highest accuracy, but slow.".to_string(), filename: "ggml-large-v3-q5_0.bin".to_string(), url: Some("https://blob.handy.computer/ggml-large-v3-q5_0.bin".to_string()), size_mb: 1080, // Approximate size is_downloaded: false, }, ); let manager = Self { app_handle, models_dir, available_models: Mutex::new(available_models), }; // Migrate any bundled models to user directory manager.migrate_bundled_models()?; // Check which models are already downloaded manager.update_download_status()?; Ok(manager) } pub fn get_available_models(&self) -> Vec { let models = self.available_models.lock().unwrap(); models.values().cloned().collect() } pub fn get_model_info(&self, model_id: &str) -> Option { let models = self.available_models.lock().unwrap(); models.get(model_id).cloned() } fn migrate_bundled_models(&self) -> Result<()> { // Check for bundled models and copy them to user directory let bundled_models = ["ggml-small.bin"]; // Add other bundled models here if any for filename in &bundled_models { let bundled_path = self.app_handle.path().resolve( &format!("resources/models/{}", filename), tauri::path::BaseDirectory::Resource, ); if let Ok(bundled_path) = bundled_path { if bundled_path.exists() { let user_path = self.models_dir.join(filename); // Only copy if user doesn't already have the model if !user_path.exists() { println!("Migrating bundled model {} to user directory", filename); fs::copy(&bundled_path, &user_path)?; println!("Successfully migrated {}", filename); } } } } Ok(()) } fn update_download_status(&self) -> Result<()> { let mut models = self.available_models.lock().unwrap(); for model in models.values_mut() { let model_path = self.models_dir.join(&model.filename); model.is_downloaded = model_path.exists(); } Ok(()) } pub async fn download_model(&self, model_id: &str) -> Result<()> { let model_info = { let models = self.available_models.lock().unwrap(); models.get(model_id).cloned() }; let model_info = model_info.ok_or_else(|| anyhow::anyhow!("Model not found: {}", model_id))?; let url = model_info .url .ok_or_else(|| anyhow::anyhow!("No download URL for model"))?; let model_path = self.models_dir.join(&model_info.filename); // Don't download if downloaded version already exists if model_path.exists() { // Update status and return self.update_download_status()?; return Ok(()); } println!("Downloading model {} from {}", model_id, url); // Create HTTP client let client = reqwest::Client::new(); let response = client.get(&url).send().await?; if !response.status().is_success() { return Err(anyhow::anyhow!( "Failed to download model: HTTP {}", response.status() )); } let total_size = response.content_length().unwrap_or(0); let mut downloaded = 0u64; let mut stream = response.bytes_stream(); // Create the file let mut file = std::fs::File::create(&model_path)?; // Download with progress while let Some(chunk) = stream.next().await { let chunk = chunk?; file.write_all(&chunk)?; downloaded += chunk.len() as u64; let percentage = if total_size > 0 { (downloaded as f64 / total_size as f64) * 100.0 } else { 0.0 }; // Emit progress event let progress = DownloadProgress { model_id: model_id.to_string(), downloaded, total: total_size, percentage, }; let _ = self.app_handle.emit("model-download-progress", &progress); } file.flush()?; // Update download status self.update_download_status()?; // Emit completion event let _ = self.app_handle.emit("model-download-complete", model_id); println!( "Successfully downloaded model {} to {:?}", model_id, model_path ); Ok(()) } pub fn delete_model(&self, model_id: &str) -> Result<()> { println!("ModelManager: delete_model called for: {}", model_id); let model_info = { let models = self.available_models.lock().unwrap(); models.get(model_id).cloned() }; let model_info = model_info.ok_or_else(|| anyhow::anyhow!("Model not found: {}", model_id))?; println!("ModelManager: Found model info: {:?}", model_info); let model_path = self.models_dir.join(&model_info.filename); println!("ModelManager: Model path: {:?}", model_path); if model_path.exists() { println!("ModelManager: Deleting model file at: {:?}", model_path); fs::remove_file(&model_path)?; println!("ModelManager: Model file deleted successfully"); } else { return Err(anyhow::anyhow!("Model file not found")); } // Update download status self.update_download_status()?; println!("ModelManager: Download status updated"); Ok(()) } pub fn get_model_path(&self, model_id: &str) -> Result { let model_info = self .get_model_info(model_id) .ok_or_else(|| anyhow::anyhow!("Model not found: {}", model_id))?; if !model_info.is_downloaded { return Err(anyhow::anyhow!("Model not available: {}", model_id)); } let model_path = self.models_dir.join(&model_info.filename); if model_path.exists() { Ok(model_path) } else { Err(anyhow::anyhow!("Model file not found: {}", model_id)) } } }