rig_core/catalog/facts.rs
1//! The facts a completion wire encodes with, and that its replies are priced
2//! by.
3
4use std::sync::Arc;
5
6use super::{Catalog, ModelSpec};
7
8/// The facts a completion wire encodes with: the spec of the model it was
9/// connected to, and the catalog that spec came from, which answers for any
10/// other model a request names. Every completion wire holds one, and its
11/// [`Descriptor`](crate::wire::Descriptor) hands it to the completion fold,
12/// which prices usage from it.
13///
14/// The default binds no spec and answers from [`Catalog::builtin`]. Facts
15/// are not part of a wire's serialized form or of its equality: a wire read
16/// back from storage answers from the built-in catalog until it is given
17/// facts again. Clones are cheap.
18///
19/// ```
20/// use rig_core::catalog::{Catalog, ModelFacts, Pricing};
21/// use rig_core::providers::registry::ProviderId;
22///
23/// let anthropic = ProviderId::catalog("anthropic").ok_or("a known vendor")?;
24/// let mut catalog = Catalog::builtin().clone();
25/// let spec = catalog.get_exact(anthropic, "claude-haiku-4-5").ok_or("listed")?;
26/// catalog.insert(spec.clone().with_pricing(Pricing::new(0.5, 2.5)));
27///
28/// let facts = catalog.facts(anthropic, "claude-haiku-4-5-20990101");
29/// let spec = facts.for_model("anthropic", "claude-haiku-4-5-20990101").ok_or("bound")?;
30/// assert_eq!(spec.pricing.as_ref().map(|pricing| pricing.input), Some(0.5));
31/// # Ok::<(), &str>(())
32/// ```
33#[derive(Clone, Debug, Default)]
34pub struct ModelFacts {
35 bound: Option<Arc<Bound>>,
36 catalog: Option<Catalog>,
37}
38
39/// The model a wire was connected to: the id it addresses and its spec.
40#[derive(Debug)]
41struct Bound {
42 id: String,
43 spec: Arc<ModelSpec>,
44}
45
46impl ModelFacts {
47 /// Facts that bind no spec and answer from the built-in catalog, as the
48 /// default does.
49 pub(crate) fn builtin() -> &'static ModelFacts {
50 static BUILTIN: ModelFacts = ModelFacts {
51 bound: None,
52 catalog: None,
53 };
54 &BUILTIN
55 }
56
57 /// The facts of `spec`, for the model it lists. Any other model is
58 /// answered from [`Catalog::builtin`].
59 pub fn new(spec: impl Into<Arc<ModelSpec>>) -> Self {
60 let spec = spec.into();
61 Self {
62 bound: Some(Arc::new(Bound {
63 id: spec.id.clone(),
64 spec,
65 })),
66 catalog: None,
67 }
68 }
69
70 /// The same facts, answering every model but the bound one from
71 /// `catalog`.
72 pub fn with_catalog(mut self, catalog: &Catalog) -> Self {
73 self.catalog = Some(catalog.clone());
74 self
75 }
76
77 /// The spec of the model the wire was connected to, when it was
78 /// connected through the catalog or given one.
79 pub fn spec(&self) -> Option<&ModelSpec> {
80 self.bound.as_deref().map(|bound| &*bound.spec)
81 }
82
83 /// The catalog that answers for every model but the bound one.
84 pub fn catalog(&self) -> &Catalog {
85 self.catalog.as_ref().unwrap_or_else(|| Catalog::builtin())
86 }
87
88 /// The facts of `vendor`'s `model`: the bound spec when `model` is the id
89 /// the wire was connected to (or the id the spec lists), otherwise the
90 /// catalog's entry by [`Catalog::get`]'s rule. `None` for a model the
91 /// catalog does not list. `vendor` is the provider's descriptor name
92 /// (`"anthropic"`, `"aws_bedrock"`).
93 ///
94 /// A spec under a catalog vendor answers only under that vendor, since
95 /// a wire also asks about other vendors' rows for its model. A spec
96 /// under a [declared](crate::providers::registry::ProviderId::declare)
97 /// vendor answers under any: no wire carries a declared vendor's name,
98 /// so the wire it was bound to asks under its own.
99 pub fn for_model(&self, vendor: &str, model: &str) -> Option<&ModelSpec> {
100 if let Some(bound) = self.bound.as_deref()
101 && (bound.spec.provider.is_declared() || bound.spec.provider.vendor() == vendor)
102 && (bound.id == model || bound.spec.id == model)
103 {
104 return Some(&bound.spec);
105 }
106 self.catalog()
107 .get_vendor(vendor, model)
108 .map(|resolved| resolved.spec)
109 }
110
111 /// Whether `vendor`'s `model` reads images: what its facts list, or, for
112 /// a model the catalog does not list, what `rule` (the vendor's naming
113 /// rule) says of its id. Every wire that filters images reads this, so
114 /// an id the catalog does not list (a gateway's spelling, another case,
115 /// a deployment name) keeps its vendor's naming rule.
116 pub(crate) fn reads_images_or(
117 &self,
118 vendor: &str,
119 model: &str,
120 rule: impl FnOnce(&str) -> bool,
121 ) -> bool {
122 self.for_model(vendor, model)
123 .map_or_else(|| rule(model), |spec| spec.input.image)
124 }
125}
126
127/// Facts take no part in a wire's equality, as they take none in its
128/// serialized form: two wires that send the same requests are equal.
129impl PartialEq for ModelFacts {
130 fn eq(&self, _other: &Self) -> bool {
131 true
132 }
133}
134
135impl Eq for ModelFacts {}
136
137impl Catalog {
138 /// The facts of `provider`'s `model` from this catalog: its entry by
139 /// [`Self::get`]'s rule, bound to `model` as written, and this catalog
140 /// for every other model a request names. A model the catalog does not
141 /// list binds nothing.
142 pub fn facts(
143 &self,
144 provider: crate::providers::registry::ProviderId,
145 model: &str,
146 ) -> ModelFacts {
147 let bound = self.get_vendor(provider.vendor(), model).map(|found| {
148 Arc::new(Bound {
149 id: model.to_owned(),
150 spec: found.shared(),
151 })
152 });
153 ModelFacts {
154 bound,
155 catalog: Some(self.clone()),
156 }
157 }
158}
159
160#[cfg(test)]
161mod tests;