1use serde::{Deserialize, Serialize};
35
36#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
41#[cfg_attr(feature = "uniffi", derive(uniffi::Enum))]
42pub enum ModelProvider {
43 OpenRouter,
46
47 Ollama,
50
51 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#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
71#[cfg_attr(feature = "uniffi", derive(uniffi::Record))]
72pub struct AiModel {
73 pub display_name: String,
76
77 pub identifier: String,
83
84 pub provider: ModelProvider,
86
87 pub is_free: bool,
90
91 pub context_window: u32,
94}
95
96impl AiModel {
97 #[must_use]
106 pub fn available_models() -> Vec<AiModel> {
107 vec![
108 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 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 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 #[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 #[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 #[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}