rig_core/catalog/

mod.rs

//! Model facts by provider and model id: limits, input modalities, reasoning, caching, prices.
//!
//! A catalog is built in layers, each sharing the rows it does not change:
//! 1. [`Catalog::builtin`]: generated from models.dev by `cargo xtask catalog sync`, its main
//!    providers reviewed by hand;
//! 2. a newer models.dev copy ([`Catalog::from_models_dev`], lenient; laid on with
//!    [`Catalog::with_models_dev`], which keeps the hand-reviewed fields);
//! 3. the user's overrides ([`Catalog::from_overrides`], strict, with a close alternative for
//!    every mistake; laid on with [`Catalog::with_overrides`]).
//!
//! A model connected earlier keeps the facts it was connected with.
//!
//! ```no_run
//! use std::path::Path;
//! use rig_core::catalog::Catalog;
//!
//! fn catalog(cache: &Path, config: &Path) -> Result<Catalog, Box<dyn std::error::Error>> {
//!     let mut catalog = Catalog::builtin().clone();
//!     let fetched = cache.join("models-dev.json");
//!     if std::fs::metadata(&fetched)
//!         .and_then(|file| file.modified())
//!         .is_ok_and(|fetched_at| fetched_at > Catalog::generated_at())
//!     {
//!         let (models_dev, _skipped) = Catalog::from_models_dev(&std::fs::read_to_string(&fetched)?)?;
//!         catalog = catalog.with_models_dev(&models_dev);
//!     }
//!     if let Ok(json) = std::fs::read_to_string(config.join("models.json")) {
//!         catalog = catalog.with_overrides(&Catalog::from_overrides(&json, &catalog)?);
//!     }
//!     Ok(catalog)
//! }
//! ```
//!
//! [`Catalog::connect`] builds a model from `vendor/model` with the key from the environment; it
//! encodes and prices with the [`ModelFacts`] it was connected with. [`ModelSpec::validate`]
//! checks a request's options before anything is sent:
//!
//! ```
//! use rig_core::catalog::Catalog;
//! use rig_core::completion::{Effort, GenerationOptions};
//!
//! let haiku = Catalog::builtin().resolve("anthropic/claude-haiku-4-5")?.spec;
//! let options = GenerationOptions::default().reasoning(Effort::High);
//! assert!(haiku.validate(&options).is_err(), "Haiku 4.5 takes a budget");
//! # Ok::<(), rig_core::catalog::NotFound>(())
//! ```

mod connector;
mod facts;
mod lookup;
mod read;
mod row;
mod spec;

use std::sync::{Arc, OnceLock};
use std::time::{Duration, SystemTime, UNIX_EPOCH};

use crate::providers::registry::{Format, ProviderId};
use row::Row;

pub use connector::{Connector, ModelStatus, Provide, SignIn, Transport};
pub use facts::ModelFacts;
pub use lookup::{Matched, NotFound, Resolved};
pub use read::{OverrideError, OverrideErrorKind, OverrideErrors, SkipReason, Skipped};
pub use spec::{
    CacheSupport, Compat, Modalities, ModelSpec, Pricing, ReasoningChoice, ReasoningSupport,
    Sampling, UnknownReasoning,
};

/// The checked-in catalog data, generated by `cargo xtask catalog sync`.
const BUILTIN: &str = include_str!("models.json");

/// When `cargo xtask catalog sync` read the models.dev data [`BUILTIN`] was
/// generated from, in seconds since the Unix epoch. The sync writes it as
/// an RFC 3339 UTC time; a file in any other form fails the build.
const GENERATED_AT: u64 = match unix_seconds(include_str!("generated_at").trim_ascii_end()) {
    Some(seconds) => seconds,
    #[allow(clippy::panic)]
    None => panic!("catalog/generated_at is not a `YYYY-MM-DDTHH:MM:SSZ` time"),
};

/// Model facts by provider and model id. Cloning is cheap: the rows are
/// shared, and a change copies only the list of handles to them.
#[derive(Clone, Debug, Default)]
pub struct Catalog {
    /// Sorted by vendor, then model id.
    entries: Arc<Vec<Entry>>,
}

/// One model: its spec, and the row it was built from, which
/// [`Catalog::with_overrides`] lays an override on.
#[derive(Clone, Debug)]
struct Entry {
    row: Arc<Row>,
    spec: Arc<ModelSpec>,
}

/// Why catalog data did not load.
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum CatalogError {
    /// The text is not a JSON object.
    #[error("not a models.dev-style catalog: {0}")]
    Json(#[from] serde_json::Error),
}

/// The models.dev provider keys whose rig vendor name differs. Any other key
/// is read as a rig vendor name.
const KEYS: [(&str, &str); 9] = [
    ("azure", "azure.openai"),
    ("google", "gcp.gemini"),
    ("google-vertex", "vertexai"),
    ("amazon-bedrock", "aws_bedrock"),
    ("togetherai", "together"),
    ("moonshotai", "moonshot"),
    ("xiaomi", "xiaomimimo"),
    ("ollama-cloud", "ollama"),
    ("github-copilot", "copilot"),
];

impl Catalog {
    /// The catalog this build ships.
    ///
    /// # Panics
    ///
    /// Panics on first use if the data compiled into this build does not
    /// parse whole. The data is checked in and a unit test parses it, so a build
    /// that passed its tests never panics here.
    pub fn builtin() -> &'static Catalog {
        static BUILTIN_CATALOG: OnceLock<Catalog> = OnceLock::new();
        BUILTIN_CATALOG.get_or_init(|| match Catalog::from_models_dev(BUILTIN) {
            Ok((catalog, skipped)) if skipped.is_empty() => catalog,
            // A missing row would silently drop that model's facts, and an
            // empty catalog every model's, so a bad build fails loudly
            // instead. `the_checked_in_catalog_loads_whole` keeps these
            // branches unreachable in a tested build.
            #[allow(clippy::panic)]
            Ok((_, skipped)) => panic!(
                "rig-core's built-in catalog data skips {}: {:?}",
                skipped.len(),
                skipped.first()
            ),
            #[allow(clippy::panic)]
            Err(error) => panic!("rig-core's built-in catalog data does not parse: {error}"),
        })
    }

    /// When the models.dev data behind [`Self::builtin`] was read. Lay a
    /// fetched copy on only when it is newer, so an old cache never hides
    /// what this rig shipped.
    pub fn generated_at() -> SystemTime {
        UNIX_EPOCH + Duration::from_secs(GENERATED_AT)
    }

    /// The model `provider` serves as `model`, by the one rule every
    /// lookup in rig follows: the id as listed, or else the longest listed
    /// id that `model` extends with `-20` and a year, a dated snapshot of it
    /// (`claude-sonnet-4-5-20250929`, `gpt-5-2025-08-07`,
    /// `claude-opus-5-5-20260601-v1:0`). Any selection of a vendor finds
    /// that vendor's models.
    pub fn get(&self, provider: ProviderId, model: &str) -> Option<Resolved<'_>> {
        self.get_vendor(provider.vendor(), model)
    }

    /// The model `provider` lists as exactly `model`, for a caller that
    /// must not read a snapshot as its base model.
    pub fn get_exact(&self, provider: ProviderId, model: &str) -> Option<&ModelSpec> {
        self.exact(provider.vendor(), model)
    }

    /// The model a reference names, by [`Self::get`]'s rule. A reference is
    /// `vendor/model`, split at the first `/`, so the model id keeps any
    /// `/` and `:` of its own (`openrouter/anthropic/claude-sonnet-4.5`,
    /// `ollama/qwen3:4b`). The older `vendor/format:model` spelling is still
    /// read; its format plays no part in the lookup.
    ///
    /// # Errors
    ///
    /// [`NotFound`] when the reference names no listed model, with up to
    /// five close references to models that are not deprecated, or else
    /// the closest ones of its vendor.
    pub fn resolve(&self, reference: &str) -> Result<Resolved<'_>, NotFound> {
        Reference::parse(reference)
            .and_then(|parsed| self.get_vendor(parsed.vendor, parsed.model))
            .ok_or_else(|| self.not_found(reference))
    }

    /// The miss [`Self::resolve`] reports for `reference`.
    pub(crate) fn not_found(&self, reference: &str) -> NotFound {
        let mut suggestions = lookup::suggestions(reference, self.iter());
        if suggestions.is_empty() {
            suggestions = lookup::vendor_models(reference, self.iter());
        }
        NotFound {
            reference: reference.to_owned(),
            suggestions,
        }
    }

    /// Every model, by vendor then id.
    pub fn iter(&self) -> impl Iterator<Item = &ModelSpec> {
        self.entries.iter().map(|entry| &*entry.spec)
    }

    /// This catalog with `overrides` laid over it, field by field. A model in
    /// both keeps every field the override leaves out; a model only in
    /// `overrides` is added. `self` is unchanged, and the result shares every
    /// row the override does not touch.
    ///
    /// Every field the override sets wins, including the facts rig reviewed
    /// by hand, so lay a fetched copy of models.dev on with
    /// [`Self::with_models_dev`] instead.
    pub fn with_overrides(&self, overrides: &Catalog) -> Catalog {
        let mut catalog = self.clone();
        for entry in overrides.entries.iter() {
            catalog.put(entry.spec.provider, &entry.spec.id, Row::clone(&entry.row));
        }
        catalog
    }

    /// This catalog with `fetched`, a newer copy of models.dev read with
    /// [`Self::from_models_dev`], laid over it as [`Self::with_overrides`]
    /// lays an override, except for the fields of a row that rig's own data
    /// pins. The built-in data pins the fields rig reviewed by hand and the
    /// ones its generator joins from another row (a gateway's Claude row
    /// takes Anthropic's reasoning options), so a refresh updates prices and
    /// limits models.dev changes but keeps those. A model only in `fetched`
    /// is added as models.dev lists it.
    pub fn with_models_dev(&self, fetched: &Catalog) -> Catalog {
        let mut catalog = self.clone();
        for entry in fetched.entries.iter() {
            let (provider, id) = (entry.spec.provider, entry.spec.id.as_str());
            let row = Row::clone(&entry.row);
            let row = match catalog.position(provider.vendor(), id) {
                Ok(index) => match catalog.entries.get(index) {
                    Some(base) => row.without(base.row.pinned()),
                    None => row,
                },
                Err(_) => row,
            };
            catalog.put(provider, id, row);
        }
        catalog
    }

    /// Add `spec`, replacing any model of the same vendor and id whole. A
    /// later [`Self::with_overrides`] lays an override on it field by field,
    /// as on any other row.
    pub fn insert(&mut self, spec: ModelSpec) {
        let position = self.position(spec.provider.vendor(), &spec.id);
        let entry = Entry {
            row: Arc::new(Row::from_spec(&spec)),
            spec: Arc::new(spec),
        };
        let entries = Arc::make_mut(&mut self.entries);
        match position {
            Ok(index) => {
                if let Some(slot) = entries.get_mut(index) {
                    *slot = entry;
                }
            }
            Err(index) => entries.insert(index, entry),
        }
    }

    /// [`Self::get`] by vendor name.
    pub(crate) fn get_vendor(&self, vendor: &str, model: &str) -> Option<Resolved<'_>> {
        if let Some(shared) = self.entry(vendor, model) {
            return Some(Resolved {
                spec: shared,
                matched: Matched::Exact,
                shared,
            });
        }
        lookup::snapshot_bases(model).find_map(|base| {
            self.entry(vendor, base).map(|shared| Resolved {
                spec: shared,
                matched: Matched::SnapshotOf(shared.id.clone()),
                shared,
            })
        })
    }

    fn exact(&self, vendor: &str, model: &str) -> Option<&ModelSpec> {
        self.entry(vendor, model).map(|spec| &**spec)
    }

    fn entry(&self, vendor: &str, model: &str) -> Option<&Arc<ModelSpec>> {
        self.position(vendor, model)
            .ok()
            .and_then(|index| self.entries.get(index))
            .map(|entry| &entry.spec)
    }

    fn position(&self, vendor: &str, model: &str) -> Result<usize, usize> {
        self.entries.binary_search_by(|entry| {
            (entry.spec.provider.vendor(), entry.spec.id.as_str()).cmp(&(vendor, model))
        })
    }

    /// Add `row` for `provider`'s `id`, laid over any row already there.
    fn put(&mut self, provider: ProviderId, id: &str, row: Row) {
        let position = self.position(provider.vendor(), id);
        let entries = Arc::make_mut(&mut self.entries);
        match position {
            Ok(index) => {
                if let Some(entry) = entries.get_mut(index) {
                    let row = Row::clone(&entry.row).overlay(row);
                    entry.spec = Arc::new(row.spec(entry.spec.provider, id));
                    entry.row = Arc::new(row);
                }
            }
            Err(index) => {
                let spec = Arc::new(row.spec(provider, id));
                entries.insert(
                    index,
                    Entry {
                        row: Arc::new(row),
                        spec,
                    },
                );
            }
        }
    }
}

/// The seconds since the Unix epoch at `time`, written
/// `YYYY-MM-DDTHH:MM:SSZ`, or `None` when it is not in that form.
const fn unix_seconds(time: &str) -> Option<u64> {
    const fn number(mut digits: &[u8]) -> Option<u64> {
        let mut value = 0;
        while let [digit, rest @ ..] = digits {
            if !digit.is_ascii_digit() {
                return None;
            }
            value = value * 10 + (*digit - b'0') as u64;
            digits = rest;
        }
        Some(value)
    }
    let &[
        y0,
        y1,
        y2,
        y3,
        b'-',
        m0,
        m1,
        b'-',
        d0,
        d1,
        b'T',
        h0,
        h1,
        b':',
        n0,
        n1,
        b':',
        s0,
        s1,
        b'Z',
    ] = time.as_bytes()
    else {
        return None;
    };
    let (Some(year), Some(month), Some(day), Some(hour), Some(minute), Some(second)) = (
        number(&[y0, y1, y2, y3]),
        number(&[m0, m1]),
        number(&[d0, d1]),
        number(&[h0, h1]),
        number(&[n0, n1]),
        number(&[s0, s1]),
    ) else {
        return None;
    };
    if year < 1970 || month < 1 || month > 12 || day < 1 || day > 31 {
        return None;
    }
    if hour > 23 || minute > 59 || second > 59 {
        return None;
    }
    // Days from the civil date, by Howard Hinnant's algorithm, with the
    // year starting in March so the leap day falls last.
    let year = if month <= 2 { year - 1 } else { year };
    let era = year / 400;
    let of_era = year - era * 400;
    let shifted_month = if month > 2 { month - 3 } else { month + 9 };
    let of_year = (153 * shifted_month + 2) / 5 + day - 1;
    let of_era_days = of_era * 365 + of_era / 4 - of_era / 100 + of_year;
    let days = era * 146_097 + of_era_days - 719_468;
    Some(days * 86_400 + hour * 3600 + minute * 60 + second)
}

/// The rig vendor a models.dev provider key names.
fn vendor_of(key: &str) -> &str {
    KEYS.iter()
        .find_map(|(models_dev, vendor)| (*models_dev == key).then_some(*vendor))
        .unwrap_or(key)
}

/// A model reference as text spells it: `vendor/model`, or the older
/// `vendor/format:model`, which also names a protocol family. Every reader
/// of a reference string in rig parses it here.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Reference<'a> {
    /// The text before the first `/`.
    pub(crate) vendor: &'a str,
    /// The provider's model id.
    pub(crate) model: &'a str,
    /// The family the older spelling names, if it was used.
    pub(crate) format: Option<Format>,
}

impl<'a> Reference<'a> {
    /// `text` split at its first `/`, or `None` when either side is empty.
    /// The rest is read as `format:model` only when the text before its
    /// first `:` is a [`Format`]'s name.
    pub(crate) fn parse(text: &'a str) -> Option<Self> {
        let (vendor, rest) = text.split_once('/')?;
        if vendor.is_empty() || rest.is_empty() {
            return None;
        }
        if let Some((format, model)) = rest.split_once(':')
            && let Some(format) = Format::named(format)
        {
            return (!model.is_empty()).then_some(Self {
                vendor,
                model,
                format: Some(format),
            });
        }
        Some(Self {
            vendor,
            model: rest,
            format: None,
        })
    }
}

#[cfg(test)]
mod tests;