rpi_cli/
model_registry.rs

1//! Model registry runtime - dynamic model registration and refresh.
2//!
3//! Mirrors the native Pi `ModelRegistry` and `ModelRuntime` functionality:
4//! - Dynamic provider registration/unregistration
5//! - Runtime model catalog refresh
6//! - Provider authentication status tracking
7//! - Extension provider integration
8//! - Real-time catalog change notifications
9
10use std::collections::HashMap;
11use std::sync::{Arc, RwLock};
12use tokio::sync::broadcast;
13
14use rpi_ai::model::Model;
15use rpi_ai::Provider;
16
17/// Model registry event types
18#[derive(Debug, Clone)]
19pub enum ModelRegistryEvent {
20    /// A new provider was registered
21    ProviderRegistered { provider_id: String },
22    /// A provider was unregistered
23    ProviderUnregistered { provider_id: String },
24    /// Model catalog was refreshed
25    CatalogRefreshed,
26    /// Provider auth status changed
27    AuthStatusChanged { provider_id: String, authenticated: bool },
28}
29
30/// Provider authentication status
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub enum AuthStatus {
33    /// Not authenticated
34    Unauthenticated,
35    /// Authenticated with valid credentials
36    Authenticated,
37    /// Authentication expired/invalid
38    Expired,
39    /// Authentication status unknown
40    Unknown,
41}
42
43/// Provider information in the registry
44#[derive(Debug, Clone)]
45pub struct ProviderInfo {
46    /// Provider ID
47    pub id: String,
48    /// Display name
49    pub name: String,
50    /// Whether this is a built-in provider
51    pub builtin: bool,
52    /// Authentication status
53    pub auth_status: AuthStatus,
54    /// Available models
55    pub models: Vec<Model>,
56}
57
58/// Model registry - manages providers and models at runtime
59pub struct ModelRegistry {
60    /// Registered providers (metadata)
61    providers: Arc<RwLock<HashMap<String, ProviderInfo>>>,
62    /// Provider instances (for actual API calls)
63    provider_instances: Arc<RwLock<HashMap<String, Arc<dyn Provider>>>>,
64    /// Event broadcaster
65    event_tx: broadcast::Sender<ModelRegistryEvent>,
66    /// Models.json last refresh time
67    last_refresh: Arc<RwLock<Option<std::time::Instant>>>,
68}
69
70impl ModelRegistry {
71    /// Create a new model registry
72    pub fn new() -> Self {
73        let (event_tx, _) = broadcast::channel(100);
74        Self {
75            providers: Arc::new(RwLock::new(HashMap::new())),
76            provider_instances: Arc::new(RwLock::new(HashMap::new())),
77            event_tx,
78            last_refresh: Arc::new(RwLock::new(None)),
79        }
80    }
81
82    /// Register a built-in provider
83    pub fn register_builtin_provider(
84        &self,
85        id: String,
86        name: String,
87        models: Vec<Model>,
88    ) -> Result<(), RegistryError> {
89        let info = ProviderInfo {
90            id: id.clone(),
91            name,
92            builtin: true,
93            auth_status: AuthStatus::Unknown,
94            models,
95        };
96
97        let mut providers = self.providers.write().map_err(|_| RegistryError::LockPoisoned)?;
98        providers.insert(id.clone(), info);
99
100        let _ = self.event_tx.send(ModelRegistryEvent::ProviderRegistered {
101            provider_id: id,
102        });
103
104        Ok(())
105    }
106
107    /// Register an extension provider
108    pub fn register_extension_provider(
109        &self,
110        provider: Arc<dyn Provider>,
111    ) -> Result<(), RegistryError> {
112        let id = provider.id().to_string();
113        let models = provider.models().to_vec();
114        
115        let info = ProviderInfo {
116            id: id.clone(),
117            name: id.clone(), // Extension providers use id as name
118            builtin: false,
119            auth_status: AuthStatus::Unknown,
120            models,
121        };
122
123        let mut providers = self.providers.write().map_err(|_| RegistryError::LockPoisoned)?;
124        providers.insert(id.clone(), info);
125
126        let mut instances = self.provider_instances.write().map_err(|_| RegistryError::LockPoisoned)?;
127        instances.insert(id.clone(), provider);
128
129        let _ = self.event_tx.send(ModelRegistryEvent::ProviderRegistered {
130            provider_id: id,
131        });
132
133        Ok(())
134    }
135
136    /// Unregister a provider
137    pub fn unregister_provider(&self, provider_id: &str) -> Result<(), RegistryError> {
138        let mut providers = self.providers.write().map_err(|_| RegistryError::LockPoisoned)?;
139        let mut instances = self.provider_instances.write().map_err(|_| RegistryError::LockPoisoned)?;
140        
141        if providers.remove(provider_id).is_some() {
142            instances.remove(provider_id);
143            let _ = self.event_tx.send(ModelRegistryEvent::ProviderUnregistered {
144                provider_id: provider_id.to_string(),
145            });
146            Ok(())
147        } else {
148            Err(RegistryError::ProviderNotFound(provider_id.to_string()))
149        }
150    }
151
152    /// Get provider info
153    pub fn get_provider(&self, provider_id: &str) -> Option<ProviderInfo> {
154        let providers = self.providers.read().ok()?;
155        providers.get(provider_id).cloned()
156    }
157
158    /// Get all providers
159    pub fn list_providers(&self) -> Vec<ProviderInfo> {
160        let providers = self.providers.read().ok().map(|p| p.values().cloned().collect());
161        providers.unwrap_or_default()
162    }
163
164    /// Get all models across all providers
165    pub fn list_models(&self) -> Vec<Model> {
166        let providers = match self.providers.read() {
167            Ok(p) => p,
168            Err(_) => return Vec::new(),
169        };
170
171        providers
172            .values()
173            .flat_map(|info| info.models.clone())
174            .collect()
175    }
176
177    /// Get models for a specific provider
178    pub fn list_provider_models(&self, provider_id: &str) -> Vec<Model> {
179        let providers = match self.providers.read() {
180            Ok(p) => p,
181            Err(_) => return Vec::new(),
182        };
183
184        providers
185            .get(provider_id)
186            .map(|info| info.models.clone())
187            .unwrap_or_default()
188    }
189
190    /// Update provider authentication status
191    pub fn update_auth_status(
192        &self,
193        provider_id: &str,
194        status: AuthStatus,
195    ) -> Result<(), RegistryError> {
196        let mut providers = self.providers.write().map_err(|_| RegistryError::LockPoisoned)?;
197        
198        if let Some(info) = providers.get_mut(provider_id) {
199            info.auth_status = status.clone();
200            
201            let authenticated = matches!(status, AuthStatus::Authenticated);
202            let _ = self.event_tx.send(ModelRegistryEvent::AuthStatusChanged {
203                provider_id: provider_id.to_string(),
204                authenticated,
205            });
206            
207            Ok(())
208        } else {
209            Err(RegistryError::ProviderNotFound(provider_id.to_string()))
210        }
211    }
212
213    /// Get provider authentication status
214    pub fn get_auth_status(&self, provider_id: &str) -> Option<AuthStatus> {
215        let providers = self.providers.read().ok()?;
216        providers.get(provider_id).map(|info| info.auth_status.clone())
217    }
218
219    /// Refresh model catalog from models.json
220    pub async fn refresh_catalog(&self) -> Result<(), RegistryError> {
221        // Read models.json from config directory
222        let config = crate::config::load_models_config()
223            .map_err(|e| RegistryError::InvalidConfiguration(e.to_string()))?;
224        
225        // Parse and validate models
226        let mut providers = self.providers.write().map_err(|_| RegistryError::LockPoisoned)?;
227        
228        // Update existing providers and register providers that appear only in
229        // models.json (native `models.json` custom providers). Previously only
230        // pre-registered providers were updated, so a user-defined provider was
231        // silently ignored.
232        for (provider_id, provider_config) in &config.providers {
233            if let Some(models) = crate::config::provider_to_models(provider_id, provider_config) {
234                if let Some(info) = providers.get_mut(provider_id) {
235                    info.models = models;
236                    continue;
237                }
238            }
239            // New provider: register it from models.json.
240            let models = crate::config::provider_to_models(provider_id, provider_config)
241                .unwrap_or_default();
242            providers.insert(
243                provider_id.clone(),
244                ProviderInfo {
245                    id: provider_id.clone(),
246                    name: provider_config
247                        .name
248                        .clone()
249                        .unwrap_or_else(|| provider_id.clone()),
250                    builtin: false,
251                    auth_status: AuthStatus::Unknown,
252                    models,
253                },
254            );
255        }
256        
257        drop(providers);
258        
259        // Update last refresh time
260        let mut last_refresh = self.last_refresh.write().map_err(|_| RegistryError::LockPoisoned)?;
261        *last_refresh = Some(std::time::Instant::now());
262        
263        // Broadcast CatalogRefreshed event
264        let _ = self.event_tx.send(ModelRegistryEvent::CatalogRefreshed);
265
266        Ok(())
267    }
268
269    /// Get last catalog refresh time
270    pub fn last_refresh_time(&self) -> Option<std::time::Instant> {
271        self.last_refresh.read().ok().and_then(|t| *t)
272    }
273
274    /// Subscribe to registry events
275    pub fn subscribe(&self) -> broadcast::Receiver<ModelRegistryEvent> {
276        self.event_tx.subscribe()
277    }
278
279    /// Get provider instance for API calls
280    pub fn get_provider_instance(&self, provider_id: &str) -> Option<Arc<dyn Provider>> {
281        let instances = self.provider_instances.read().ok()?;
282        instances.get(provider_id).cloned()
283    }
284
285    /// Check if a provider is registered
286    pub fn has_provider(&self, provider_id: &str) -> bool {
287        let providers = match self.providers.read() {
288            Ok(p) => p,
289            Err(_) => return false,
290        };
291        providers.contains_key(provider_id)
292    }
293
294    /// Get provider count
295    pub fn provider_count(&self) -> usize {
296        self.providers.read().map(|p| p.len()).unwrap_or(0)
297    }
298
299    /// Get model count
300    pub fn model_count(&self) -> usize {
301        let providers = match self.providers.read() {
302            Ok(p) => p,
303            Err(_) => return 0,
304        };
305        providers.values().map(|info| info.models.len()).sum()
306    }
307}
308
309impl Default for ModelRegistry {
310    fn default() -> Self {
311        Self::new()
312    }
313}
314
315/// Registry errors
316#[derive(Debug, thiserror::Error)]
317pub enum RegistryError {
318    #[error("Lock poisoned")]
319    LockPoisoned,
320    
321    #[error("Provider not found: {0}")]
322    ProviderNotFound(String),
323    
324    #[error("Provider already registered: {0}")]
325    ProviderAlreadyRegistered(String),
326    
327    #[error("Invalid provider configuration: {0}")]
328    InvalidConfiguration(String),
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334
335    #[test]
336    fn test_registry_creation() {
337        let registry = ModelRegistry::new();
338        assert_eq!(registry.provider_count(), 0);
339        assert_eq!(registry.model_count(), 0);
340    }
341
342    #[test]
343    fn test_provider_registration() {
344        let registry = ModelRegistry::new();
345        
346        assert!(registry.register_builtin_provider(
347            "test".to_string(),
348            "Test Provider".to_string(),
349            vec![]
350        ).is_ok());
351        assert_eq!(registry.provider_count(), 1);
352        assert!(registry.has_provider("test"));
353    }
354
355    #[test]
356    fn test_provider_unregistration() {
357        let registry = ModelRegistry::new();
358        
359        registry.register_builtin_provider(
360            "test".to_string(),
361            "Test Provider".to_string(),
362            vec![]
363        ).unwrap();
364        assert!(registry.unregister_provider("test").is_ok());
365        assert_eq!(registry.provider_count(), 0);
366    }
367
368    #[test]
369    fn test_auth_status_update() {
370        let registry = ModelRegistry::new();
371        
372        registry.register_builtin_provider(
373            "test".to_string(),
374            "Test Provider".to_string(),
375            vec![]
376        ).unwrap();
377        
378        assert!(registry.update_auth_status("test", AuthStatus::Authenticated).is_ok());
379        assert_eq!(
380            registry.get_auth_status("test"),
381            Some(AuthStatus::Authenticated)
382        );
383    }
384}
rpi_cli/model_registry.rs

rpi_cli/
model_registry.rs
