aither_models/
lib.rs

1//! Model registry with capabilities, context windows, and metadata for popular LLMs.
2//!
3//! This crate provides a database of known AI models sourced from `LiteLLM` upstream,
4//! with their capabilities, context window sizes, pricing, and provider information.
5//!
6//! # Example
7//!
8//! ```rust
9//! use aither_models::{lookup, Ability};
10//!
11//! if let Some(entry) = lookup("gpt-4o") {
12//!     assert!(entry.max_input_tokens().unwrap_or(0) > 0);
13//!     assert!(entry.has_ability(Ability::Vision));
14//! }
15//! ```
16
17pub mod convert;
18mod registry;
19#[cfg(feature = "remote")]
20pub mod remote;
21pub(crate) mod tier;
22pub mod types;
23
24pub use aither_core::llm::model::Ability;
25pub use registry::ModelRegistry;
26pub use tier::ModelTier;
27pub use types::{ModelEntry, ModelMode, Pricing, Provider};
28
29/// Look up a model by ID or alias.
30///
31/// Delegates to the bundled [`ModelRegistry`].
32#[must_use]
33pub fn lookup(model_id: &str) -> Option<&'static ModelEntry> {
34    ModelRegistry::bundled().lookup(model_id)
35}
36
37/// Get all models for a provider.
38pub fn models_for_provider(provider: &Provider) -> impl Iterator<Item = &'static ModelEntry> + '_ {
39    ModelRegistry::bundled().models_for_provider(provider)
40}
41
42/// Get all models with a specific ability.
43pub fn models_with_ability(ability: Ability) -> impl Iterator<Item = &'static ModelEntry> {
44    ModelRegistry::bundled().models_with_ability(ability)
45}
46
47/// Get all models matching a mode.
48pub fn models_by_mode(mode: ModelMode) -> impl Iterator<Item = &'static ModelEntry> {
49    ModelRegistry::bundled().models_by_mode(mode)
50}
51
52/// Get all known models.
53pub fn all_models() -> impl Iterator<Item = &'static ModelEntry> {
54    ModelRegistry::bundled().all()
55}
56
57/// Capability/cost tier of a model in the bundled registry.
58///
59/// Unknown models classify as [`ModelTier::Balanced`].
60#[must_use]
61pub fn tier(model_id: &str) -> ModelTier {
62    ModelRegistry::bundled().tier(model_id)
63}
64
65#[cfg(test)]
66mod tests {
67    use super::*;
68
69    #[test]
70    fn test_lookup_gpt4o() {
71        let entry = lookup("gpt-4o").expect("gpt-4o should exist");
72        assert_eq!(entry.id(), "gpt-4o");
73        assert!(entry.max_input_tokens().unwrap_or(0) >= 128_000);
74        assert!(entry.has_ability(Ability::Vision));
75        assert!(entry.has_ability(Ability::ToolUse));
76    }
77
78    #[test]
79    fn test_lookup_claude() {
80        let entry = lookup("claude-sonnet-4-20250514").expect("claude-sonnet-4 should exist");
81        assert!(entry.max_input_tokens().unwrap_or(0) >= 200_000);
82        assert!(entry.has_ability(Ability::ToolUse));
83    }
84
85    #[test]
86    fn test_lookup_gemini() {
87        let entry = lookup("gemini/gemini-2.5-flash").expect("gemini-2.5-flash should exist");
88        assert!(entry.max_input_tokens().unwrap_or(0) >= 1_000_000);
89    }
90
91    #[test]
92    fn test_lookup_case_insensitive() {
93        let entry = lookup("GPT-4O").expect("case-insensitive lookup should work");
94        assert_eq!(entry.id(), "gpt-4o");
95    }
96
97    #[test]
98    fn test_lookup_prefix_match() {
99        // A dated variant should match the base model via prefix
100        let entry = lookup("gpt-4o-2024-05-13");
101        assert!(entry.is_some());
102    }
103
104    #[test]
105    fn test_models_with_ability() {
106        let vision_models: Vec<_> = models_with_ability(Ability::Vision).collect();
107        assert!(!vision_models.is_empty());
108        assert!(vision_models.iter().all(|m| m.has_ability(Ability::Vision)));
109    }
110
111    #[test]
112    fn test_models_by_mode() {
113        let chat_models: Vec<_> = models_by_mode(ModelMode::Chat).collect();
114        assert!(!chat_models.is_empty());
115        assert!(chat_models.iter().all(|m| m.mode() == ModelMode::Chat));
116    }
117
118    #[test]
119    fn test_registry_not_empty() {
120        let registry = ModelRegistry::bundled();
121        assert!(
122            registry.len() > 100,
123            "Expected 100+ models, got {}",
124            registry.len()
125        );
126    }
127
128    #[test]
129    fn test_pricing_populated() {
130        let entry = lookup("gpt-4o").expect("gpt-4o should exist");
131        assert!(entry.pricing().input_per_token() > 0.0);
132        assert!(entry.pricing().output_per_token() > 0.0);
133    }
134}
aither_models/lib.rs

aither_models/
lib.rs
