Skip to main content

cookie_agent_models/authoring/
model.rs

1use 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/// Provider wire identity, deliberately independent of local model-key syntax.
12#[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}