Skip to main content

rig_core/providers/openai/extension/
chat.rs

1//! The Chat Completions section of OpenAI's options, and the token details
2//! only a Chat reply carries.
3
4use std::collections::BTreeMap;
5
6use serde::{Deserialize, Serialize};
7
8/// The fields only Chat Completions takes
9/// (<https://developers.openai.com/api/reference/resources/chat/subresources/completions/methods/create>).
10#[non_exhaustive]
11#[derive(Clone, Debug, Default, PartialEq, Serialize)]
12pub struct ChatOptions {
13    /// Bias per token id, from -100 to 100.
14    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
15    pub logit_bias: BTreeMap<String, i32>,
16    /// Predicted output, which speeds up a regeneration that mostly repeats
17    /// it.
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub prediction: Option<Prediction>,
20    /// Whether the reply carries each output token's log probability.
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub logprobs: Option<bool>,
23    /// How many of the most likely tokens to report per position, 0 to 20.
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub top_logprobs: Option<u8>,
26    /// Penalizes tokens by how often they already appear, -2.0 to 2.0.
27    #[serde(skip_serializing_if = "Option::is_none")]
28    pub frequency_penalty: Option<f64>,
29    /// Penalizes tokens that already appear, -2.0 to 2.0.
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub presence_penalty: Option<f64>,
32    /// The output modalities, for models that also speak.
33    #[serde(skip_serializing_if = "Vec::is_empty")]
34    pub modalities: Vec<Modality>,
35    /// The spoken output, when `modalities` holds audio.
36    #[serde(skip_serializing_if = "Option::is_none")]
37    pub audio: Option<AudioOutput>,
38    /// Web search, for the search models.
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub web_search_options: Option<WebSearchOptions>,
41}
42
43impl ChatOptions {
44    /// No option set.
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Bias token id `token` by `bias`.
50    pub fn logit_bias(mut self, token: u32, bias: i32) -> Self {
51        self.logit_bias.insert(token.to_string(), bias);
52        self
53    }
54
55    /// Predict the output as `content`.
56    pub fn prediction(mut self, content: impl Into<String>) -> Self {
57        self.prediction = Some(Prediction::content(content));
58        self
59    }
60
61    /// Whether the reply carries log probabilities.
62    pub fn logprobs(mut self, logprobs: bool) -> Self {
63        self.logprobs = Some(logprobs);
64        self
65    }
66
67    /// Report the `count` most likely tokens per position.
68    pub fn top_logprobs(mut self, count: u8) -> Self {
69        self.top_logprobs = Some(count);
70        self
71    }
72
73    /// Set the frequency penalty.
74    pub fn frequency_penalty(mut self, penalty: f64) -> Self {
75        self.frequency_penalty = Some(penalty);
76        self
77    }
78
79    /// Set the presence penalty.
80    pub fn presence_penalty(mut self, penalty: f64) -> Self {
81        self.presence_penalty = Some(penalty);
82        self
83    }
84
85    /// Answer in `modalities`.
86    pub fn modalities(mut self, modalities: impl IntoIterator<Item = Modality>) -> Self {
87        self.modalities = modalities.into_iter().collect();
88        self
89    }
90
91    /// Speak with `voice` in `format`.
92    pub fn audio(mut self, voice: impl Into<String>, format: AudioFormat) -> Self {
93        self.audio = Some(AudioOutput {
94            voice: voice.into(),
95            format,
96        });
97        self
98    }
99
100    /// Search the web with `options`.
101    pub fn web_search_options(mut self, options: WebSearchOptions) -> Self {
102        self.web_search_options = Some(options);
103        self
104    }
105}
106
107/// Predicted output.
108#[non_exhaustive]
109#[derive(Clone, Debug, PartialEq, Serialize)]
110#[serde(tag = "type", rename_all = "lowercase")]
111pub enum Prediction {
112    /// Text the output is expected to mostly repeat.
113    Content {
114        /// The predicted text.
115        content: String,
116    },
117}
118
119impl Prediction {
120    /// A prediction of `content`.
121    pub fn content(content: impl Into<String>) -> Self {
122        Self::Content {
123            content: content.into(),
124        }
125    }
126}
127
128/// An output modality.
129#[non_exhaustive]
130#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
131#[serde(rename_all = "lowercase")]
132pub enum Modality {
133    /// Text.
134    Text,
135    /// Speech.
136    Audio,
137}
138
139/// The spoken output's voice and encoding.
140#[non_exhaustive]
141#[derive(Clone, Debug, PartialEq, Serialize)]
142pub struct AudioOutput {
143    /// The voice, such as `alloy`.
144    pub voice: String,
145    /// The encoding.
146    pub format: AudioFormat,
147}
148
149/// A spoken output's encoding.
150#[non_exhaustive]
151#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
152#[serde(rename_all = "lowercase")]
153pub enum AudioFormat {
154    /// WAV.
155    Wav,
156    /// AAC.
157    Aac,
158    /// MP3.
159    Mp3,
160    /// FLAC.
161    Flac,
162    /// Opus.
163    Opus,
164    /// Raw 16-bit PCM.
165    Pcm16,
166}
167
168/// How a search model searches the web.
169#[non_exhaustive]
170#[derive(Clone, Debug, Default, PartialEq, Serialize)]
171pub struct WebSearchOptions {
172    /// How much search context goes into the prompt.
173    #[serde(skip_serializing_if = "Option::is_none")]
174    pub search_context_size: Option<SearchContextSize>,
175    /// Where the user is, to localize results.
176    #[serde(skip_serializing_if = "Option::is_none")]
177    pub user_location: Option<UserLocation>,
178}
179
180impl WebSearchOptions {
181    /// The defaults.
182    pub fn new() -> Self {
183        Self::default()
184    }
185
186    /// Put `size` search context into the prompt.
187    pub fn search_context_size(mut self, size: SearchContextSize) -> Self {
188        self.search_context_size = Some(size);
189        self
190    }
191
192    /// Localize results to `location`.
193    pub fn user_location(mut self, location: ApproximateLocation) -> Self {
194        self.user_location = Some(UserLocation::Approximate {
195            approximate: location,
196        });
197        self
198    }
199}
200
201/// How much search context goes into the prompt.
202#[non_exhaustive]
203#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)]
204#[serde(rename_all = "lowercase")]
205pub enum SearchContextSize {
206    /// Least context, cheapest.
207    Low,
208    /// The default.
209    Medium,
210    /// Most context.
211    High,
212}
213
214/// Where the user is.
215#[non_exhaustive]
216#[derive(Clone, Debug, PartialEq, Serialize)]
217#[serde(tag = "type", rename_all = "lowercase")]
218pub enum UserLocation {
219    /// An approximate location.
220    Approximate {
221        /// The location.
222        approximate: ApproximateLocation,
223    },
224}
225
226/// An approximate location; every field is optional.
227#[non_exhaustive]
228#[derive(Clone, Debug, Default, PartialEq, Serialize)]
229pub struct ApproximateLocation {
230    /// The city, such as `San Francisco`.
231    #[serde(skip_serializing_if = "Option::is_none")]
232    pub city: Option<String>,
233    /// The two-letter ISO country code, such as `US`.
234    #[serde(skip_serializing_if = "Option::is_none")]
235    pub country: Option<String>,
236    /// The region, such as `California`.
237    #[serde(skip_serializing_if = "Option::is_none")]
238    pub region: Option<String>,
239    /// The IANA time zone, such as `America/Los_Angeles`.
240    #[serde(skip_serializing_if = "Option::is_none")]
241    pub timezone: Option<String>,
242}
243
244impl ApproximateLocation {
245    /// No field set.
246    pub fn new() -> Self {
247        Self::default()
248    }
249
250    /// Set the city.
251    pub fn city(mut self, city: impl Into<String>) -> Self {
252        self.city = Some(city.into());
253        self
254    }
255
256    /// Set the two-letter country code.
257    pub fn country(mut self, country: impl Into<String>) -> Self {
258        self.country = Some(country.into());
259        self
260    }
261
262    /// Set the region.
263    pub fn region(mut self, region: impl Into<String>) -> Self {
264        self.region = Some(region.into());
265        self
266    }
267
268    /// Set the IANA time zone.
269    pub fn timezone(mut self, timezone: impl Into<String>) -> Self {
270        self.timezone = Some(timezone.into());
271        self
272    }
273}
274
275/// Where a Chat reply's prompt tokens went.
276#[non_exhaustive]
277#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
278pub struct PromptTokensDetails {
279    /// Tokens read from the prompt cache.
280    #[serde(default)]
281    pub cached_tokens: Option<u64>,
282    /// Audio input tokens.
283    #[serde(default)]
284    pub audio_tokens: Option<u64>,
285    /// Tokens written to the prompt cache.
286    #[serde(default)]
287    pub cache_write_tokens: Option<u64>,
288}
289
290/// Where a Chat reply's completion tokens went.
291#[non_exhaustive]
292#[derive(Clone, Debug, Default, PartialEq, Eq, Deserialize)]
293pub struct CompletionTokensDetails {
294    /// Tokens spent reasoning.
295    #[serde(default)]
296    pub reasoning_tokens: Option<u64>,
297    /// Audio output tokens.
298    #[serde(default)]
299    pub audio_tokens: Option<u64>,
300    /// Predicted tokens that appeared in the output.
301    #[serde(default)]
302    pub accepted_prediction_tokens: Option<u64>,
303    /// Predicted tokens that did not, billed as output.
304    #[serde(default)]
305    pub rejected_prediction_tokens: Option<u64>,
306}