cookie_agent_models/authoring/
model.rs1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4
5use super::{HeaderName, PartialRequestDefaults, SafeStaticHeaderValue};
6use crate::{
7 CompactionCapability, MediaCapability, MediaKind, Modality, ModelCapabilities, ProviderOptions,
8 ReasoningBehavior, ReplayCapability, RequestDefaults, catalog::PicoUsdPerMillion,
9};
10
11#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
13#[serde(transparent)]
14pub struct WireModelId(String);
15
16impl WireModelId {
17 pub fn new(value: impl Into<String>) -> Result<Self, &'static str> {
18 let value = value.into();
19 oven_sdk::ModelId::new(&value)
20 .validate()
21 .map_err(|_| "invalid wire model_id")?;
22 if value.trim() != value || value.len() > 2048 {
23 return Err(
24 "model_id must be nonempty, at most 2048 bytes, and have no surrounding whitespace",
25 );
26 }
27 Ok(Self(value))
28 }
29
30 pub fn as_str(&self) -> &str {
31 &self.0
32 }
33}
34
35impl<'de> Deserialize<'de> for WireModelId {
36 fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
37 Self::new(String::deserialize(deserializer)?).map_err(serde::de::Error::custom)
38 }
39}
40
41#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
42#[serde(rename_all = "snake_case")]
43pub enum RequestEndpoint {
44 Completions,
45 Responses,
46}
47
48#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
49#[serde(deny_unknown_fields)]
50pub struct AdaptorOptions {
51 pub request_endpoint: Option<RequestEndpoint>,
52 pub beta: Option<Vec<String>>,
53 pub organization: Option<String>,
54 pub project: Option<String>,
55 pub store: Option<bool>,
56}
57
58impl AdaptorOptions {
59 pub(crate) fn apply(&self, base: &mut ProviderOptions) {
60 base.request_endpoint = self.request_endpoint.or(base.request_endpoint);
61 if let Some(beta) = &self.beta {
62 base.beta.clone_from(beta);
63 }
64 if let Some(value) = &self.organization {
65 base.organization = Some(value.clone());
66 }
67 if let Some(value) = &self.project {
68 base.project = Some(value.clone());
69 }
70 base.store = self.store.or(base.store);
71 }
72
73 pub(crate) fn resolve(&self) -> ProviderOptions {
74 let mut result = ProviderOptions::default();
75 self.apply(&mut result);
76 result
77 }
78}
79
80impl PartialRequestDefaults {
81 pub(crate) fn apply(&self, base: &mut RequestDefaults) {
82 base.temperature = self.temperature.or(base.temperature);
83 base.top_p = self.top_p.or(base.top_p);
84 base.max_output_tokens = self.max_output_tokens.or(base.max_output_tokens);
85 if let Some(stop) = &self.stop {
86 base.stop.clone_from(stop);
87 }
88 base.seed = self.seed.or(base.seed);
89 if let Some(choice) = &self.tool_choice {
90 base.tool_choice = Some(choice.clone());
91 }
92 }
93
94 pub(crate) fn resolve(&self) -> RequestDefaults {
95 let mut result = RequestDefaults::default();
96 self.apply(&mut result);
97 result
98 }
99}
100
101#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
102#[serde(deny_unknown_fields)]
103pub struct VariantDefinition {
104 pub enabled: Option<bool>,
105 pub model_id: Option<WireModelId>,
106 pub display_name: Option<String>,
107 pub generation_options: Option<PartialRequestDefaults>,
108 pub adaptor_options: Option<AdaptorOptions>,
109 pub reasoning: Option<ReasoningBehavior>,
110 #[serde(
111 default,
112 deserialize_with = "super::deserialize_optional_model_headers"
113 )]
114 pub headers: Option<BTreeMap<HeaderName, SafeStaticHeaderValue>>,
115}
116
117impl VariantDefinition {
118 pub(crate) fn disabled_has_settings(&self) -> bool {
119 self.model_id.is_some()
120 || self.display_name.is_some()
121 || self.generation_options.is_some()
122 || self.adaptor_options.is_some()
123 || self.reasoning.is_some()
124 || self.headers.is_some()
125 }
126
127 pub fn headers(&self) -> &BTreeMap<HeaderName, SafeStaticHeaderValue> {
128 static EMPTY: BTreeMap<HeaderName, SafeStaticHeaderValue> = BTreeMap::new();
129 self.headers.as_ref().unwrap_or(&EMPTY)
130 }
131}
132
133#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
134#[serde(deny_unknown_fields)]
135pub struct AuthoredCapabilities {
136 pub input: BTreeSet<Modality>,
137 pub output: BTreeSet<Modality>,
138 pub context_tokens: u64,
139 pub output_tokens: u64,
140 #[serde(default = "super::yes")]
141 pub tool_calling: bool,
142 #[serde(default = "super::yes")]
143 pub parallel_tool_calls: bool,
144 #[serde(default = "super::yes")]
145 pub structured_output: bool,
146 pub reasoning: bool,
147 pub temperature: bool,
148 pub top_p: bool,
149 pub seed: bool,
150 #[serde(default)]
151 pub compaction: CompactionCapability,
152 pub native_replay: Option<ReplayCapability>,
153 pub media: BTreeMap<MediaKind, MediaCapability>,
154}
155
156impl AuthoredCapabilities {
157 pub(crate) fn resolve(&self, adapter: crate::adapters::OvenAdapterFamily) -> ModelCapabilities {
158 ModelCapabilities {
159 input: self.input.clone(),
160 output: self.output.clone(),
161 context_tokens: self.context_tokens,
162 output_tokens: self.output_tokens,
163 tool_calling: self.tool_calling,
164 parallel_tool_calls: self.parallel_tool_calls,
165 structured_output: self.structured_output,
166 reasoning: self.reasoning,
167 temperature: self.temperature,
168 top_p: self.top_p,
169 seed: self.seed,
170 compaction: self.compaction,
171 native_replay: self
172 .native_replay
173 .unwrap_or_else(|| adapter.automatic_replay(self.reasoning)),
174 cancellation: crate::CancellationCapability::LocalOnly,
175 media: self.media.clone(),
176 }
177 }
178}
179
180#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
181#[serde(deny_unknown_fields)]
182pub struct ModelPricing {
183 pub input_per_million_usd: Option<PicoUsdPerMillion>,
184 pub output_per_million_usd: Option<PicoUsdPerMillion>,
185 pub reasoning_per_million_usd: Option<PicoUsdPerMillion>,
186 pub cache_read_per_million_usd: Option<PicoUsdPerMillion>,
187 pub cache_write_per_million_usd: Option<PicoUsdPerMillion>,
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193
194 #[test]
195 fn sparse_options_preserve_absence_and_apply_false_and_empty_lists() {
196 let mut options = ProviderOptions {
197 store: Some(true),
198 beta: vec!["inherited".into()],
199 organization: Some("org".into()),
200 ..ProviderOptions::default()
201 };
202 let empty: AdaptorOptions = toml::from_str("").unwrap();
203 empty.apply(&mut options);
204 assert_eq!(options.store, Some(true));
205 assert_eq!(options.beta, ["inherited"]);
206 let overlay: AdaptorOptions = toml::from_str("store = false\nbeta = []").unwrap();
207 overlay.apply(&mut options);
208 assert_eq!(options.store, Some(false));
209 assert!(options.beta.is_empty());
210 assert_eq!(options.organization.as_deref(), Some("org"));
211 }
212}