rig-core 0.44.0

//! A catalog of model facts, keyed by provider and model id: context and
//! output limits, input modalities, the reasoning and caching each model
//! takes, and its prices.
//!
//! The built-in catalog is data: `cargo xtask catalog sync` generates it from
//! models.dev, and the main providers' entries are reviewed by hand. A
//! harness can read an override file in the same shape with
//! [`Catalog::from_json`] and lay it over the built-in one with
//! [`Catalog::merge`]. [`ModelSpec::validate`] checks a request's
//! [`GenerationOptions`](crate::completion::GenerationOptions) against what
//! the model takes 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")
//!     .ok_or("listed")?;
//! let options = GenerationOptions::default().reasoning(Effort::High);
//! assert!(haiku.validate(&options).is_err(), "Haiku 4.5 takes a budget");
//! # Ok::<(), &str>(())
//! ```

mod row;
mod spec;

use std::collections::BTreeMap;
use std::sync::OnceLock;

use serde::Deserialize;

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

#[doc(hidden)]
pub use spec::Compat;
pub use spec::{CacheSupport, Modalities, ModelSpec, Pricing, ReasoningSupport, Sampling};

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

/// Model facts by provider and model id.
#[derive(Clone, Debug, Default)]
pub struct Catalog {
    /// Sorted by vendor, then model id.
    entries: Vec<Entry>,
}

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

/// Why catalog data did not load.
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum CatalogError {
    /// The text is not a models.dev-style catalog: an object of provider
    /// keys, each with a `models` object keyed by model id.
    #[error("not a models.dev-style catalog: {0}")]
    Json(#[from] serde_json::Error),
}

/// A provider's section of the data. models.dev's other provider keys
/// (`name`, `env`, `doc`, ...) are ignored.
#[derive(Deserialize)]
struct Section {
    #[serde(default)]
    models: BTreeMap<String, Row>,
}

/// 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.
    pub fn builtin() -> &'static Catalog {
        static BUILTIN_CATALOG: OnceLock<Catalog> = OnceLock::new();
        // The data is checked in, and a test proves it parses.
        BUILTIN_CATALOG.get_or_init(|| Catalog::from_json(BUILTIN).unwrap_or_default())
    }

    /// The model `provider` serves as `model`, by exact id. Any selection of
    /// a vendor finds that vendor's models.
    pub fn get(&self, provider: ProviderId, model: &str) -> Option<&ModelSpec> {
        self.exact(provider.vendor(), model)
    }

    /// The model a reference names: `vendor/model`
    /// (`anthropic/claude-opus-5-5`, `openrouter/anthropic/claude-sonnet-4.5`)
    /// or `vendor[/format]:model` as [`ProviderRef`] spells it. The second
    /// form applies when the text before the first `:` is a vendor or a
    /// `vendor/format` pair, so a model id holding a `:` (`ollama/qwen3:4b`)
    /// still reads as the first.
    ///
    /// [`ProviderRef`]: crate::providers::registry::ProviderRef
    pub fn resolve(&self, reference: &str) -> Option<&ModelSpec> {
        let (vendor, model) = split_reference(reference)?;
        self.exact(vendor, model)
    }

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

    /// Read a models.dev-style catalog: an object of provider keys, each
    /// with a `models` object keyed by model id. A provider key is a
    /// models.dev key (`google`, `amazon-bedrock`) or a rig vendor name
    /// (`gcp.gemini`); a key that is neither is skipped, so models.dev's own
    /// `api.json` reads whole. Rig's hand-entered facts go under each row's
    /// `rig` object.
    pub fn from_json(json: &str) -> Result<Catalog, CatalogError> {
        let sections: BTreeMap<String, Section> = serde_json::from_str(json)?;
        let mut catalog = Catalog::default();
        for (key, section) in sections {
            let Some(provider) = ProviderId::catalog(vendor_of(&key)) else {
                continue;
            };
            for (id, row) in section.models {
                catalog.put(provider, &id, row);
            }
        }
        Ok(catalog)
    }

    /// This catalog with `overrides` laid over it. A model in both keeps
    /// every field the override leaves out; a model only in `overrides` is
    /// added.
    pub fn merge(mut self, overrides: Catalog) -> Catalog {
        for entry in overrides.entries {
            self.put(entry.spec.provider, &entry.spec.id, entry.row);
        }
        self
    }

    /// The model `provider` serves as `model`, as the encoders look it up:
    /// by exact id or, failing that, without a dated snapshot suffix
    /// (`-20251001`, `-2025-08-07`), so `claude-sonnet-4-5-20250929` finds
    /// `claude-sonnet-4-5`. [`Self::get`] matches the exact id only.
    pub fn find(&self, provider: ProviderId, model: &str) -> Option<&ModelSpec> {
        self.find_vendor(provider.vendor(), model)
    }

    /// [`Self::find`] by vendor name.
    pub(crate) fn find_vendor(&self, vendor: &str, model: &str) -> Option<&ModelSpec> {
        self.exact(vendor, model)
            .or_else(|| self.exact(vendor, undated(model)?))
    }

    fn exact(&self, vendor: &str, model: &str) -> Option<&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) {
        match self.position(provider.vendor(), id) {
            Ok(index) => {
                if let Some(entry) = self.entries.get_mut(index) {
                    let row = std::mem::take(&mut entry.row).overlay(row);
                    entry.spec = row.spec(entry.spec.provider, id);
                    entry.row = row;
                }
            }
            Err(index) => {
                let spec = row.spec(provider, id);
                self.entries.insert(index, Entry { row, spec });
            }
        }
    }
}

/// The model of the built-in catalog `vendor` serves as `model`, looked up
/// as the encoders look models up ([`Catalog::find`]).
pub(crate) fn lookup(vendor: &str, model: &str) -> Option<&'static ModelSpec> {
    Catalog::builtin().find_vendor(vendor, model)
}

/// The model of the built-in catalog `vendor` serves as `model`, or as a
/// snapshot of it: a listed id followed by `-20` and anything
/// (`claude-opus-4-1-20250805`, `claude-opus-5-5-20260601-v1:0`), the
/// longest such id.
pub(crate) fn lookup_snapshot(vendor: &str, model: &str) -> Option<&'static ModelSpec> {
    let catalog = Catalog::builtin();
    catalog.exact(vendor, model).or_else(|| {
        model
            .match_indices("-20")
            .filter_map(|(at, _)| catalog.exact(vendor, model.get(..at)?))
            .last()
    })
}

/// Whether `vendor`'s `model` reads images: what its catalog entry lists,
/// or, for a model the catalog does not list, what `rule` (the vendor's
/// naming rule) says of its id. Every wire that filters images reads this,
/// so an id the catalog does not list (a gateway's spelling, another case,
/// a deployment name) keeps its vendor's naming rule.
pub(crate) fn reads_images_or(vendor: &str, model: &str, rule: impl FnOnce(&str) -> bool) -> bool {
    lookup(vendor, model).map_or_else(|| rule(model), |spec| spec.input.image)
}

/// 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)
}

/// `(vendor, model)` from a reference, by the grammar [`Catalog::resolve`]
/// documents.
pub(crate) fn split_reference(reference: &str) -> Option<(&str, &str)> {
    if let Some((selection, model)) = reference.split_once(':') {
        let vendor = match selection.split_once('/') {
            None => Some(selection),
            Some((vendor, format)) => Format::named(format).map(|_| vendor),
        };
        if let Some(vendor) = vendor.filter(|vendor| !vendor.is_empty()) {
            return (!model.is_empty()).then_some((vendor, model));
        }
    }
    reference
        .split_once('/')
        .filter(|(vendor, model)| !vendor.is_empty() && !model.is_empty())
}

/// `model` without a trailing dated snapshot: `-YYYYMMDD` or `-YYYY-MM-DD`.
fn undated(model: &str) -> Option<&str> {
    let digits = |text: &str| !text.is_empty() && text.bytes().all(|byte| byte.is_ascii_digit());
    let (rest, last) = model.rsplit_once('-')?;
    if last.len() == 8 && digits(last) {
        return Some(rest);
    }
    let (rest, month) = rest.rsplit_once('-')?;
    let (rest, year) = rest.rsplit_once('-')?;
    (year.len() == 4 && month.len() == 2 && last.len() == 2)
        .then_some(())
        .filter(|()| digits(year) && digits(month) && digits(last))
        .map(|()| rest)
}

#[cfg(test)]
mod tests;