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    //! The bundled snapshot is refreshed from `LiteLLM` upstream, so these
69    //! tests only assert properties every valid snapshot has. Lookup and
70    //! tier behaviour is tested on a fixture in `registry` and `tier`.
71
72    use super::*;
73
74    #[test]
75    fn bundled_matches_runtime_conversion_of_snapshot() {
76        let runtime =
77            ModelRegistry::from_litellm_json(include_bytes!("../data/litellm_snapshot.json"));
78        let bundled: Vec<_> = all_models().map(ModelEntry::litellm_id).collect();
79        let converted: Vec<_> = runtime.all().map(ModelEntry::litellm_id).collect();
80        assert_eq!(bundled, converted);
81    }
82
83    #[test]
84    fn bundled_registry_is_populated() {
85        let registry = ModelRegistry::bundled();
86        assert!(
87            registry.len() > 100,
88            "Expected 100+ models, got {}",
89            registry.len()
90        );
91        assert_ne!(models_by_mode(ModelMode::Chat).count(), 0);
92    }
93
94    #[test]
95    fn every_entry_is_reachable_by_its_litellm_id() {
96        for entry in all_models() {
97            assert!(
98                lookup(entry.litellm_id()).is_some(),
99                "{} is not reachable by its own key",
100                entry.litellm_id()
101            );
102            assert!(
103                entry.litellm_id().ends_with(entry.id()),
104                "{} does not end with its canonical id {}",
105                entry.litellm_id(),
106                entry.id()
107            );
108        }
109    }
110
111    #[test]
112    fn query_helpers_return_only_matching_entries() {
113        assert!(models_with_ability(Ability::Vision).all(|m| m.has_ability(Ability::Vision)));
114        assert!(models_by_mode(ModelMode::Chat).all(|m| m.mode() == ModelMode::Chat));
115        assert!(models_for_provider(&Provider::OpenAI).all(|m| *m.provider() == Provider::OpenAI));
116    }
117
118    #[test]
119    fn image_generation_entries_classify_as_image_tier() {
120        for entry in models_by_mode(ModelMode::ImageGeneration) {
121            assert!(entry.has_ability(Ability::ImageGeneration));
122            assert_eq!(
123                ModelRegistry::bundled().tier_of(entry),
124                ModelTier::ImageGeneration
125            );
126        }
127    }
128}
aither_models/lib.rs

aither_models/
lib.rs
