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

aither_models/
lib.rs
