offline_intelligence/model_management/
mod.rs

1//! Model Management System
2//!
3//! Provides comprehensive model lifecycle management including:
4//! - Model registry and metadata storage
5//! - Download from multiple sources (HuggingFace, OpenRouter)
6//! - Local storage management in AppData
7//! - Hardware-aware model recommendations
8//! - Progress tracking for downloads
9
10pub mod registry;
11pub mod downloader;
12pub mod storage;
13pub mod recommendation;
14pub mod progress;
15pub mod hf_access;
16
17pub use registry::{ModelRegistry, ModelInfo, ModelPricing, ModelStatus};
18pub use downloader::{ModelDownloader, DownloadSource};
19pub use storage::{ModelStorage, StorageLocation};
20pub use recommendation::ModelRecommender;
21pub use progress::{DownloadProgress, ProgressTracker};
22pub use hf_access::{check_hf_gated_access, HfAccessStatus};
23
24use anyhow::Result;
25use std::sync::Arc;
26use tokio::sync::RwLock;
27
28use crate::config::Config;
29
30/// Main model management service
31pub struct ModelManager {
32    pub registry: Arc<RwLock<ModelRegistry>>,
33    pub downloader: Arc<ModelDownloader>,
34    pub storage: Arc<ModelStorage>,
35    pub recommender: Arc<ModelRecommender>,
36}
37
38impl ModelManager {
39    pub fn new() -> Result<Self> {
40        let storage = Arc::new(ModelStorage::new()?);
41        let registry = Arc::new(RwLock::new(ModelRegistry::new(storage.clone())?));
42        let downloader = Arc::new(ModelDownloader::new(storage.clone()));
43        let recommender = Arc::new(ModelRecommender::new());
44
45        Ok(Self {
46            registry,
47            downloader,
48            storage,
49            recommender,
50        })
51    }
52
53    /// Initialize the model manager and scan for existing models
54    pub async fn initialize(&self, cfg: &Config) -> Result<()> {
55        // Scan storage for existing models and populate registry
56        self.registry.write().await.scan_storage().await?;
57
58        // Refresh model catalogs from remote sources on startup
59        self.refresh_catalogs(cfg).await?;
60
61        // After scanning storage, compute compatibility scores for local models
62        // based on the current hardware profile so that the UI can sort by
63        // "Best Match" using the compatibility_score field.
64        let hardware = ModelRecommender::detect_hardware_profile(cfg);
65        {
66            let mut registry = self.registry.write().await;
67            let recommender = &*self.recommender;
68            registry.update_compatibility_scores(recommender, &hardware);
69            // Persist updated registry so compatibility scores survive restarts
70            let _ = registry.save_registry().await;
71        }
72
73        Ok(())
74    }
75
76    /// Refresh model catalogs from remote sources (HuggingFace, OpenRouter)
77    pub async fn refresh_catalogs(&self, cfg: &Config) -> Result<()> {
78        let env_key = std::env::var("OPENROUTER_API_KEY").ok();
79        let api_key = if !cfg.openrouter_api_key.is_empty() && !cfg.openrouter_api_key.trim().is_empty() {
80            Some(cfg.openrouter_api_key.as_str())
81        } else if let Some(ref env_value) = env_key {
82            Some(env_value.as_str())
83        } else {
84            None
85        };
86
87        // Refresh OpenRouter catalog if API key is available
88        let mut registry = self.registry.write().await;
89        if let Some(key) = api_key {
90            if let Err(e) = registry.refresh_openrouter_catalog_from_api(key).await {
91                tracing::warn!("Failed to refresh OpenRouter catalog: {}", e);
92                // Continue with default catalog even if API refresh fails
93            }
94        } else {
95            // Load default popular OpenRouter models so users can see what's available
96            registry.populate_default_openrouter_models().await;
97        }
98        
99        if let Err(e) = registry.save_registry().await {
100            tracing::error!("Failed to save model registry after OpenRouter refresh: {}", e);
101        }
102
103        // Refresh Hugging Face GGUF/GGML catalog
104        // Fetch more models (500 instead of 100) and include more quantization options
105        if let Err(e) = registry.refresh_huggingface_catalog_from_api(500).await {
106            tracing::warn!("Failed to refresh Hugging Face catalog: {}", e);
107        }
108        if let Err(e) = registry.save_registry().await {
109            tracing::error!("Failed to save model registry after HuggingFace refresh: {}", e);
110        }
111
112        Ok(())
113    }
114}
offline_intelligence/model_management/mod.rs

offline_intelligence/model_management/
mod.rs
