aptu_core/ai/
models.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! AI model registry and provider abstraction.
4//!
5//! This module provides a static registry of available AI models across multiple providers
6//! (`OpenRouter`, `Ollama`, `Mlx`). It enables:
7//! - Model discovery and filtering by provider
8//! - Default model selection for free tier
9//! - Model lookup by identifier for configuration validation
10//! - Extensibility for future providers
11//!
12//! # Examples
13//!
14//! ```
15//! use aptu_core::ai::models::{AiModel, ModelProvider};
16//!
17//! // Get all available models
18//! let models = AiModel::available_models();
19//! assert!(!models.is_empty());
20//!
21//! // Get default free model
22//! let default = AiModel::default_free();
23//! assert!(default.is_free);
24//!
25//! // Filter by provider
26//! let openrouter_models = AiModel::for_provider(ModelProvider::OpenRouter);
27//! assert!(!openrouter_models.is_empty());
28//!
29//! // Find model by identifier
30//! let model = AiModel::find_by_identifier("mistralai/devstral-2512:free");
31//! assert!(model.is_some());
32//! ```
33
34use serde::{Deserialize, Serialize};
35
36/// AI provider identifier.
37///
38/// Represents different AI service providers that Aptu can integrate with.
39/// Each provider has different capabilities, pricing, and deployment models.
40#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
41#[cfg_attr(feature = "uniffi", derive(uniffi::Enum))]
42pub enum ModelProvider {
43    /// `OpenRouter` - Unified API for multiple AI providers
44    /// Supports free and paid models from Mistral, Anthropic, xAI, and others.
45    OpenRouter,
46
47    /// `Ollama` - Local AI model runner
48    /// Runs models locally without API calls or costs.
49    Ollama,
50
51    /// `MLX` - Apple Silicon optimized models (future iOS support)
52    /// Runs models natively on iOS devices.
53    Mlx,
54}
55
56impl std::fmt::Display for ModelProvider {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        match self {
59            ModelProvider::OpenRouter => write!(f, "OpenRouter"),
60            ModelProvider::Ollama => write!(f, "Ollama"),
61            ModelProvider::Mlx => write!(f, "MLX"),
62        }
63    }
64}
65
66/// AI model metadata and configuration.
67///
68/// Represents a single AI model with its capabilities, pricing, and provider information.
69/// Used for model selection, validation, and UI display.
70#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
71#[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
72pub struct AiModel {
73    /// Human-readable model name for UI display
74    /// Example: "Devstral 2", "Claude Sonnet 4.5"
75    pub display_name: String,
76
77    /// Provider-specific model identifier
78    /// Used in API requests to specify which model to use.
79    /// Examples:
80    /// - `OpenRouter`: "mistralai/devstral-2512:free"
81    /// - `Ollama`: "mistral:7b"
82    pub identifier: String,
83
84    /// AI service provider
85    pub provider: ModelProvider,
86
87    /// Whether this model is free to use
88    /// Free models have no API cost (either free tier or local execution).
89    pub is_free: bool,
90
91    /// Maximum context window size in tokens
92    /// Determines how much input text the model can process.
93    pub context_window: u32,
94}
95
96impl AiModel {
97    /// Returns all available AI models across all providers.
98    ///
99    /// This is the authoritative registry of models that Aptu supports.
100    /// Models are organized by provider and tier (free/paid).
101    ///
102    /// # Returns
103    ///
104    /// A vector of all available models, sorted by provider and tier.
105    #[must_use]
106    pub fn available_models() -> Vec<AiModel> {
107        vec![
108            // ================================================================
109            // OpenRouter - Free Tier Models
110            // ================================================================
111            AiModel {
112                display_name: "Devstral 2".to_string(),
113                identifier: "mistralai/devstral-2512:free".to_string(),
114                provider: ModelProvider::OpenRouter,
115                is_free: true,
116                context_window: 262_000,
117            },
118            AiModel {
119                display_name: "Mistral Small 3.1".to_string(),
120                identifier: "mistralai/mistral-small-3.1-24b-instruct:free".to_string(),
121                provider: ModelProvider::OpenRouter,
122                is_free: true,
123                context_window: 128_000,
124            },
125            // ================================================================
126            // OpenRouter - Paid Tier Models
127            // ================================================================
128            AiModel {
129                display_name: "Grok Code Fast".to_string(),
130                identifier: "x-ai/grok-code-fast-1".to_string(),
131                provider: ModelProvider::OpenRouter,
132                is_free: false,
133                context_window: 256_000,
134            },
135            AiModel {
136                display_name: "Claude Sonnet 4.5".to_string(),
137                identifier: "anthropic/claude-sonnet-4.5".to_string(),
138                provider: ModelProvider::OpenRouter,
139                is_free: false,
140                context_window: 1_000_000,
141            },
142            // ================================================================
143            // Ollama - Local Models
144            // ================================================================
145            AiModel {
146                display_name: "Mistral 7B (Local)".to_string(),
147                identifier: "mistral:7b".to_string(),
148                provider: ModelProvider::Ollama,
149                is_free: true,
150                context_window: 32_000,
151            },
152        ]
153    }
154
155    /// Returns the default free model for new users.
156    ///
157    /// Selects the first free `OpenRouter` model from the registry.
158    /// This is the recommended starting point for users without API keys.
159    ///
160    /// # Panics
161    ///
162    /// Panics if no free `OpenRouter` models are available in the registry.
163    /// This should never happen in practice as the registry is hardcoded.
164    ///
165    /// # Returns
166    ///
167    /// The default free model (Devstral 2).
168    #[must_use]
169    pub fn default_free() -> AiModel {
170        Self::available_models()
171            .into_iter()
172            .find(|m| m.is_free && m.provider == ModelProvider::OpenRouter)
173            .expect("Registry must contain at least one free OpenRouter model")
174    }
175
176    /// Filters models by provider.
177    ///
178    /// Returns all models from a specific provider, useful for UI dropdowns
179    /// or provider-specific configuration.
180    ///
181    /// # Arguments
182    ///
183    /// * `provider` - The provider to filter by
184    ///
185    /// # Returns
186    ///
187    /// A vector of models from the specified provider, or empty if none exist.
188    #[must_use]
189    pub fn for_provider(provider: ModelProvider) -> Vec<AiModel> {
190        Self::available_models()
191            .into_iter()
192            .filter(|m| m.provider == provider)
193            .collect()
194    }
195
196    /// Finds a model by its identifier.
197    ///
198    /// Used for configuration validation and model lookup from user input.
199    /// Identifiers are case-sensitive and must match exactly.
200    ///
201    /// # Arguments
202    ///
203    /// * `identifier` - The model identifier to search for
204    ///
205    /// # Returns
206    ///
207    /// Some(model) if found, None otherwise.
208    ///
209    /// # Examples
210    ///
211    /// ```
212    /// use aptu_core::ai::models::AiModel;
213    ///
214    /// let model = AiModel::find_by_identifier("mistralai/devstral-2512:free");
215    /// assert!(model.is_some());
216    /// assert_eq!(model.unwrap().display_name, "Devstral 2");
217    /// ```
218    #[must_use]
219    pub fn find_by_identifier(identifier: &str) -> Option<AiModel> {
220        Self::available_models()
221            .into_iter()
222            .find(|m| m.identifier == identifier)
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn test_available_models_not_empty() {
232        let models = AiModel::available_models();
233        assert!(
234            !models.is_empty(),
235            "Registry must contain at least one model"
236        );
237    }
238
239    #[test]
240    fn test_available_models_have_unique_identifiers() {
241        let models = AiModel::available_models();
242        let mut identifiers = Vec::new();
243        for model in &models {
244            assert!(
245                !identifiers.contains(&model.identifier),
246                "Duplicate identifier: {}",
247                model.identifier
248            );
249            identifiers.push(model.identifier.clone());
250        }
251    }
252
253    #[test]
254    fn test_default_free_is_free() {
255        let model = AiModel::default_free();
256        assert!(model.is_free, "Default model must be free");
257    }
258
259    #[test]
260    fn test_default_free_is_openrouter() {
261        let model = AiModel::default_free();
262        assert_eq!(
263            model.provider,
264            ModelProvider::OpenRouter,
265            "Default model must be from OpenRouter"
266        );
267    }
268
269    #[test]
270    fn test_for_provider_openrouter() {
271        let models = AiModel::for_provider(ModelProvider::OpenRouter);
272        assert!(!models.is_empty(), "OpenRouter should have models");
273        assert!(
274            models
275                .iter()
276                .all(|m| m.provider == ModelProvider::OpenRouter),
277            "All returned models should be from OpenRouter"
278        );
279    }
280
281    #[test]
282    fn test_for_provider_ollama() {
283        let models = AiModel::for_provider(ModelProvider::Ollama);
284        assert!(!models.is_empty(), "Ollama should have models");
285        assert!(
286            models.iter().all(|m| m.provider == ModelProvider::Ollama),
287            "All returned models should be from Ollama"
288        );
289    }
290
291    #[test]
292    fn test_for_provider_mlx_empty() {
293        let models = AiModel::for_provider(ModelProvider::Mlx);
294        assert!(
295            models.is_empty(),
296            "MLX should have no models in Phase 1 (reserved for future)"
297        );
298    }
299
300    #[test]
301    fn test_find_by_identifier_devstral() {
302        let model = AiModel::find_by_identifier("mistralai/devstral-2512:free");
303        assert!(model.is_some(), "Should find Devstral model");
304        let model = model.unwrap();
305        assert_eq!(model.display_name, "Devstral 2");
306        assert!(model.is_free);
307    }
308
309    #[test]
310    fn test_find_by_identifier_claude() {
311        let model = AiModel::find_by_identifier("anthropic/claude-sonnet-4.5");
312        assert!(model.is_some(), "Should find Claude model");
313        let model = model.unwrap();
314        assert_eq!(model.display_name, "Claude Sonnet 4.5");
315        assert!(!model.is_free);
316    }
317
318    #[test]
319    fn test_find_by_identifier_not_found() {
320        let model = AiModel::find_by_identifier("nonexistent/model");
321        assert!(model.is_none(), "Should not find nonexistent model");
322    }
323
324    #[test]
325    fn test_find_by_identifier_case_sensitive() {
326        let model = AiModel::find_by_identifier("MISTRALAI/DEVSTRAL-2512:FREE");
327        assert!(
328            model.is_none(),
329            "Identifier lookup should be case-sensitive"
330        );
331    }
332
333    #[test]
334    fn test_model_provider_display() {
335        assert_eq!(ModelProvider::OpenRouter.to_string(), "OpenRouter");
336        assert_eq!(ModelProvider::Ollama.to_string(), "Ollama");
337        assert_eq!(ModelProvider::Mlx.to_string(), "MLX");
338    }
339
340    #[test]
341    fn test_free_models_have_reasonable_context() {
342        let free_models = AiModel::available_models()
343            .into_iter()
344            .filter(|m| m.is_free)
345            .collect::<Vec<_>>();
346
347        assert!(!free_models.is_empty(), "Should have free models");
348        for model in free_models {
349            assert!(
350                model.context_window >= 32_000,
351                "Free model {} should have at least 32K context",
352                model.display_name
353            );
354        }
355    }
356
357    #[test]
358    fn test_paid_models_have_larger_context() {
359        let paid_models = AiModel::available_models()
360            .into_iter()
361            .filter(|m| !m.is_free)
362            .collect::<Vec<_>>();
363
364        assert!(!paid_models.is_empty(), "Should have paid models");
365        for model in paid_models {
366            assert!(
367                model.context_window >= 256_000,
368                "Paid model {} should have at least 256K context",
369                model.display_name
370            );
371        }
372    }
373
374    #[test]
375    fn test_model_serialization() {
376        let model = AiModel::default_free();
377        let json = serde_json::to_string(&model).expect("Should serialize");
378        let deserialized: AiModel = serde_json::from_str(&json).expect("Should deserialize");
379        assert_eq!(model, deserialized);
380    }
381
382    #[test]
383    fn test_model_provider_serialization() {
384        let provider = ModelProvider::OpenRouter;
385        let json = serde_json::to_string(&provider).expect("Should serialize");
386        let deserialized: ModelProvider = serde_json::from_str(&json).expect("Should deserialize");
387        assert_eq!(provider, deserialized);
388    }
389}