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,
};
const BUILTIN: &str = include_str!("models.json");
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"),
};
#[derive(Clone, Debug, Default)]
pub struct Catalog {
entries: Arc<Vec<Entry>>,
}
#[derive(Clone, Debug)]
struct Entry {
row: Arc<Row>,
spec: Arc<ModelSpec>,
}
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum CatalogError {
#[error("not a models.dev-style catalog: {0}")]
Json(#[from] serde_json::Error),
}
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 {
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,
#[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}"),
})
}
pub fn generated_at() -> SystemTime {
UNIX_EPOCH + Duration::from_secs(GENERATED_AT)
}
pub fn get(&self, provider: ProviderId, model: &str) -> Option<Resolved<'_>> {
self.get_vendor(provider.vendor(), model)
}
pub fn get_exact(&self, provider: ProviderId, model: &str) -> Option<&ModelSpec> {
self.exact(provider.vendor(), model)
}
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))
}
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,
}
}
pub fn iter(&self) -> impl Iterator<Item = &ModelSpec> {
self.entries.iter().map(|entry| &*entry.spec)
}
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
}
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
}
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),
}
}
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))
})
}
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,
},
);
}
}
}
}
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;
}
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)
}
fn vendor_of(key: &str) -> &str {
KEYS.iter()
.find_map(|(models_dev, vendor)| (*models_dev == key).then_some(*vendor))
.unwrap_or(key)
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct Reference<'a> {
pub(crate) vendor: &'a str,
pub(crate) model: &'a str,
pub(crate) format: Option<Format>,
}
impl<'a> Reference<'a> {
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;