Skip to main content

yapi_types/
model.rs

1//! Catalog models: what a provider serves and how to talk to it.
2//!
3//! Mirrors `Model` and the compat interfaces in `packages/ai/src/types.ts` in pi
4//! `v1.0.0`. The same shape is used by the built-in catalog, by models resolved from
5//! `models.json` and by extension-registered providers.
6#![allow(
7    missing_docs,
8    reason = "fields mirror pi's TypeScript types; contracts are noted where they differ from the name"
9)]
10
11use indexmap::IndexMap;
12use serde::de::DeserializeOwned;
13use serde::{Deserialize, Serialize};
14use serde_json::{Map, Value};
15
16use crate::message::ThinkingLevel;
17use crate::models::{CostTier, ImageResize, InputKind, InputLimits, PromptCache};
18
19/// A chat model.
20#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
21#[serde(rename_all = "camelCase")]
22pub struct Model {
23    pub id: String,
24    pub name: String,
25    /// Wire API id, such as `anthropic-messages`.
26    pub api: String,
27    pub provider: String,
28    pub base_url: String,
29    pub reasoning: bool,
30    pub input: Vec<InputKind>,
31    pub cost: Pricing,
32    pub context_window: u64,
33    pub max_tokens: u64,
34    /// pi thinking level to provider value; `None` marks the level unsupported.
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub thinking_level_map: Option<IndexMap<String, Option<String>>>,
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub prompt_cache: Option<PromptCache>,
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub sampling_params: Option<Map<String, Value>>,
41    /// Wire-API-specific options; read through [`Model::compat`].
42    #[serde(default, skip_serializing_if = "Option::is_none")]
43    pub compat: Option<Map<String, Value>>,
44    #[serde(default, skip_serializing_if = "Option::is_none")]
45    pub input_limits: Option<InputLimits>,
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub headers: Option<IndexMap<String, String>>,
48    /// `"chat"` or absent.
49    #[serde(default, rename = "type", skip_serializing_if = "Option::is_none")]
50    pub kind: Option<String>,
51}
52
53/// An image-generation model, usable with `generateImages()` only. Fields
54/// follow pi's catalog order.
55#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
56#[serde(rename_all = "camelCase")]
57pub struct ImageModel {
58    /// Always `"image"`.
59    #[serde(rename = "type")]
60    pub kind: String,
61    pub id: String,
62    pub name: String,
63    pub api: String,
64    pub provider: String,
65    pub base_url: String,
66    pub input: Vec<InputKind>,
67    /// Output modalities; always includes `image`.
68    pub output: Vec<InputKind>,
69    pub cost: Pricing,
70    #[serde(default, skip_serializing_if = "Option::is_none")]
71    pub input_limits: Option<InputLimits>,
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    pub headers: Option<IndexMap<String, String>>,
74}
75
76/// A structured classifier model, usable with `classify()` only. Fields
77/// follow pi's catalog order.
78#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
79#[serde(rename_all = "camelCase")]
80pub struct ClassifierModel {
81    /// Always `"classifier"`.
82    #[serde(rename = "type")]
83    pub kind: String,
84    pub id: String,
85    pub name: String,
86    pub api: String,
87    pub provider: String,
88    pub base_url: String,
89    pub input: Vec<InputKind>,
90    pub cost: Pricing,
91    pub context_window: u64,
92    #[serde(default, skip_serializing_if = "Option::is_none")]
93    pub input_limits: Option<InputLimits>,
94    #[serde(default, skip_serializing_if = "Option::is_none")]
95    pub headers: Option<IndexMap<String, String>>,
96}
97
98impl ImageModel {
99    /// `provider/id`.
100    pub fn reference(&self) -> String {
101        format!("{}/{}", self.provider, self.id)
102    }
103}
104
105impl ClassifierModel {
106    /// `provider/id`.
107    pub fn reference(&self) -> String {
108        format!("{}/{}", self.provider, self.id)
109    }
110}
111
112impl Model {
113    /// Whether this is model `id` of `provider`.
114    pub fn is(&self, provider: &str, id: &str) -> bool {
115        self.provider == provider && self.id == id
116    }
117
118    /// `provider/id`, the reference pi prints and accepts.
119    pub fn reference(&self) -> String {
120        format!("{}/{}", self.provider, self.id)
121    }
122
123    /// Reads `compat` as the view for this model's wire API. Unknown or invalid keys
124    /// are ignored, as in pi.
125    pub fn compat<T: DeserializeOwned + Default>(&self) -> T {
126        self.compat
127            .as_ref()
128            .and_then(|compat| T::deserialize(Value::Object(compat.clone())).ok())
129            .unwrap_or_default()
130    }
131
132    /// Whether the model accepts image input.
133    pub fn accepts_images(&self) -> bool {
134        self.input.contains(&InputKind::Image)
135    }
136
137    /// The resize profile for images that enter the conversation, from
138    /// `inputLimits.images.resize`.
139    pub fn image_resize(&self) -> Option<&ImageResize> {
140        self.input_limits.as_ref()?.images.as_ref()?.resize.as_ref()
141    }
142
143    /// The provider value for a thinking level: `Some(None)` when the map marks it
144    /// unsupported, `None` when the map does not mention it.
145    pub fn thinking_level_value(&self, level: ThinkingLevel) -> Option<Option<&str>> {
146        self.thinking_level_map
147            .as_ref()?
148            .get(level.as_str())
149            .map(Option::as_deref)
150    }
151}
152
153/// US dollars per million tokens. pi calls this `ModelCost`.
154#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
155#[serde(rename_all = "camelCase")]
156pub struct Pricing {
157    pub input: f64,
158    pub output: f64,
159    pub cache_read: f64,
160    pub cache_write: f64,
161    /// Request-wide rates; the highest tier below the input size applies.
162    #[serde(default, skip_serializing_if = "Option::is_none")]
163    pub tiers: Option<Vec<CostTier>>,
164}
165
166/// `compat` for `anthropic-messages`.
167#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
168#[serde(rename_all = "camelCase")]
169pub struct AnthropicMessagesCompat {
170    pub supports_eager_tool_input_streaming: Option<bool>,
171    pub supports_long_cache_retention: Option<bool>,
172    pub send_session_affinity_headers: Option<bool>,
173    pub session_affinity_format: Option<String>,
174    pub supports_cache_control_on_tools: Option<bool>,
175    pub supports_temperature: Option<bool>,
176    pub force_adaptive_thinking: Option<bool>,
177    pub allow_empty_signature: Option<bool>,
178    pub supports_strict_tools: Option<bool>,
179    pub supports_mid_convo_effort: Option<bool>,
180    pub supports_mid_convo_system_messages: Option<bool>,
181    pub supports_mid_convo_tool_changes: Option<bool>,
182}
183
184/// `compat` for `openai-completions`.
185#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
186#[serde(rename_all = "camelCase")]
187pub struct OpenAiCompletionsCompat {
188    pub supports_store: Option<bool>,
189    pub supports_developer_role: Option<bool>,
190    pub supports_reasoning_effort: Option<bool>,
191    pub supports_usage_in_streaming: Option<bool>,
192    pub supports_finish_reason: Option<bool>,
193    /// `max_completion_tokens` or `max_tokens`.
194    pub max_tokens_field: Option<String>,
195    pub requires_tool_result_name: Option<bool>,
196    pub requires_assistant_after_tool_result: Option<bool>,
197    pub requires_thinking_as_text: Option<bool>,
198    pub requires_reasoning_content_on_assistant_messages: Option<bool>,
199    pub thinking_format: Option<String>,
200    pub chat_template_kwargs: Option<Map<String, Value>>,
201    pub chat_template_args: Option<Map<String, Value>>,
202    pub open_router_routing: Option<Map<String, Value>>,
203    pub vercel_gateway_routing: Option<Map<String, Value>>,
204    pub zai_tool_stream: Option<bool>,
205    pub thinking_token_budget_field: Option<String>,
206    pub supports_thinking_token_budget: Option<bool>,
207    #[serde(rename = "supportsOpenAIGrammarTools")]
208    pub supports_open_ai_grammar_tools: Option<bool>,
209    pub supports_mid_convo_system_messages: Option<bool>,
210    pub supports_mid_convo_tool_additions: Option<bool>,
211    pub supports_strict_mode: Option<bool>,
212    pub cache_control_format: Option<String>,
213    pub send_session_affinity_headers: Option<bool>,
214    pub session_affinity_format: Option<String>,
215    pub supports_long_cache_retention: Option<bool>,
216    pub vllm_priority: Option<f64>,
217}
218
219/// `compat` for the OpenAI Responses family.
220#[derive(Clone, Debug, Default, PartialEq, Deserialize)]
221#[serde(rename_all = "camelCase")]
222pub struct OpenAiResponsesCompat {
223    pub supports_developer_role: Option<bool>,
224    pub supports_mid_convo_system_messages: Option<bool>,
225    pub session_affinity_format: Option<String>,
226    pub supports_long_cache_retention: Option<bool>,
227    pub supports_strict_mode: Option<bool>,
228    #[serde(rename = "supportsOpenAIGrammarTools")]
229    pub supports_open_ai_grammar_tools: Option<bool>,
230    pub supports_additional_tools: Option<bool>,
231    pub supports_tool_search: Option<bool>,
232    pub supports_explicit_prompt_cache_mode: Option<bool>,
233    pub supports_max_output_tokens: Option<bool>,
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    #[test]
241    fn reads_compat_views() {
242        let model: Model = serde_json::from_str(
243            r#"{"id":"m","name":"M","api":"openai-completions","provider":"p","baseUrl":"u",
244                "reasoning":true,"input":["text"],"cost":{"input":1,"output":2,"cacheRead":0,"cacheWrite":0},
245                "contextWindow":10,"maxTokens":5,
246                "compat":{"maxTokensField":"max_tokens","supportsStore":false,"supportsOpenAIGrammarTools":true,"unknown":1},
247                "thinkingLevelMap":{"minimal":null,"high":"high"}}"#,
248        )
249        .unwrap();
250        let compat: OpenAiCompletionsCompat = model.compat();
251        assert_eq!(compat.max_tokens_field.as_deref(), Some("max_tokens"));
252        assert_eq!(compat.supports_store, Some(false));
253        assert_eq!(compat.supports_open_ai_grammar_tools, Some(true));
254        assert_eq!(compat.supports_developer_role, None);
255        assert_eq!(
256            model.thinking_level_value(ThinkingLevel::Minimal),
257            Some(None)
258        );
259        assert_eq!(
260            model.thinking_level_value(ThinkingLevel::High),
261            Some(Some("high"))
262        );
263        assert_eq!(model.thinking_level_value(ThinkingLevel::Low), None);
264    }
265}