Skip to main content

oino_ai/
model.rs

1//! The `Model` and related catalogue types.
2
3use std::fmt;
4
5use indexmap::IndexMap;
6use serde::{Deserialize, Deserializer, Serialize, Serializer};
7
8use crate::json::{JsonObject, JsonValue, js_f64, js_f64_opt};
9use crate::primitives::{InputModality, ThinkingLevel};
10
11/// Chat model.
12#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
13#[serde(rename_all = "camelCase")]
14pub struct Model {
15    /// Model id within the provider.
16    pub id: String,
17    /// Display name.
18    pub name: String,
19    /// Wire API.
20    pub api: String,
21    /// Provider id.
22    pub provider: String,
23    /// Base URL.
24    #[serde(default)]
25    pub base_url: String,
26    /// Accepted input modalities.
27    #[serde(default)]
28    pub input: Vec<InputModality>,
29    /// Request size / image limits.
30    #[serde(default, skip_serializing_if = "Option::is_none")]
31    pub input_limits: Option<InputLimits>,
32    /// Pricing per million tokens.
33    #[serde(default)]
34    pub cost: ModelCost,
35    /// Extra request headers.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub headers: Option<IndexMap<String, String>>,
38    /// `"chat"` when present.
39    #[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
40    pub kind: Option<ModelKindChat>,
41    /// Whether the model supports reasoning.
42    #[serde(default)]
43    pub reasoning: bool,
44    /// Mapping of thinking levels to provider effort strings (`null` = unsupported).
45    #[serde(default, skip_serializing_if = "Option::is_none")]
46    pub thinking_level_map: Option<ThinkingLevelMap>,
47    /// Prompt cache pricing/retention hints.
48    #[serde(default, skip_serializing_if = "Option::is_none")]
49    pub prompt_cache: Option<PromptCache>,
50    /// Context window in tokens.
51    #[serde(default)]
52    pub context_window: u64,
53    /// Maximum output tokens.
54    #[serde(default)]
55    pub max_tokens: u64,
56    /// Default sampling parameters.
57    #[serde(default, skip_serializing_if = "Option::is_none")]
58    pub sampling_params: Option<JsonObject>,
59    /// Sampling parameters per thinking level.
60    #[serde(default, skip_serializing_if = "Option::is_none")]
61    pub sampling_params_by_thinking_level: Option<IndexMap<ThinkingLevel, JsonObject>>,
62    /// Provider compatibility flags (opaque to core).
63    #[serde(default, skip_serializing_if = "Option::is_none")]
64    pub compat: Option<JsonValue>,
65}
66
67/// The literal `"chat"` model type.
68#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
69pub enum ModelKindChat {
70    /// `"chat"`.
71    #[serde(rename = "chat")]
72    Chat,
73}
74
75/// `"provider/id"` model key.
76#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
77pub struct ModelKey {
78    /// Provider id.
79    pub provider: String,
80    /// Model id.
81    pub id: String,
82}
83
84impl fmt::Display for ModelKey {
85    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
86        write!(f, "{}/{}", self.provider, self.id)
87    }
88}
89
90impl Model {
91    /// `provider/id` key.
92    #[must_use]
93    pub fn key(&self) -> ModelKey {
94        ModelKey {
95            provider: self.provider.clone(),
96            id: self.id.clone(),
97        }
98    }
99
100    /// `[Off]` for non-reasoning models; otherwise every level
101    /// not mapped to `null`, with `xhigh`/`max` only when explicitly mapped.
102    #[must_use]
103    pub fn available_thinking_levels(&self) -> Vec<ThinkingLevel> {
104        if !self.reasoning {
105            return vec![ThinkingLevel::Off];
106        }
107        ThinkingLevel::ALL
108            .into_iter()
109            .filter(|level| {
110                let mapped = self
111                    .thinking_level_map
112                    .as_ref()
113                    .map_or(MapEntry::Absent, |map| map.get(*level));
114                match mapped {
115                    MapEntry::Null => false,
116                    MapEntry::Absent => !matches!(level, ThinkingLevel::Xhigh | ThinkingLevel::Max),
117                    MapEntry::Value(_) => true,
118                }
119            })
120            .collect()
121    }
122}
123
124/// Pricing (USD per million tokens).
125#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
126#[serde(rename_all = "camelCase")]
127pub struct ModelCost {
128    /// Input price.
129    #[serde(with = "js_f64")]
130    pub input: f64,
131    /// Output price.
132    #[serde(with = "js_f64")]
133    pub output: f64,
134    /// Cache read price.
135    #[serde(with = "js_f64")]
136    pub cache_read: f64,
137    /// Cache write price.
138    #[serde(with = "js_f64")]
139    pub cache_write: f64,
140    /// Tiered prices above input-token thresholds.
141    #[serde(default, skip_serializing_if = "Option::is_none")]
142    pub tiers: Option<Vec<ModelCostTier>>,
143}
144
145/// One price tier.
146#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
147#[serde(rename_all = "camelCase")]
148pub struct ModelCostTier {
149    /// Input price.
150    #[serde(with = "js_f64")]
151    pub input: f64,
152    /// Output price.
153    #[serde(with = "js_f64")]
154    pub output: f64,
155    /// Cache read price.
156    #[serde(with = "js_f64")]
157    pub cache_read: f64,
158    /// Cache write price.
159    #[serde(with = "js_f64")]
160    pub cache_write: f64,
161    /// Threshold in input tokens.
162    pub input_tokens_above: u64,
163}
164
165/// Value of one [`ThinkingLevelMap`] slot.
166#[derive(Clone, Copy, Debug, PartialEq, Eq)]
167pub enum MapEntry<'a> {
168    /// Key absent.
169    Absent,
170    /// Explicit `null`: level unsupported.
171    Null,
172    /// Provider effort string.
173    Value(&'a str),
174}
175
176/// Each key is absent, `null`, or a provider string.
177#[allow(
178    clippy::option_option,
179    reason = "absent, null and string are three distinct states"
180)]
181#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
182pub struct ThinkingLevelMap {
183    /// `off`.
184    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
185    pub off: Option<Option<String>>,
186    /// `minimal`.
187    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
188    pub minimal: Option<Option<String>>,
189    /// `low`.
190    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
191    pub low: Option<Option<String>>,
192    /// `medium`.
193    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
194    pub medium: Option<Option<String>>,
195    /// `high`.
196    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
197    pub high: Option<Option<String>>,
198    /// `xhigh`.
199    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
200    pub xhigh: Option<Option<String>>,
201    /// `max`.
202    #[serde(default, skip_serializing_if = "Option::is_none", with = "tri")]
203    pub max: Option<Option<String>>,
204}
205
206impl ThinkingLevelMap {
207    /// Slot for `level`.
208    #[must_use]
209    pub fn get(&self, level: ThinkingLevel) -> MapEntry<'_> {
210        let slot = match level {
211            ThinkingLevel::Off => &self.off,
212            ThinkingLevel::Minimal => &self.minimal,
213            ThinkingLevel::Low => &self.low,
214            ThinkingLevel::Medium => &self.medium,
215            ThinkingLevel::High => &self.high,
216            ThinkingLevel::Xhigh => &self.xhigh,
217            ThinkingLevel::Max => &self.max,
218        };
219        match slot {
220            None => MapEntry::Absent,
221            Some(None) => MapEntry::Null,
222            Some(Some(value)) => MapEntry::Value(value),
223        }
224    }
225}
226
227/// Serde for `Option<Option<String>>`: absent / `null` / string.
228#[allow(
229    clippy::option_option,
230    reason = "absent, null and string are three distinct states"
231)]
232mod tri {
233    use super::{Deserialize, Deserializer, Serializer};
234
235    #[allow(
236        clippy::ref_option,
237        reason = "signature required by serde(serialize_with)"
238    )]
239    pub(super) fn serialize<S: Serializer>(
240        value: &Option<Option<String>>,
241        serializer: S,
242    ) -> Result<S::Ok, S::Error> {
243        match value {
244            Some(Some(text)) => serializer.serialize_str(text),
245            _ => serializer.serialize_none(),
246        }
247    }
248
249    pub(super) fn deserialize<'de, D: Deserializer<'de>>(
250        deserializer: D,
251    ) -> Result<Option<Option<String>>, D::Error> {
252        Option::<String>::deserialize(deserializer).map(Some)
253    }
254}
255
256/// Prompt cache pricing hints.
257#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
258pub struct PromptCache {
259    /// Short-retention multiplier.
260    #[serde(default, skip_serializing_if = "Option::is_none", with = "js_f64_opt")]
261    pub short: Option<f64>,
262    /// Long-retention multiplier.
263    #[serde(default, skip_serializing_if = "Option::is_none", with = "js_f64_opt")]
264    pub long: Option<f64>,
265}
266
267/// Request limits.
268#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
269#[serde(rename_all = "camelCase")]
270pub struct InputLimits {
271    /// Maximum request body size.
272    #[serde(default, skip_serializing_if = "Option::is_none")]
273    pub max_request_bytes: Option<u64>,
274    /// Image limits.
275    #[serde(default, skip_serializing_if = "Option::is_none")]
276    pub images: Option<ImageInputLimits>,
277}
278
279/// Image limits.
280#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
281#[serde(rename_all = "camelCase")]
282pub struct ImageInputLimits {
283    /// Resize policy.
284    #[serde(default, skip_serializing_if = "Option::is_none")]
285    pub resize: Option<ImageResize>,
286    /// Images per message.
287    #[serde(default, skip_serializing_if = "Option::is_none")]
288    pub max_per_message: Option<u32>,
289    /// Images per request.
290    #[serde(default, skip_serializing_if = "Option::is_none")]
291    pub max_per_request: Option<u32>,
292}
293
294/// Image resize policy.
295#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
296#[serde(rename_all = "camelCase")]
297pub struct ImageResize {
298    /// Maximum width.
299    #[serde(default, skip_serializing_if = "Option::is_none")]
300    pub max_width: Option<u32>,
301    /// Maximum height.
302    #[serde(default, skip_serializing_if = "Option::is_none")]
303    pub max_height: Option<u32>,
304    /// Maximum encoded size.
305    #[serde(default, skip_serializing_if = "Option::is_none")]
306    pub max_bytes: Option<u64>,
307    /// JPEG quality used when re-encoding.
308    #[serde(default, skip_serializing_if = "Option::is_none")]
309    pub jpeg_quality: Option<u8>,
310}
311
312/// Image-generation model.
313#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
314#[serde(rename_all = "camelCase")]
315pub struct ImageModel {
316    /// Model id.
317    pub id: String,
318    /// Display name.
319    pub name: String,
320    /// Wire API.
321    pub api: String,
322    /// Provider id.
323    pub provider: String,
324    /// Base URL.
325    #[serde(default)]
326    pub base_url: String,
327    /// Input modalities.
328    #[serde(default)]
329    pub input: Vec<InputModality>,
330    /// Output modalities (`text`, `image`).
331    #[serde(default)]
332    pub output: Vec<InputModality>,
333    /// Pricing.
334    #[serde(default)]
335    pub cost: ModelCost,
336    /// Extra headers.
337    #[serde(default, skip_serializing_if = "Option::is_none")]
338    pub headers: Option<IndexMap<String, String>>,
339}
340
341/// Classifier model.
342#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
343#[serde(rename_all = "camelCase")]
344pub struct ClassifierModel {
345    /// Model id.
346    pub id: String,
347    /// Display name.
348    pub name: String,
349    /// Wire API.
350    pub api: String,
351    /// Provider id.
352    pub provider: String,
353    /// Base URL.
354    #[serde(default)]
355    pub base_url: String,
356    /// Input modalities.
357    #[serde(default)]
358    pub input: Vec<InputModality>,
359    /// Pricing.
360    #[serde(default)]
361    pub cost: ModelCost,
362    /// Extra headers.
363    #[serde(default, skip_serializing_if = "Option::is_none")]
364    pub headers: Option<IndexMap<String, String>>,
365    /// Context window.
366    #[serde(default)]
367    pub context_window: u64,
368    /// Provider compatibility flags.
369    #[serde(default, skip_serializing_if = "Option::is_none")]
370    pub compat: Option<JsonValue>,
371}
372
373/// Any registry model, discriminated by `type` (absent ⇒ chat).
374#[derive(Clone, Debug, PartialEq)]
375pub enum AnyModel {
376    /// Chat model (`type` absent or `"chat"`).
377    Chat(Box<Model>),
378    /// `type: "image"`.
379    Image(Box<ImageModel>),
380    /// `type: "classifier"`.
381    Classifier(Box<ClassifierModel>),
382}
383
384impl Serialize for AnyModel {
385    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
386        use serde::ser::Error;
387        let (mut value, tag) = match self {
388            Self::Chat(model) => return model.serialize(serializer),
389            Self::Image(model) => (
390                serde_json::to_value(model).map_err(S::Error::custom)?,
391                "image",
392            ),
393            Self::Classifier(model) => (
394                serde_json::to_value(model).map_err(S::Error::custom)?,
395                "classifier",
396            ),
397        };
398        if let Some(object) = value.as_object_mut() {
399            object.insert("type".into(), JsonValue::String(tag.into()));
400        }
401        value.serialize(serializer)
402    }
403}
404
405impl<'de> Deserialize<'de> for AnyModel {
406    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
407        use serde::de::Error;
408        let value = JsonValue::deserialize(deserializer)?;
409        let kind = value
410            .get("type")
411            .and_then(JsonValue::as_str)
412            .map(str::to_owned);
413        match kind.as_deref() {
414            None | Some("chat") => serde_json::from_value(value)
415                .map(|m| Self::Chat(Box::new(m)))
416                .map_err(D::Error::custom),
417            Some("image") => serde_json::from_value(value)
418                .map(|m| Self::Image(Box::new(m)))
419                .map_err(D::Error::custom),
420            Some("classifier") => serde_json::from_value(value)
421                .map(|m| Self::Classifier(Box::new(m)))
422                .map_err(D::Error::custom),
423            Some(other) => Err(D::Error::custom(format!("unknown model type: {other}"))),
424        }
425    }
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    use serde_json::json;
432
433    fn model(raw: JsonValue) -> Option<Model> {
434        serde_json::from_value(raw).ok()
435    }
436
437    #[test]
438    fn thinking_levels_follow_map() {
439        let base = json!({"id":"m","name":"M","api":"a","provider":"p","baseUrl":"","input":["text"],
440            "cost":{"input":3,"output":15,"cacheRead":0.3,"cacheWrite":3.75},"reasoning":true,
441            "contextWindow":200_000,"maxTokens":8192});
442        let plain = model(base.clone());
443        assert_eq!(plain.map(|m| m.available_thinking_levels().len()), Some(5));
444        let mut mapped = base;
445        mapped["thinkingLevelMap"] = json!({"minimal": null, "xhigh": "max"});
446        let levels = model(mapped.clone()).map(|m| m.available_thinking_levels());
447        assert_eq!(
448            levels,
449            Some(vec![
450                ThinkingLevel::Off,
451                ThinkingLevel::Low,
452                ThinkingLevel::Medium,
453                ThinkingLevel::High,
454                ThinkingLevel::Xhigh
455            ])
456        );
457        assert_eq!(
458            model(mapped.clone()).and_then(|m| serde_json::to_value(m).ok()),
459            Some(mapped)
460        );
461    }
462
463    #[test]
464    fn any_model_dispatches_on_type() {
465        let image: Option<AnyModel> = serde_json::from_value(json!({"type":"image","id":"i","name":"I","api":"x",
466            "provider":"p","baseUrl":"","input":["text"],"output":["image"],"cost":{"input":0,"output":0,"cacheRead":0,"cacheWrite":0}})).ok();
467        assert!(matches!(image, Some(AnyModel::Image(_))));
468        let raw = image
469            .and_then(|m| serde_json::to_value(m).ok())
470            .unwrap_or_default();
471        assert_eq!(raw["type"], json!("image"));
472    }
473}