rpi_cli/
model_registry.rs

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

rpi_cli/
model_registry.rs
