Skip to main content

kcode_openai_api/
model.rs

1use std::fmt;
2
3use serde_json::{Value, json};
4
5use crate::{DEFAULT_TRANSCRIPTION_PROMPT, Error, Result};
6
7pub(crate) const MAX_AUDIO_BYTES: usize = 25 * 1024 * 1024;
8const MAX_AUDIO_FILENAME_CHARACTERS: usize = 120;
9const MAX_TRANSCRIPTION_PROMPT_CHARACTERS: usize = 32_000;
10const MAX_IMAGE_PROMPT_CHARACTERS: usize = 32_000;
11const MIN_IMAGE_PIXELS: u64 = 655_360;
12const MAX_IMAGE_PIXELS: u64 = 8_294_400;
13const MAX_IMAGE_EDGE: u16 = 3_840;
14
15/// One in-memory audio file accepted by `gpt-4o-transcribe`.
16#[derive(Clone, Eq, PartialEq)]
17pub struct AudioInput {
18    file_name: String,
19    mime_type: String,
20    data: Vec<u8>,
21}
22
23impl AudioInput {
24    /// Constructs and validates an audio input.
25    pub fn new(
26        file_name: impl Into<String>,
27        mime_type: impl Into<String>,
28        data: Vec<u8>,
29    ) -> Result<Self> {
30        let value = Self {
31            file_name: file_name.into(),
32            mime_type: mime_type.into().to_ascii_lowercase(),
33            data,
34        };
35        value.validate()?;
36        Ok(value)
37    }
38
39    /// Returns the validated upload filename.
40    pub fn file_name(&self) -> &str {
41        &self.file_name
42    }
43
44    /// Returns the validated audio MIME type.
45    pub fn mime_type(&self) -> &str {
46        &self.mime_type
47    }
48
49    /// Returns the raw audio bytes.
50    pub fn data(&self) -> &[u8] {
51        &self.data
52    }
53
54    /// Returns the raw audio byte count.
55    pub fn len(&self) -> usize {
56        self.data.len()
57    }
58
59    /// Returns whether the audio input contains no bytes.
60    pub fn is_empty(&self) -> bool {
61        self.data.is_empty()
62    }
63
64    pub(crate) fn into_parts(self) -> (String, String, Vec<u8>) {
65        (self.file_name, self.mime_type, self.data)
66    }
67
68    fn validate(&self) -> Result<()> {
69        if self.file_name.is_empty()
70            || self.file_name.chars().count() > MAX_AUDIO_FILENAME_CHARACTERS
71            || !self.file_name.chars().all(|character| {
72                character.is_ascii_alphanumeric() || matches!(character, '.' | '-' | '_')
73            })
74            || matches!(self.file_name.as_str(), "." | "..")
75        {
76            return Err(Error::InvalidInput(format!(
77                "audio filename must contain 1 through {MAX_AUDIO_FILENAME_CHARACTERS} ASCII letters, digits, dots, hyphens, or underscores"
78            )));
79        }
80        if !matches!(
81            self.mime_type.as_str(),
82            "audio/flac"
83                | "audio/x-flac"
84                | "audio/m4a"
85                | "audio/mp3"
86                | "audio/mp4"
87                | "audio/mpeg"
88                | "audio/mpga"
89                | "audio/ogg"
90                | "audio/opus"
91                | "audio/wav"
92                | "audio/x-wav"
93                | "audio/webm"
94                | "application/ogg"
95                | "video/mp4"
96                | "video/webm"
97        ) {
98            return Err(Error::InvalidInput(
99                "audio MIME type must describe a supported FLAC, MP3, MP4, M4A, OGG, WAV, or WebM recording".into(),
100            ));
101        }
102        let extension = self
103            .file_name
104            .rsplit_once('.')
105            .map(|(_, extension)| extension.to_ascii_lowercase());
106        if !matches!(
107            extension.as_deref(),
108            Some("flac" | "mp3" | "mp4" | "mpeg" | "mpga" | "m4a" | "ogg" | "wav" | "webm")
109        ) {
110            return Err(Error::InvalidInput(
111                "audio filename must use a supported flac, mp3, mp4, mpeg, mpga, m4a, ogg, wav, or webm extension".into(),
112            ));
113        }
114        if self.data.is_empty() || self.data.len() > MAX_AUDIO_BYTES {
115            return Err(Error::InvalidInput(format!(
116                "audio must contain between 1 and {MAX_AUDIO_BYTES} bytes"
117            )));
118        }
119        Ok(())
120    }
121}
122
123impl fmt::Debug for AudioInput {
124    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
125        f.debug_struct("AudioInput")
126            .field("file_name", &self.file_name)
127            .field("mime_type", &self.mime_type)
128            .field("bytes", &self.data.len())
129            .finish()
130    }
131}
132
133/// One audio transcription request.
134#[derive(Clone, Debug, Eq, PartialEq)]
135pub struct TranscriptionRequest {
136    /// Audio file to transcribe.
137    pub audio: AudioInput,
138    /// Optional text that guides transcription style and context.
139    pub prompt: Option<String>,
140    /// Optional ISO-639-1 input language code, such as `en`.
141    pub language: Option<String>,
142}
143
144impl TranscriptionRequest {
145    /// Constructs a request with Kennedy's current faithful-transcription prompt.
146    pub fn new(audio: AudioInput) -> Self {
147        Self {
148            audio,
149            prompt: Some(DEFAULT_TRANSCRIPTION_PROMPT.into()),
150            language: None,
151        }
152    }
153
154    pub(crate) fn validate(&self) -> Result<()> {
155        self.audio.validate()?;
156        if let Some(prompt) = &self.prompt
157            && (prompt.trim().is_empty()
158                || prompt.chars().count() > MAX_TRANSCRIPTION_PROMPT_CHARACTERS)
159        {
160            return Err(Error::InvalidInput(format!(
161                "transcription prompt must contain 1 through {MAX_TRANSCRIPTION_PROMPT_CHARACTERS} characters when supplied"
162            )));
163        }
164        if let Some(language) = &self.language
165            && (language.len() != 2 || !language.bytes().all(|value| value.is_ascii_lowercase()))
166        {
167            return Err(Error::InvalidInput(
168                "transcription language must be a two-letter lowercase ISO-639-1 code".into(),
169            ));
170        }
171        Ok(())
172    }
173}
174
175/// Modality detail for token-billed audio transcription input.
176#[derive(Clone, Debug, Default, Eq, PartialEq)]
177pub struct TranscriptionTokenDetails {
178    /// Audio tokens billed for the request, when reported.
179    pub audio_tokens: Option<u64>,
180    /// Prompt text tokens billed for the request, when reported.
181    pub text_tokens: Option<u64>,
182}
183
184/// Token-billed transcription usage.
185#[derive(Clone, Debug, Eq, PartialEq)]
186pub struct TranscriptionTokenUsage {
187    /// Input tokens billed for the request.
188    pub input_tokens: u64,
189    /// Output tokens generated by the request.
190    pub output_tokens: u64,
191    /// Total input and output tokens.
192    pub total_tokens: u64,
193    /// Optional input modality breakdown.
194    pub input_details: Option<TranscriptionTokenDetails>,
195}
196
197/// Usage returned for a transcription request.
198#[derive(Clone, Debug, PartialEq)]
199pub enum TranscriptionUsage {
200    /// Usage billed in tokens.
201    Tokens(TranscriptionTokenUsage),
202    /// Usage billed from audio duration in seconds.
203    DurationSeconds(f64),
204}
205
206/// Normalized `gpt-4o-transcribe` result.
207#[derive(Clone, Debug, PartialEq)]
208pub struct Transcription {
209    /// Complete non-empty transcript text.
210    pub text: String,
211    /// Provider usage, when returned.
212    pub usage: Option<TranscriptionUsage>,
213    /// OpenAI request identifier, when returned.
214    pub request_id: Option<String>,
215}
216
217/// Generated image dimensions.
218#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
219pub enum ImageSize {
220    /// Let GPT Image choose the dimensions.
221    #[default]
222    Auto,
223    /// Explicit width and height in pixels.
224    Dimensions {
225        /// Width in pixels.
226        width: u16,
227        /// Height in pixels.
228        height: u16,
229    },
230}
231
232impl ImageSize {
233    /// Constructs explicit dimensions after validating GPT Image 2 constraints.
234    pub fn dimensions(width: u16, height: u16) -> Result<Self> {
235        let value = Self::Dimensions { width, height };
236        value.validate()?;
237        Ok(value)
238    }
239
240    pub(crate) fn as_api_value(self) -> String {
241        match self {
242            Self::Auto => "auto".into(),
243            Self::Dimensions { width, height } => format!("{width}x{height}"),
244        }
245    }
246
247    pub(crate) fn validate(self) -> Result<()> {
248        let Self::Dimensions { width, height } = self else {
249            return Ok(());
250        };
251        let pixels = u64::from(width).saturating_mul(u64::from(height));
252        let short = width.min(height);
253        let long = width.max(height);
254        if width % 16 != 0
255            || height % 16 != 0
256            || width > MAX_IMAGE_EDGE
257            || height > MAX_IMAGE_EDGE
258            || short == 0
259            || u32::from(long) > u32::from(short).saturating_mul(3)
260            || !(MIN_IMAGE_PIXELS..=MAX_IMAGE_PIXELS).contains(&pixels)
261        {
262            return Err(Error::InvalidInput(format!(
263                "GPT Image 2 dimensions must be multiples of 16, no edge may exceed {MAX_IMAGE_EDGE}, the aspect ratio must be at most 3:1, and total pixels must be between {MIN_IMAGE_PIXELS} and {MAX_IMAGE_PIXELS}"
264            )));
265        }
266        Ok(())
267    }
268}
269
270/// GPT Image rendering quality.
271#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
272pub enum ImageQuality {
273    /// Let GPT Image select quality.
274    #[default]
275    Auto,
276    /// Fast draft quality.
277    Low,
278    /// Balanced quality.
279    Medium,
280    /// Highest supported quality.
281    High,
282}
283
284impl ImageQuality {
285    pub(crate) const fn as_str(self) -> &'static str {
286        match self {
287            Self::Auto => "auto",
288            Self::Low => "low",
289            Self::Medium => "medium",
290            Self::High => "high",
291        }
292    }
293
294    pub(crate) fn parse(value: &str) -> Option<Self> {
295        match value {
296            "auto" => Some(Self::Auto),
297            "low" => Some(Self::Low),
298            "medium" => Some(Self::Medium),
299            "high" => Some(Self::High),
300            _ => None,
301        }
302    }
303}
304
305/// GPT Image output file format.
306#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
307pub enum ImageFormat {
308    /// PNG output.
309    #[default]
310    Png,
311    /// JPEG output.
312    Jpeg,
313    /// WebP output.
314    WebP,
315}
316
317impl ImageFormat {
318    /// Returns the output MIME type.
319    pub const fn mime_type(self) -> &'static str {
320        match self {
321            Self::Png => "image/png",
322            Self::Jpeg => "image/jpeg",
323            Self::WebP => "image/webp",
324        }
325    }
326
327    pub(crate) const fn as_str(self) -> &'static str {
328        match self {
329            Self::Png => "png",
330            Self::Jpeg => "jpeg",
331            Self::WebP => "webp",
332        }
333    }
334
335    pub(crate) fn parse(value: &str) -> Option<Self> {
336        match value {
337            "png" => Some(Self::Png),
338            "jpeg" => Some(Self::Jpeg),
339            "webp" => Some(Self::WebP),
340            _ => None,
341        }
342    }
343}
344
345/// Background selection supported by GPT Image 2.
346#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
347pub enum ImageBackground {
348    /// Let GPT Image select the background.
349    #[default]
350    Auto,
351    /// Require an opaque background.
352    Opaque,
353}
354
355impl ImageBackground {
356    pub(crate) const fn as_str(self) -> &'static str {
357        match self {
358            Self::Auto => "auto",
359            Self::Opaque => "opaque",
360        }
361    }
362}
363
364/// Provider content-moderation level.
365#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq)]
366pub enum Moderation {
367    /// Standard provider moderation.
368    #[default]
369    Auto,
370    /// Less restrictive provider moderation.
371    Low,
372}
373
374impl Moderation {
375    pub(crate) const fn as_str(self) -> &'static str {
376        match self {
377            Self::Auto => "auto",
378            Self::Low => "low",
379        }
380    }
381}
382
383/// One GPT Image 2 text-to-image request.
384#[derive(Clone, Debug, Eq, PartialEq)]
385pub struct ImageGenerationRequest {
386    /// Text description of the desired image.
387    pub prompt: String,
388    /// Output dimensions.
389    pub size: ImageSize,
390    /// Rendering quality.
391    pub quality: ImageQuality,
392    /// Output file format.
393    pub output_format: ImageFormat,
394    /// JPEG or WebP compression from 0 through 100, when explicitly requested.
395    pub output_compression: Option<u8>,
396    /// Background behavior. GPT Image 2 does not support transparency.
397    pub background: ImageBackground,
398    /// Provider moderation level.
399    pub moderation: Moderation,
400    /// Optional stable end-user identifier supplied to OpenAI abuse monitoring.
401    pub user: Option<String>,
402}
403
404impl ImageGenerationRequest {
405    /// Constructs a single-image request with automatic size and quality and PNG output.
406    pub fn new(prompt: impl Into<String>) -> Self {
407        Self {
408            prompt: prompt.into(),
409            size: ImageSize::Auto,
410            quality: ImageQuality::Auto,
411            output_format: ImageFormat::Png,
412            output_compression: None,
413            background: ImageBackground::Auto,
414            moderation: Moderation::Auto,
415            user: None,
416        }
417    }
418
419    pub(crate) fn validate(&self) -> Result<()> {
420        if self.prompt.trim().is_empty()
421            || self.prompt.chars().count() > MAX_IMAGE_PROMPT_CHARACTERS
422        {
423            return Err(Error::InvalidInput(format!(
424                "image prompt must contain 1 through {MAX_IMAGE_PROMPT_CHARACTERS} characters"
425            )));
426        }
427        self.size.validate()?;
428        if self.output_compression.is_some_and(|value| value > 100) {
429            return Err(Error::InvalidInput(
430                "output compression must be between 0 and 100".into(),
431            ));
432        }
433        if self.output_compression.is_some() && self.output_format == ImageFormat::Png {
434            return Err(Error::InvalidInput(
435                "output compression is supported only for JPEG and WebP images".into(),
436            ));
437        }
438        if let Some(user) = &self.user
439            && (user.trim().is_empty()
440                || user.chars().count() > 512
441                || user.chars().any(char::is_control))
442        {
443            return Err(Error::InvalidInput(
444                "image user identifier must contain 1 through 512 non-control characters when supplied".into(),
445            ));
446        }
447        Ok(())
448    }
449
450    pub(crate) fn payload(&self) -> Value {
451        let mut payload = json!({
452            "model": crate::GPT_IMAGE_2,
453            "prompt": self.prompt,
454            "n": 1,
455            "size": self.size.as_api_value(),
456            "quality": self.quality.as_str(),
457            "output_format": self.output_format.as_str(),
458            "background": self.background.as_str(),
459            "moderation": self.moderation.as_str(),
460            "stream": false
461        });
462        let object = payload.as_object_mut().expect("image payload is an object");
463        if let Some(compression) = self.output_compression {
464            object.insert("output_compression".into(), json!(compression));
465        }
466        if let Some(user) = &self.user {
467            object.insert("user".into(), json!(user));
468        }
469        payload
470    }
471}
472
473/// One generated image held in memory.
474#[derive(Clone, Eq, PartialEq)]
475pub struct GeneratedImage {
476    /// Decoded image bytes.
477    pub data: Vec<u8>,
478    /// Returned image file format.
479    pub format: ImageFormat,
480}
481
482impl fmt::Debug for GeneratedImage {
483    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
484        f.debug_struct("GeneratedImage")
485            .field("format", &self.format)
486            .field("bytes", &self.data.len())
487            .finish()
488    }
489}
490
491/// Text and image token detail for GPT Image usage.
492#[derive(Clone, Debug, Default, Eq, PartialEq)]
493pub struct ImageTokenDetails {
494    /// Text tokens reported for the modality side.
495    pub text_tokens: u64,
496    /// Image tokens reported for the modality side.
497    pub image_tokens: u64,
498}
499
500/// Token usage returned by GPT Image.
501#[derive(Clone, Debug, Eq, PartialEq)]
502pub struct ImageUsage {
503    /// Total input tokens.
504    pub input_tokens: u64,
505    /// Total output tokens.
506    pub output_tokens: u64,
507    /// Total input and output tokens.
508    pub total_tokens: u64,
509    /// Input text and image token breakdown.
510    pub input_details: ImageTokenDetails,
511    /// Output text and image token breakdown, when returned.
512    pub output_details: Option<ImageTokenDetails>,
513}
514
515/// Normalized single-image GPT Image 2 result.
516#[derive(Clone, Debug, PartialEq)]
517pub struct ImageGeneration {
518    /// Provider creation time as Unix seconds.
519    pub created: u64,
520    /// Decoded generated image.
521    pub image: GeneratedImage,
522    /// Actual provider-selected size, when returned.
523    pub size: Option<String>,
524    /// Actual provider-selected quality, when returned.
525    pub quality: Option<ImageQuality>,
526    /// Provider token usage, when returned.
527    pub usage: Option<ImageUsage>,
528    /// OpenAI request identifier, when returned.
529    pub request_id: Option<String>,
530}
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535
536    #[test]
537    fn audio_debug_omits_bytes_and_rejects_unsafe_names() {
538        let audio = AudioInput::new("note.webm", "audio/webm", vec![7, 8, 9]).unwrap();
539        let debug = format!("{audio:?}");
540        assert!(debug.contains("bytes: 3"));
541        assert!(!debug.contains("7, 8, 9"));
542        assert!(AudioInput::new("../note.webm", "audio/webm", vec![1]).is_err());
543    }
544
545    #[test]
546    fn image_dimensions_enforce_current_gpt_image_2_constraints() {
547        assert_eq!(
548            ImageSize::dimensions(2048, 2048).unwrap(),
549            ImageSize::Dimensions {
550                width: 2048,
551                height: 2048
552            }
553        );
554        assert!(ImageSize::dimensions(1000, 1000).is_err());
555        assert!(ImageSize::dimensions(3840, 3840).is_err());
556        assert!(ImageSize::dimensions(3072, 1024).is_ok());
557        assert!(ImageSize::dimensions(3088, 1024).is_err());
558    }
559
560    #[test]
561    fn png_rejects_compression_but_jpeg_accepts_it() {
562        let mut request = ImageGenerationRequest::new("draw a lighthouse");
563        request.output_compression = Some(80);
564        assert!(request.validate().is_err());
565        request.output_format = ImageFormat::Jpeg;
566        assert!(request.validate().is_ok());
567        request.output_compression = Some(101);
568        assert!(request.validate().is_err());
569    }
570
571    #[test]
572    fn image_payload_is_single_shot_gpt_image_2() {
573        let request = ImageGenerationRequest::new("draw a lighthouse");
574        let payload = request.payload();
575        assert_eq!(payload["model"], "gpt-image-2");
576        assert_eq!(payload["n"], 1);
577        assert_eq!(payload["stream"], false);
578        assert_eq!(payload["background"], "auto");
579        assert_eq!(payload["output_format"], "png");
580    }
581}