zai_rs/model/audio_to_text/
data.rs1use 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
32pub struct SpeechToTextStream {
35 inner: crate::model::sse_parser::DecodedSseStream<SpeechToTextEvent>,
36}
37
38impl SpeechToTextStream {
39 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
55pub 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 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 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 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 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 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 pub fn body(&self) -> &AudioToTextBody<N> {
149 &self.body
150 }
151
152 pub fn file_path(&self) -> Option<&Path> {
154 self.file_path.as_deref()
155 }
156
157 pub fn has_file_base64(&self) -> bool {
162 self.file_base64.is_some()
163 }
164
165 pub fn with_file_path(mut self, path: impl Into<PathBuf>) -> Self {
168 self.file_path = Some(path.into());
169 self
170 }
171
172 pub fn with_file_base64(mut self, encoded: impl Into<String>) -> Self {
175 self.file_base64 = Some(encoded.into());
176 self
177 }
178
179 pub fn with_prompt(mut self, prompt: impl Into<String>) -> Self {
182 self.body = self.body.with_prompt(prompt);
183 self
184 }
185
186 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 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 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 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}