rig_core/catalog/
lookup.rs

1//! The one rule every catalog lookup follows, what a lookup returns, and the
2//! suggestions a miss carries.
3
4use std::fmt;
5use std::sync::Arc;
6
7use super::ModelSpec;
8
9/// How a requested model id matched its catalog entry.
10#[non_exhaustive]
11#[derive(Clone, Debug, PartialEq, Eq)]
12pub enum Matched {
13    /// The id is listed as given.
14    Exact,
15    /// The id is a dated snapshot of the listed id this holds: that id
16    /// followed by `-20` and a year (`claude-sonnet-4-5-20250929`,
17    /// `gpt-5-2025-08-07`, `claude-opus-5-5-20260601-v1:0`).
18    SnapshotOf(String),
19}
20
21/// A model the catalog found, and how the requested id matched it.
22#[non_exhaustive]
23#[derive(Clone, Debug, PartialEq)]
24pub struct Resolved<'a> {
25    /// The catalog entry.
26    pub spec: &'a ModelSpec,
27    /// How the requested id matched [`Self::spec`].
28    pub matched: Matched,
29    /// The entry as the catalog keeps it.
30    pub(super) shared: &'a Arc<ModelSpec>,
31}
32
33impl Resolved<'_> {
34    /// The catalog entry as the catalog keeps it, to hold beyond the
35    /// lookup without copying it.
36    pub fn shared(&self) -> Arc<ModelSpec> {
37        Arc::clone(self.shared)
38    }
39}
40
41/// A reference the catalog does not list, with up to five listed models
42/// whose references are close to it, or else the closest models of its
43/// vendor.
44#[non_exhaustive]
45#[derive(Clone, Debug, PartialEq, Eq)]
46pub struct NotFound {
47    /// The reference as given.
48    pub reference: String,
49    /// Close references to models that are not deprecated, closest first;
50    /// with none close, the vendor's closest ones.
51    pub suggestions: Vec<String>,
52}
53
54impl fmt::Display for NotFound {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        write!(f, "the catalog lists no model `{}`", self.reference)?;
57        match self.suggestions.as_slice() {
58            [] => Ok(()),
59            [only] => write!(f, "; did you mean `{only}`?"),
60            many => write!(f, "; did you mean one of `{}`?", many.join("`, `")),
61        }
62    }
63}
64
65impl std::error::Error for NotFound {}
66
67/// The most suggestions a [`NotFound`] carries.
68const SUGGESTIONS: usize = 5;
69
70/// The listed ids `model` may be a dated snapshot of, longest first: each
71/// prefix of `model` that is followed by `-20` and two more digits.
72pub(super) fn snapshot_bases(model: &str) -> impl Iterator<Item = &str> {
73    let mut bases: Vec<&str> = model
74        .match_indices("-20")
75        .filter(|(at, _)| {
76            model
77                .get(at + 3..at + 5)
78                .is_some_and(|year| year.bytes().all(|byte| byte.is_ascii_digit()))
79        })
80        .filter_map(|(at, _)| model.get(..at))
81        .filter(|base| !base.is_empty())
82        .collect();
83    bases.reverse();
84    bases.into_iter()
85}
86
87/// Up to [`SUGGESTIONS`] references from `candidates` close to `reference`,
88/// closest first. A candidate is close when it contains the reference's
89/// model id or its edit distance is at most a third of the reference's
90/// length.
91pub(super) fn suggestions<'a>(
92    reference: &str,
93    candidates: impl Iterator<Item = &'a ModelSpec>,
94) -> Vec<String> {
95    let wanted = reference.to_lowercase();
96    let model = wanted
97        .split_once('/')
98        .map_or(wanted.as_str(), |(_, model)| model);
99    let limit = wanted.chars().count() / 3 + 1;
100    closest(&wanted, candidates, |spec, distance| {
101        let contains = !model.is_empty() && spec.id.to_lowercase().contains(model);
102        contains || distance <= limit
103    })
104}
105
106/// Up to [`SUGGESTIONS`] models of `reference`'s vendor from `candidates`,
107/// closest first however far, for a miss no listed model is close to.
108pub(super) fn vendor_models<'a>(
109    reference: &str,
110    candidates: impl Iterator<Item = &'a ModelSpec>,
111) -> Vec<String> {
112    let wanted = reference.to_lowercase();
113    let vendor = wanted.split_once('/').map_or("", |(vendor, _)| vendor);
114    closest(&wanted, candidates, |spec, _| {
115        spec.provider.vendor().eq_ignore_ascii_case(vendor)
116    })
117}
118
119/// Up to [`SUGGESTIONS`] references from `candidates` that are not
120/// deprecated and that `keep` takes, given their edit distance from
121/// `wanted` (lowercase), closest first.
122fn closest<'a>(
123    wanted: &str,
124    candidates: impl Iterator<Item = &'a ModelSpec>,
125    keep: impl Fn(&ModelSpec, usize) -> bool,
126) -> Vec<String> {
127    let mut close: Vec<(usize, String)> = candidates
128        .filter(|spec| !spec.deprecated)
129        .filter_map(|spec| {
130            let candidate = format!("{}/{}", spec.provider.vendor(), spec.id);
131            let distance = edit_distance(wanted, &candidate.to_lowercase());
132            keep(spec, distance).then_some((distance, candidate))
133        })
134        .collect();
135    close.sort();
136    close.truncate(SUGGESTIONS);
137    close.into_iter().map(|(_, candidate)| candidate).collect()
138}
139
140/// The edit distance between `a` and `b`, by character: insertions,
141/// deletions, substitutions and swaps of two adjacent characters each count
142/// one (optimal string alignment), so `inptu` is one edit from `input`.
143pub(crate) fn edit_distance(a: &str, b: &str) -> usize {
144    let a: Vec<char> = a.chars().collect();
145    let b: Vec<char> = b.chars().collect();
146    let at = |row: &[usize], k: usize| row.get(k).copied().unwrap_or(usize::MAX);
147    let mut before: Vec<usize> = Vec::new();
148    let mut previous: Vec<usize> = (0..=b.len()).collect();
149    for (i, left) in a.iter().enumerate() {
150        let mut current = Vec::with_capacity(b.len() + 1);
151        current.push(i + 1);
152        for (j, right) in b.iter().enumerate() {
153            let substitute = at(&previous, j).saturating_add(usize::from(left != right));
154            let delete = at(&previous, j + 1).saturating_add(1);
155            let insert = at(&current, j).saturating_add(1);
156            let mut best = substitute.min(delete).min(insert);
157            let swapped = i > 0
158                && j > 0
159                && a.get(i - 1) == Some(right)
160                && b.get(j - 1) == Some(left)
161                && left != right;
162            if swapped {
163                best = best.min(at(&before, j - 1).saturating_add(1));
164            }
165            current.push(best);
166        }
167        before = std::mem::replace(&mut previous, current);
168    }
169    previous.last().copied().unwrap_or(0)
170}
rig_core/catalog/lookup.rs

rig_core/catalog/
lookup.rs
