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 `LiteLLM` ID or a dated/versioned alias.
30///
31/// Returns `None` when the ID is unknown to the registry — see
32/// [`ModelRegistry::lookup`] for the normalization rules.
33#[must_use]
34pub fn lookup(model_id: &str) -> Option<&'static ModelEntry> {
35    ModelRegistry::bundled().lookup(model_id)
36}
37
38/// Get all models for a provider.
39pub fn models_for_provider(provider: &Provider) -> impl Iterator<Item = &'static ModelEntry> + '_ {
40    ModelRegistry::bundled().models_for_provider(provider)
41}
42
43/// Get all models with a specific ability.
44pub fn models_with_ability(ability: Ability) -> impl Iterator<Item = &'static ModelEntry> {
45    ModelRegistry::bundled().models_with_ability(ability)
46}
47
48/// Get all models matching a mode.
49pub fn models_by_mode(mode: ModelMode) -> impl Iterator<Item = &'static ModelEntry> {
50    ModelRegistry::bundled().models_by_mode(mode)
51}
52
53/// Get all known models.
54pub fn all_models() -> impl Iterator<Item = &'static ModelEntry> {
55    ModelRegistry::bundled().all()
56}
57
58/// Capability/cost tier of a model in the bundled registry.
59///
60/// Unknown models classify as [`ModelTier::Balanced`].
61#[must_use]
62pub fn tier(model_id: &str) -> ModelTier {
63    ModelRegistry::bundled().tier(model_id)
64}
65
66#[cfg(test)]
67mod tests {
68    use super::*;
69
70    #[test]
71    fn test_lookup_gpt4o() {
72        let entry = lookup("gpt-4o").expect("gpt-4o should exist");
73        assert_eq!(entry.id(), "gpt-4o");
74        assert!(entry.max_input_tokens().unwrap_or(0) >= 128_000);
75        assert!(entry.has_ability(Ability::Vision));
76        assert!(entry.has_ability(Ability::ToolUse));
77    }
78
79    #[test]
80    fn test_lookup_claude() {
81        // Canary ID: must name a model the committed snapshot really
82        // publishes — a retired ID is a lookup miss, not a fallback.
83        let entry = lookup("claude-sonnet-5-5").expect("claude-sonnet-5-5 should exist");
84        assert!(entry.max_input_tokens().unwrap_or(0) >= 200_000);
85        assert!(entry.has_ability(Ability::ToolUse));
86    }
87
88    #[test]
89    fn test_lookup_retired_id_returns_none() {
90        // `claude-sonnet-4-20250514` and bare `claude-sonnet-4` were pruned
91        // by LiteLLM upstream; the only remaining `claude-sonnet-4` record
92        // is GitHub Copilot's, and a bare query must not bind it.
93        assert!(lookup("claude-sonnet-4-20250514").is_none());
94        assert!(lookup("claude-sonnet-4").is_none());
95        // Provider-scoped queries still normalize inside their own namespace.
96        let entry = lookup("github_copilot/claude-sonnet-4-20250514")
97            .expect("provider-scoped dated alias should resolve");
98        assert_eq!(entry.litellm_id(), "github_copilot/claude-sonnet-4");
99    }
100
101    #[test]
102    fn test_lookup_unknown_id_returns_none() {
103        assert!(lookup("no-such-model-xyz").is_none());
104    }
105
106    #[test]
107    fn test_lookup_gemini() {
108        let entry = lookup("gemini/gemini-2.5-flash").expect("gemini-2.5-flash should exist");
109        assert!(entry.max_input_tokens().unwrap_or(0) >= 1_000_000);
110        // `gemini/` is Google's own namespace, so bare IDs resolve to
111        // first-party records there too.
112        assert!(lookup("gemini-robotics-er-2-preview").is_some());
113    }
114
115    #[test]
116    fn test_lookup_case_insensitive() {
117        let entry = lookup("GPT-4O").expect("case-insensitive lookup should work");
118        assert_eq!(entry.id(), "gpt-4o");
119    }
120
121    #[test]
122    fn test_lookup_dated_alias() {
123        // A dated snapshot ID exact-matches its own record, while a date
124        // suffix with no dedicated record normalizes to the base model.
125        assert!(lookup("gpt-4o-2024-05-13").is_some());
126        let entry = lookup("gpt-4o-21000101").expect("dated alias should resolve");
127        assert_eq!(entry.id(), "gpt-4o");
128    }
129
130    #[test]
131    fn test_models_with_ability() {
132        let vision_models: Vec<_> = models_with_ability(Ability::Vision).collect();
133        assert_ne!(vision_models.len(), 0);
134        assert!(vision_models.iter().all(|m| m.has_ability(Ability::Vision)));
135    }
136
137    #[test]
138    fn test_models_by_mode() {
139        let chat_models: Vec<_> = models_by_mode(ModelMode::Chat).collect();
140        assert_ne!(chat_models.len(), 0);
141        assert!(chat_models.iter().all(|m| m.mode() == ModelMode::Chat));
142    }
143
144    #[test]
145    fn test_registry_not_empty() {
146        let registry = ModelRegistry::bundled();
147        assert!(
148            registry.len() > 100,
149            "Expected 100+ models, got {}",
150            registry.len()
151        );
152    }
153
154    #[test]
155    fn test_pricing_populated() {
156        let entry = lookup("gpt-4o").expect("gpt-4o should exist");
157        assert!(entry.pricing().input_per_token() > 0.0);
158        assert!(entry.pricing().output_per_token() > 0.0);
159    }
160}
aither_models/lib.rs

aither_models/
lib.rs
