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 effort;
19mod registry;
20#[cfg(feature = "remote")]
21pub mod remote;
22pub(crate) mod tier;
23pub mod types;
24
25pub use aither_core::llm::model::Ability;
26pub use effort::default_reasoning_effort;
27pub use registry::ModelRegistry;
28pub use tier::ModelTier;
29pub use types::{ModelEntry, ModelMode, Pricing, Provider};
30
31/// Look up a model by `LiteLLM` ID or a dated/versioned alias
32/// (`-YYYYMMDD`, `-MMDD`, `vMAJOR-MINOR`).
33///
34/// Returns `None` when the ID is unknown to the registry — see
35/// [`ModelRegistry::lookup`] for the normalization rules.
36#[must_use]
37pub fn lookup(model_id: &str) -> Option<&'static ModelEntry> {
38    ModelRegistry::bundled().lookup(model_id)
39}
40
41/// Get all models for a provider.
42pub fn models_for_provider(provider: &Provider) -> impl Iterator<Item = &'static ModelEntry> + '_ {
43    ModelRegistry::bundled().models_for_provider(provider)
44}
45
46/// Get all models with a specific ability.
47pub fn models_with_ability(ability: Ability) -> impl Iterator<Item = &'static ModelEntry> {
48    ModelRegistry::bundled().models_with_ability(ability)
49}
50
51/// Get all models matching a mode.
52pub fn models_by_mode(mode: ModelMode) -> impl Iterator<Item = &'static ModelEntry> {
53    ModelRegistry::bundled().models_by_mode(mode)
54}
55
56/// Get all known models.
57pub fn all_models() -> impl Iterator<Item = &'static ModelEntry> {
58    ModelRegistry::bundled().all()
59}
60
61/// Capability/cost tier of a model in the bundled registry.
62///
63/// Unknown models classify as [`ModelTier::Balanced`].
64#[must_use]
65pub fn tier(model_id: &str) -> ModelTier {
66    ModelRegistry::bundled().tier(model_id)
67}
68
69#[cfg(test)]
70mod tests {
71    //! The bundled snapshot is refreshed from `LiteLLM` upstream, so these
72    //! tests only assert properties every valid snapshot has. Lookup and
73    //! tier behaviour is tested on a fixture in `registry` and `tier`.
74
75    use super::*;
76
77    #[test]
78    fn bundled_matches_runtime_conversion_of_snapshot() {
79        let runtime =
80            ModelRegistry::from_litellm_json(include_bytes!("../data/litellm_snapshot.json"));
81        let bundled: Vec<_> = all_models().map(ModelEntry::litellm_id).collect();
82        let converted: Vec<_> = runtime.all().map(ModelEntry::litellm_id).collect();
83        assert_eq!(bundled, converted);
84    }
85
86    #[test]
87    fn bundled_registry_is_populated() {
88        let registry = ModelRegistry::bundled();
89        assert!(
90            registry.len() > 100,
91            "Expected 100+ models, got {}",
92            registry.len()
93        );
94        assert_ne!(models_by_mode(ModelMode::Chat).count(), 0);
95    }
96
97    #[test]
98    fn every_entry_is_reachable_by_its_litellm_id() {
99        for entry in all_models() {
100            assert!(
101                lookup(entry.litellm_id()).is_some(),
102                "{} is not reachable by its own key",
103                entry.litellm_id()
104            );
105            assert!(
106                entry.litellm_id().ends_with(entry.id()),
107                "{} does not end with its canonical id {}",
108                entry.litellm_id(),
109                entry.id()
110            );
111        }
112    }
113
114    #[test]
115    fn query_helpers_return_only_matching_entries() {
116        assert!(models_with_ability(Ability::Vision).all(|m| m.has_ability(Ability::Vision)));
117        assert!(models_by_mode(ModelMode::Chat).all(|m| m.mode() == ModelMode::Chat));
118        assert!(models_for_provider(&Provider::OpenAI).all(|m| *m.provider() == Provider::OpenAI));
119    }
120
121    #[test]
122    fn image_generation_entries_classify_as_image_tier() {
123        for entry in models_by_mode(ModelMode::ImageGeneration) {
124            assert!(entry.has_ability(Ability::ImageGeneration));
125            assert_eq!(
126                ModelRegistry::bundled().tier_of(entry),
127                ModelTier::ImageGeneration
128            );
129        }
130    }
131}
aither_models/lib.rs

aither_models/
lib.rs
