rig_core/catalog/
mod.rs

1//! A catalog of model facts, keyed by provider and model id: context and
2//! output limits, input modalities, the reasoning and caching each model
3//! takes, and its prices.
4//!
5//! The built-in catalog is data: `cargo xtask catalog sync` generates it from
6//! models.dev, and the main providers' entries are reviewed by hand. A
7//! harness can read an override file in the same shape with
8//! [`Catalog::from_json`] and lay it over the built-in one with
9//! [`Catalog::merge`]. [`ModelSpec::validate`] checks a request's
10//! [`GenerationOptions`](crate::completion::GenerationOptions) against what
11//! the model takes before anything is sent.
12//!
13//! ```
14//! use rig_core::catalog::Catalog;
15//! use rig_core::completion::{Effort, GenerationOptions};
16//!
17//! let haiku = Catalog::builtin()
18//!     .resolve("anthropic/claude-haiku-4-5")
19//!     .ok_or("listed")?;
20//! let options = GenerationOptions::default().reasoning(Effort::High);
21//! assert!(haiku.validate(&options).is_err(), "Haiku 4.5 takes a budget");
22//! # Ok::<(), &str>(())
23//! ```
24
25mod row;
26mod spec;
27
28use std::collections::BTreeMap;
29use std::sync::OnceLock;
30
31use serde::Deserialize;
32
33use crate::providers::registry::{Format, ProviderId};
34use row::Row;
35
36#[doc(hidden)]
37pub use spec::Compat;
38pub use spec::{CacheSupport, Modalities, ModelSpec, Pricing, ReasoningSupport, Sampling};
39
40/// The checked-in catalog data, generated by `cargo xtask catalog sync`.
41const BUILTIN: &str = include_str!("models.json");
42
43/// Model facts by provider and model id.
44#[derive(Clone, Debug, Default)]
45pub struct Catalog {
46    /// Sorted by vendor, then model id.
47    entries: Vec<Entry>,
48}
49
50/// One model: its spec, and the row it was built from, which
51/// [`Catalog::merge`] lays an override on.
52#[derive(Clone, Debug)]
53struct Entry {
54    row: Row,
55    spec: ModelSpec,
56}
57
58/// Why catalog data did not load.
59#[non_exhaustive]
60#[derive(Debug, thiserror::Error)]
61pub enum CatalogError {
62    /// The text is not a models.dev-style catalog: an object of provider
63    /// keys, each with a `models` object keyed by model id.
64    #[error("not a models.dev-style catalog: {0}")]
65    Json(#[from] serde_json::Error),
66}
67
68/// A provider's section of the data. models.dev's other provider keys
69/// (`name`, `env`, `doc`, ...) are ignored.
70#[derive(Deserialize)]
71struct Section {
72    #[serde(default)]
73    models: BTreeMap<String, Row>,
74}
75
76/// The models.dev provider keys whose rig vendor name differs. Any other key
77/// is read as a rig vendor name.
78const KEYS: [(&str, &str); 9] = [
79    ("azure", "azure.openai"),
80    ("google", "gcp.gemini"),
81    ("google-vertex", "vertexai"),
82    ("amazon-bedrock", "aws_bedrock"),
83    ("togetherai", "together"),
84    ("moonshotai", "moonshot"),
85    ("xiaomi", "xiaomimimo"),
86    ("ollama-cloud", "ollama"),
87    ("github-copilot", "copilot"),
88];
89
90impl Catalog {
91    /// The catalog this build ships.
92    pub fn builtin() -> &'static Catalog {
93        static BUILTIN_CATALOG: OnceLock<Catalog> = OnceLock::new();
94        // The data is checked in, and a test proves it parses.
95        BUILTIN_CATALOG.get_or_init(|| Catalog::from_json(BUILTIN).unwrap_or_default())
96    }
97
98    /// The model `provider` serves as `model`, by exact id. Any selection of
99    /// a vendor finds that vendor's models.
100    pub fn get(&self, provider: ProviderId, model: &str) -> Option<&ModelSpec> {
101        self.exact(provider.vendor(), model)
102    }
103
104    /// The model a reference names: `vendor/model`
105    /// (`anthropic/claude-opus-5-5`, `openrouter/anthropic/claude-sonnet-4.5`)
106    /// or `vendor[/format]:model` as [`ProviderRef`] spells it. The second
107    /// form applies when the text before the first `:` is a vendor or a
108    /// `vendor/format` pair, so a model id holding a `:` (`ollama/qwen3:4b`)
109    /// still reads as the first.
110    ///
111    /// [`ProviderRef`]: crate::providers::registry::ProviderRef
112    pub fn resolve(&self, reference: &str) -> Option<&ModelSpec> {
113        let (vendor, model) = split_reference(reference)?;
114        self.exact(vendor, model)
115    }
116
117    /// Every model, by vendor then id.
118    pub fn iter(&self) -> impl Iterator<Item = &ModelSpec> {
119        self.entries.iter().map(|entry| &entry.spec)
120    }
121
122    /// Read a models.dev-style catalog: an object of provider keys, each
123    /// with a `models` object keyed by model id. A provider key is a
124    /// models.dev key (`google`, `amazon-bedrock`) or a rig vendor name
125    /// (`gcp.gemini`); a key that is neither is skipped, so models.dev's own
126    /// `api.json` reads whole. Rig's hand-entered facts go under each row's
127    /// `rig` object.
128    pub fn from_json(json: &str) -> Result<Catalog, CatalogError> {
129        let sections: BTreeMap<String, Section> = serde_json::from_str(json)?;
130        let mut catalog = Catalog::default();
131        for (key, section) in sections {
132            let Some(provider) = ProviderId::catalog(vendor_of(&key)) else {
133                continue;
134            };
135            for (id, row) in section.models {
136                catalog.put(provider, &id, row);
137            }
138        }
139        Ok(catalog)
140    }
141
142    /// This catalog with `overrides` laid over it. A model in both keeps
143    /// every field the override leaves out; a model only in `overrides` is
144    /// added.
145    pub fn merge(mut self, overrides: Catalog) -> Catalog {
146        for entry in overrides.entries {
147            self.put(entry.spec.provider, &entry.spec.id, entry.row);
148        }
149        self
150    }
151
152    /// The model `provider` serves as `model`, as the encoders look it up:
153    /// by exact id or, failing that, without a dated snapshot suffix
154    /// (`-20251001`, `-2025-08-07`), so `claude-sonnet-4-5-20250929` finds
155    /// `claude-sonnet-4-5`. [`Self::get`] matches the exact id only.
156    pub fn find(&self, provider: ProviderId, model: &str) -> Option<&ModelSpec> {
157        self.find_vendor(provider.vendor(), model)
158    }
159
160    /// [`Self::find`] by vendor name.
161    pub(crate) fn find_vendor(&self, vendor: &str, model: &str) -> Option<&ModelSpec> {
162        self.exact(vendor, model)
163            .or_else(|| self.exact(vendor, undated(model)?))
164    }
165
166    fn exact(&self, vendor: &str, model: &str) -> Option<&ModelSpec> {
167        self.position(vendor, model)
168            .ok()
169            .and_then(|index| self.entries.get(index))
170            .map(|entry| &entry.spec)
171    }
172
173    fn position(&self, vendor: &str, model: &str) -> Result<usize, usize> {
174        self.entries.binary_search_by(|entry| {
175            (entry.spec.provider.vendor(), entry.spec.id.as_str()).cmp(&(vendor, model))
176        })
177    }
178
179    /// Add `row` for `provider`'s `id`, laid over any row already there.
180    fn put(&mut self, provider: ProviderId, id: &str, row: Row) {
181        match self.position(provider.vendor(), id) {
182            Ok(index) => {
183                if let Some(entry) = self.entries.get_mut(index) {
184                    let row = std::mem::take(&mut entry.row).overlay(row);
185                    entry.spec = row.spec(entry.spec.provider, id);
186                    entry.row = row;
187                }
188            }
189            Err(index) => {
190                let spec = row.spec(provider, id);
191                self.entries.insert(index, Entry { row, spec });
192            }
193        }
194    }
195}
196
197/// The model of the built-in catalog `vendor` serves as `model`, looked up
198/// as the encoders look models up ([`Catalog::find`]).
199pub(crate) fn lookup(vendor: &str, model: &str) -> Option<&'static ModelSpec> {
200    Catalog::builtin().find_vendor(vendor, model)
201}
202
203/// The model of the built-in catalog `vendor` serves as `model`, or as a
204/// snapshot of it: a listed id followed by `-20` and anything
205/// (`claude-opus-4-1-20250805`, `claude-opus-5-5-20260601-v1:0`), the
206/// longest such id.
207pub(crate) fn lookup_snapshot(vendor: &str, model: &str) -> Option<&'static ModelSpec> {
208    let catalog = Catalog::builtin();
209    catalog.exact(vendor, model).or_else(|| {
210        model
211            .match_indices("-20")
212            .filter_map(|(at, _)| catalog.exact(vendor, model.get(..at)?))
213            .last()
214    })
215}
216
217/// Whether `vendor`'s `model` reads images: what its catalog entry lists,
218/// or, for a model the catalog does not list, what `rule` (the vendor's
219/// naming rule) says of its id. Every wire that filters images reads this,
220/// so an id the catalog does not list (a gateway's spelling, another case,
221/// a deployment name) keeps its vendor's naming rule.
222pub(crate) fn reads_images_or(vendor: &str, model: &str, rule: impl FnOnce(&str) -> bool) -> bool {
223    lookup(vendor, model).map_or_else(|| rule(model), |spec| spec.input.image)
224}
225
226/// The rig vendor a models.dev provider key names.
227fn vendor_of(key: &str) -> &str {
228    KEYS.iter()
229        .find_map(|(models_dev, vendor)| (*models_dev == key).then_some(*vendor))
230        .unwrap_or(key)
231}
232
233/// `(vendor, model)` from a reference, by the grammar [`Catalog::resolve`]
234/// documents.
235pub(crate) fn split_reference(reference: &str) -> Option<(&str, &str)> {
236    if let Some((selection, model)) = reference.split_once(':') {
237        let vendor = match selection.split_once('/') {
238            None => Some(selection),
239            Some((vendor, format)) => Format::named(format).map(|_| vendor),
240        };
241        if let Some(vendor) = vendor.filter(|vendor| !vendor.is_empty()) {
242            return (!model.is_empty()).then_some((vendor, model));
243        }
244    }
245    reference
246        .split_once('/')
247        .filter(|(vendor, model)| !vendor.is_empty() && !model.is_empty())
248}
249
250/// `model` without a trailing dated snapshot: `-YYYYMMDD` or `-YYYY-MM-DD`.
251fn undated(model: &str) -> Option<&str> {
252    let digits = |text: &str| !text.is_empty() && text.bytes().all(|byte| byte.is_ascii_digit());
253    let (rest, last) = model.rsplit_once('-')?;
254    if last.len() == 8 && digits(last) {
255        return Some(rest);
256    }
257    let (rest, month) = rest.rsplit_once('-')?;
258    let (rest, year) = rest.rsplit_once('-')?;
259    (year.len() == 4 && month.len() == 2 && last.len() == 2)
260        .then_some(())
261        .filter(|()| digits(year) && digits(month) && digits(last))
262        .map(|()| rest)
263}
264
265#[cfg(test)]
266mod tests;
rig_core/catalog/mod.rs

rig_core/catalog/
mod.rs
