Skip to main content

zai_rs/model/audio_to_text/
data.rs

1use std::{
2    marker::PhantomData,
3    path::{Path, PathBuf},
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use base64::Engine as _;
9use futures_util::{Stream, StreamExt};
10use validator::Validate;
11
12use super::{
13    super::traits::{AudioToText, StreamOff, StreamOn, StreamState},
14    request::AudioToTextBody,
15    response::{AudioToTextResponse, SpeechToTextEvent},
16};
17use crate::{
18    ZaiError, ZaiResult,
19    client::{ZaiClient, error::codes},
20};
21
22const ASR_FILE_MAX_BYTES: u64 = 25 * 1024 * 1024;
23const ASR_BASE64_MAX_BYTES: u64 = ASR_FILE_MAX_BYTES.div_ceil(3) * 4;
24const ASR_BASE64_MULTIPART_LIMIT: u64 =
25    ASR_BASE64_MAX_BYTES + crate::client::transport::limits::MULTIPART_FIELD_BYTES_MAX + 64;
26
27enum AudioInputRef<'a> {
28    File(&'a Path),
29    Base64(&'a str),
30}
31
32/// Authenticated typed response stream returned by
33/// [`AudioToTextRequest::stream_via`].
34pub struct SpeechToTextStream {
35    inner: crate::model::sse_parser::DecodedSseStream<SpeechToTextEvent>,
36}
37
38impl SpeechToTextStream {
39    /// Await the next transcription event, or `None` after `[DONE]`.
40    ///
41    /// Protocol and decode failures are yielded once before termination.
42    pub async fn next(&mut self) -> Option<ZaiResult<SpeechToTextEvent>> {
43        self.inner.next().await
44    }
45}
46
47impl Stream for SpeechToTextStream {
48    type Item = ZaiResult<SpeechToTextEvent>;
49
50    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
51        self.inner.as_mut().poll_next(context)
52    }
53}
54
55/// Type-state request builder for speech-to-text transcription.
56///
57/// Exactly one input must be configured with
58/// [`with_file_path`](Self::with_file_path) or
59/// [`with_file_base64`](Self::with_file_base64). Non-streaming requests use
60/// [`send_via`](Self::send_via); [`enable_stream`](Self::enable_stream)
61/// switches the builder to the streaming-only [`stream_via`](Self::stream_via)
62/// API.
63pub struct AudioToTextRequest<N, S = StreamOff>
64where
65    N: AudioToText,
66    S: StreamState,
67{
68    body: AudioToTextBody<N>,
69    file_path: Option<PathBuf>,
70    file_base64: Option<String>,
71    _stream: PhantomData<S>,
72}
73
74impl<N> AudioToTextRequest<N, StreamOff>
75where
76    N: AudioToText,
77{
78    /// Create a non-streaming ASR request for the selected model.
79    pub fn new(model: N) -> Self {
80        Self {
81            body: AudioToTextBody::new(model),
82            file_path: None,
83            file_base64: None,
84            _stream: PhantomData,
85        }
86    }
87
88    /// Switch to the typed SSE response API and serialize `stream=true`.
89    pub fn enable_stream(self) -> AudioToTextRequest<N, StreamOn> {
90        AudioToTextRequest {
91            body: self.body.with_stream(true),
92            file_path: self.file_path,
93            file_base64: self.file_base64,
94            _stream: PhantomData,
95        }
96    }
97
98    /// Submit a non-streaming request and decode its required response fields.
99    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<AudioToTextResponse> {
100        let factory = self.build_multipart().await?;
101        let route = crate::client::routes::AUDIO_TRANSCRIBE;
102        let url = client.endpoints().resolve_route(route, &[])?;
103        client
104            .send_multipart::<AudioToTextResponse>(route.method(), url, &factory)
105            .await
106    }
107}
108
109impl<N> AudioToTextRequest<N, StreamOn>
110where
111    N: AudioToText,
112{
113    /// Return to the non-streaming response API and serialize `stream=false`.
114    pub fn disable_stream(self) -> AudioToTextRequest<N, StreamOff> {
115        AudioToTextRequest {
116            body: self.body.with_stream(false),
117            file_path: self.file_path,
118            file_base64: self.file_base64,
119            _stream: PhantomData,
120        }
121    }
122
123    /// Submit the multipart request and decode typed transcription SSE events.
124    ///
125    /// The streaming POST is never retried or redirected. A missing `[DONE]`,
126    /// malformed event, in-band business error, oversized event, or idle
127    /// timeout is returned once and then terminates the stream.
128    pub async fn stream_via(&self, client: &ZaiClient) -> ZaiResult<SpeechToTextStream> {
129        let factory = self.build_multipart().await?;
130        let route = crate::client::routes::AUDIO_TRANSCRIBE;
131        let url = client.endpoints().resolve_route(route, &[])?;
132        let raw = client
133            .send_sse_multipart(route.method(), url, &factory)
134            .await?;
135        let inner = crate::model::sse_parser::decode_required_done_stream(raw, |payload| {
136            serde_json::from_slice::<SpeechToTextEvent>(payload).map_err(ZaiError::from)
137        });
138        Ok(SpeechToTextStream { inner })
139    }
140}
141
142impl<N, S> AudioToTextRequest<N, S>
143where
144    N: AudioToText,
145    S: StreamState,
146{
147    /// Borrow the multipart metadata body.
148    pub fn body(&self) -> &AudioToTextBody<N> {
149        &self.body
150    }
151
152    /// Borrow the configured local input path, if any.
153    pub fn file_path(&self) -> Option<&Path> {
154        self.file_path.as_deref()
155    }
156
157    /// Whether a base64 input is configured.
158    ///
159    /// The encoded audio itself is intentionally not exposed by an accessor or
160    /// by `Debug` output.
161    pub fn has_file_base64(&self) -> bool {
162        self.file_base64.is_some()
163    }
164
165    /// Select a local WAV or MP3 file. Supplying base64 as well is an error;
166    /// setters do not silently override one another.
167    pub fn with_file_path(mut self, path: impl Into<PathBuf>) -> Self {
168        self.file_path = Some(path.into());
169        self
170    }
171
172    /// Select standard-base64 encoded WAV or MP3 bytes. Supplying a local file
173    /// as well is an error.
174    pub fn with_file_base64(mut self, encoded: impl Into<String>) -> Self {
175        self.file_base64 = Some(encoded.into());
176        self
177    }
178
179    /// Set optional prior-transcript context. The upstream 8,000-character
180    /// recommendation remains advisory and is not hard-rejected by the SDK.
181    pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
182        self.body = self.body.with_prompt(prompt);
183        self
184    }
185
186    /// Set at most 100 non-blank hot words.
187    pub fn with_hotwords(mut self, hotwords: Vec<String>) -> ZaiResult<Self> {
188        self.body = self.body.with_hotwords(hotwords)?;
189        Ok(self)
190    }
191
192    /// Set the client request identifier (`6..=64` characters).
193    pub fn with_request_id(mut self, request_id: impl Into<String>) -> Self {
194        self.body = self.body.with_request_id(request_id);
195        self
196    }
197
198    /// Set the end-user identifier (`6..=128` characters).
199    pub fn with_user_id(mut self, user_id: impl Into<String>) -> Self {
200        self.body = self.body.with_user_id(user_id);
201        self
202    }
203
204    /// Validate metadata and the selected input without network I/O.
205    pub fn validate(&self) -> ZaiResult<()> {
206        self.validate_common()?;
207        match self.input()? {
208            AudioInputRef::File(path) => validate_local_file(path),
209            AudioInputRef::Base64(encoded) => validate_base64_audio(encoded),
210        }
211    }
212
213    fn validate_common(&self) -> ZaiResult<()> {
214        self.body.validate().map_err(ZaiError::from)?;
215        let _ = self.input()?;
216        Ok(())
217    }
218
219    fn input(&self) -> ZaiResult<AudioInputRef<'_>> {
220        match (self.file_path.as_deref(), self.file_base64.as_deref()) {
221            (Some(path), None) => Ok(AudioInputRef::File(path)),
222            (None, Some(encoded)) => Ok(AudioInputRef::Base64(encoded)),
223            (Some(_), Some(_)) => Err(validation_error(
224                "ASR input requires file XOR file_base64; both were provided",
225            )),
226            (None, None) => Err(validation_error(
227                "ASR input requires exactly one of file or file_base64",
228            )),
229        }
230    }
231
232    async fn build_multipart(
233        &self,
234    ) -> ZaiResult<crate::client::transport::multipart::MultipartBodyFactory> {
235        self.validate_common()?;
236        let mut factory = crate::client::transport::multipart::MultipartBodyFactory::new()
237            .field("model", N::NAME)?
238            .field("stream", self.body.stream.to_string())?;
239        if let Some(prompt) = &self.body.prompt {
240            factory = factory.field("prompt", prompt.clone())?;
241        }
242        for hotword in &self.body.hotwords {
243            factory = factory.field("hotwords", hotword.clone())?;
244        }
245        if let Some(request_id) = &self.body.request_id {
246            factory = factory.field("request_id", request_id.clone())?;
247        }
248        if let Some(user_id) = &self.body.user_id {
249            factory = factory.field("user_id", user_id.clone())?;
250        }
251
252        match self.input()? {
253            AudioInputRef::File(path) => {
254                let mime = audio_mime_type(path)?;
255                let part = crate::client::transport::multipart::FilePart::from_path_async(path)
256                    .await?
257                    .with_content_type(mime)?;
258                if part.len() > ASR_FILE_MAX_BYTES {
259                    return Err(file_error(
260                        codes::SDK_FILE_TOO_LARGE,
261                        "ASR audio exceeds the 25 MiB limit",
262                    ));
263                }
264                factory.file_named("file", part)
265            },
266            AudioInputRef::Base64(encoded) => {
267                validate_base64_audio(encoded)?;
268                factory.field_with_total_limit(
269                    "file_base64",
270                    encoded.to_owned(),
271                    ASR_BASE64_MULTIPART_LIMIT,
272                )
273            },
274        }
275    }
276}
277
278fn validate_local_file(path: &Path) -> ZaiResult<()> {
279    let _ = audio_mime_type(path)?;
280    let metadata = std::fs::symlink_metadata(path).map_err(|error| {
281        if error.kind() == std::io::ErrorKind::NotFound {
282            file_error(codes::SDK_FILE_NOT_FOUND, "ASR file_path was not found")
283        } else {
284            ZaiError::from(error)
285        }
286    })?;
287    if !metadata.is_file() {
288        return Err(file_error(
289            codes::SDK_FILE_NOT_FOUND,
290            "ASR file_path must be a regular, non-symlink file",
291        ));
292    }
293    if metadata.len() > ASR_FILE_MAX_BYTES {
294        return Err(file_error(
295            codes::SDK_FILE_TOO_LARGE,
296            "ASR audio exceeds the 25 MiB limit",
297        ));
298    }
299    Ok(())
300}
301
302fn validate_base64_audio(encoded: &str) -> ZaiResult<()> {
303    if encoded.len() as u64 > ASR_BASE64_MAX_BYTES {
304        return Err(file_error(
305            codes::SDK_FILE_TOO_LARGE,
306            "ASR base64 audio exceeds the 25 MiB decoded limit",
307        ));
308    }
309    let decoded = base64::engine::general_purpose::STANDARD
310        .decode(encoded)
311        .map_err(|_| validation_error("file_base64 must use valid standard base64"))?;
312    if decoded.len() as u64 > ASR_FILE_MAX_BYTES {
313        return Err(file_error(
314            codes::SDK_FILE_TOO_LARGE,
315            "ASR base64 audio exceeds the 25 MiB decoded limit",
316        ));
317    }
318    if !is_wav(&decoded) && !is_mp3(&decoded) {
319        return Err(file_error(
320            codes::SDK_FILE_TYPE_UNSUPPORTED,
321            "file_base64 must contain WAV or MP3 audio",
322        ));
323    }
324    Ok(())
325}
326
327fn is_wav(bytes: &[u8]) -> bool {
328    bytes.len() >= 12 && bytes.starts_with(b"RIFF") && &bytes[8..12] == b"WAVE"
329}
330
331fn is_mp3(bytes: &[u8]) -> bool {
332    bytes.starts_with(b"ID3") || (bytes.len() >= 2 && bytes[0] == 0xff && bytes[1] & 0xe0 == 0xe0)
333}
334
335fn audio_mime_type(path: &Path) -> ZaiResult<&'static str> {
336    match path
337        .extension()
338        .and_then(|extension| extension.to_str())
339        .map(str::to_ascii_lowercase)
340        .as_deref()
341    {
342        Some("wav") => Ok("audio/wav"),
343        Some("mp3") => Ok("audio/mpeg"),
344        _ => Err(file_error(
345            codes::SDK_FILE_TYPE_UNSUPPORTED,
346            "ASR file_path must use a .wav or .mp3 extension",
347        )),
348    }
349}
350
351fn validation_error(message: impl Into<String>) -> ZaiError {
352    ZaiError::ApiError {
353        code: codes::SDK_VALIDATION,
354        message: message.into(),
355    }
356}
357
358fn file_error(code: u16, message: impl Into<String>) -> ZaiError {
359    ZaiError::FileError {
360        code,
361        message: message.into(),
362    }
363}
364
365#[cfg(test)]
366mod tests {
367    use super::*;
368    use crate::model::audio_to_text::GlmAsr;
369
370    fn wav_base64() -> String {
371        base64::engine::general_purpose::STANDARD.encode(b"RIFF\x04\0\0\0WAVE")
372    }
373
374    #[test]
375    fn input_is_exactly_one_of_file_or_base64() {
376        let neither = AudioToTextRequest::new(GlmAsr {});
377        assert_eq!(
378            neither.validate().unwrap_err().code(),
379            Some(codes::SDK_VALIDATION)
380        );
381
382        let both = AudioToTextRequest::new(GlmAsr {})
383            .with_file_path("audio.wav")
384            .with_file_base64(wav_base64());
385        let error = both.validate().unwrap_err();
386        assert_eq!(error.code(), Some(codes::SDK_VALIDATION));
387        assert!(error.message().contains("XOR"));
388    }
389
390    #[test]
391    fn base64_requires_wav_or_mp3_and_obeys_decoded_limit() {
392        assert!(
393            AudioToTextRequest::new(GlmAsr {})
394                .with_file_base64(wav_base64())
395                .validate()
396                .is_ok()
397        );
398        assert!(
399            AudioToTextRequest::new(GlmAsr {})
400                .with_file_base64(base64::engine::general_purpose::STANDARD.encode(b"not audio"))
401                .validate()
402                .is_err()
403        );
404        assert!(
405            AudioToTextRequest::new(GlmAsr {})
406                .with_file_base64("%%%")
407                .validate()
408                .is_err()
409        );
410    }
411
412    #[test]
413    fn stream_type_state_sets_the_wire_flag() {
414        let request = AudioToTextRequest::new(GlmAsr {}).with_file_base64(wav_base64());
415        assert!(!request.body().is_streaming());
416        let request = request.enable_stream();
417        assert!(request.body().is_streaming());
418        assert!(!request.disable_stream().body().is_streaming());
419    }
420
421    #[test]
422    fn local_input_accepts_only_wav_or_mp3_up_to_25_mib() {
423        let directory = tempfile::tempdir().unwrap();
424        let wav = directory.path().join("audio.WAV");
425        std::fs::write(&wav, b"small").unwrap();
426        assert!(
427            AudioToTextRequest::new(GlmAsr {})
428                .with_file_path(&wav)
429                .validate()
430                .is_ok()
431        );
432
433        let unsupported = directory.path().join("audio.flac");
434        std::fs::write(&unsupported, b"small").unwrap();
435        assert_eq!(
436            AudioToTextRequest::new(GlmAsr {})
437                .with_file_path(&unsupported)
438                .validate()
439                .unwrap_err()
440                .code(),
441            Some(codes::SDK_FILE_TYPE_UNSUPPORTED)
442        );
443
444        let oversized = directory.path().join("large.mp3");
445        let file = std::fs::File::create(&oversized).unwrap();
446        file.set_len(ASR_FILE_MAX_BYTES + 1).unwrap();
447        assert_eq!(
448            AudioToTextRequest::new(GlmAsr {})
449                .with_file_path(&oversized)
450                .validate()
451                .unwrap_err()
452                .code(),
453            Some(codes::SDK_FILE_TOO_LARGE)
454        );
455    }
456}