zai_rs/model/text_to_audio/
data.rs

1use std::{
2    marker::PhantomData,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7use base64::Engine as _;
8use futures_util::{Stream, StreamExt};
9use validator::Validate;
10
11use super::{
12    super::traits::{StreamOff, StreamOn, StreamState, TextToAudio},
13    request::{TextToAudioBody, TtsAudioFormat, TtsEncodeFormat, Voice},
14};
15use crate::{
16    ZaiError, ZaiResult,
17    client::{ZaiClient, error::codes},
18};
19
20/// Typed PCM byte stream returned by [`TextToAudioRequest::stream_via`].
21pub struct TextToAudioStream {
22    inner: crate::model::sse_parser::DecodedSseStream<bytes::Bytes>,
23}
24
25impl TextToAudioStream {
26    /// Await the next decoded PCM chunk, or `None` after `[DONE]`.
27    ///
28    /// Protocol and decode failures are yielded once before termination.
29    pub async fn next(&mut self) -> Option<ZaiResult<bytes::Bytes>> {
30        self.inner.next().await
31    }
32}
33
34impl Stream for TextToAudioStream {
35    type Item = ZaiResult<bytes::Bytes>;
36
37    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
38        self.inner.as_mut().poll_next(context)
39    }
40}
41
42/// Type-state request builder for text-to-speech synthesis.
43///
44/// Non-streaming requests return complete WAV or PCM bytes through
45/// [`send_via`](Self::send_via). [`enable_stream`](Self::enable_stream)
46/// switches to PCM-only SSE output decoded by [`stream_via`](Self::stream_via).
47pub struct TextToAudioRequest<N, S = StreamOff>
48where
49    N: TextToAudio,
50    S: StreamState,
51{
52    body: TextToAudioBody<N>,
53    _stream: PhantomData<S>,
54}
55
56impl<N> TextToAudioRequest<N, StreamOff>
57where
58    N: TextToAudio,
59{
60    /// Create a non-streaming TTS request for the selected model.
61    pub fn new(model: N) -> Self {
62        Self {
63            body: TextToAudioBody::new(model),
64            _stream: PhantomData,
65        }
66    }
67
68    /// Select WAV or PCM for a complete non-streaming response.
69    pub fn with_response_format(mut self, format: TtsAudioFormat) -> Self {
70        self.body = self.body.with_response_format(format);
71        self
72    }
73
74    /// Switch to PCM-only SSE output with base64 payloads by default.
75    pub fn enable_stream(self) -> TextToAudioRequest<N, StreamOn> {
76        TextToAudioRequest {
77            body: self.body.with_stream(true),
78            _stream: PhantomData,
79        }
80    }
81
82    /// Submit a non-streaming request and return its complete audio bytes.
83    pub async fn send_via(&self, client: &ZaiClient) -> ZaiResult<bytes::Bytes> {
84        self.validate()?;
85        let route = crate::client::routes::AUDIO_SYNTHESIZE;
86        let url = client.endpoints().resolve_route(route, &[])?;
87        client
88            .send_json_bytes(route.method(), url, &self.body)
89            .await
90    }
91}
92
93impl<N> TextToAudioRequest<N, StreamOn>
94where
95    N: TextToAudio,
96{
97    /// Select standard base64 or hexadecimal SSE audio payloads.
98    pub fn with_encode_format(mut self, format: TtsEncodeFormat) -> Self {
99        self.body = self.body.with_encode_format(format);
100        self
101    }
102
103    /// Return to non-streaming PCM output and remove `encode_format`.
104    pub fn disable_stream(self) -> TextToAudioRequest<N, StreamOff> {
105        TextToAudioRequest {
106            body: self.body.with_stream(false),
107            _stream: PhantomData,
108        }
109    }
110
111    /// Submit the request and yield decoded PCM chunks from its SSE response.
112    ///
113    /// The streaming POST is never retried or redirected. A missing `[DONE]`,
114    /// malformed encoded chunk, in-band business error, oversized event, or
115    /// idle timeout is returned once and then terminates the stream.
116    pub async fn stream_via(&self, client: &ZaiClient) -> ZaiResult<TextToAudioStream> {
117        self.validate()?;
118        let Some(encoding) = self.body.encode_format else {
119            return Err(ZaiError::ApiError {
120                code: codes::SDK_VALIDATION,
121                message: "streaming TTS requires encode_format".to_string(),
122            });
123        };
124        let route = crate::client::routes::AUDIO_SYNTHESIZE;
125        let url = client.endpoints().resolve_route(route, &[])?;
126        let raw = client
127            .send_sse_json(route.method(), url, &self.body)
128            .await?;
129        let inner = crate::model::sse_parser::decode_required_done_stream(raw, move |payload| {
130            decode_audio_payload(payload, encoding)
131        });
132        Ok(TextToAudioStream { inner })
133    }
134}
135
136impl<N, S> TextToAudioRequest<N, S>
137where
138    N: TextToAudio,
139    S: StreamState,
140{
141    /// Borrow the request body without exposing its type-state invariants for
142    /// mutation.
143    pub fn body(&self) -> &TextToAudioBody<N> {
144        &self.body
145    }
146
147    /// Set the required input text (`1..=1024` Unicode scalar values).
148    pub fn with_input(mut self, input: impl Into<String>) -> Self {
149        self.body = self.body.with_input(input);
150        self
151    }
152
153    /// Select a built-in or cloned voice.
154    pub fn with_voice(mut self, voice: Voice) -> Self {
155        self.body = self.body.with_voice(voice);
156        self
157    }
158
159    /// Set playback speed in `0.5..=2.0`.
160    pub fn with_speed(mut self, speed: f32) -> Self {
161        self.body = self.body.with_speed(speed);
162        self
163    }
164
165    /// Set playback volume in `0 < volume <= 10`.
166    pub fn with_volume(mut self, volume: f32) -> Self {
167        self.body = self.body.with_volume(volume);
168        self
169    }
170
171    /// Enable or disable the service's audio watermark.
172    pub fn with_watermark_enabled(mut self, enabled: bool) -> Self {
173        self.body = self.body.with_watermark_enabled(enabled);
174        self
175    }
176
177    /// Validate all wire constraints without network I/O.
178    pub fn validate(&self) -> ZaiResult<()> {
179        self.body.validate().map_err(ZaiError::from)
180    }
181}
182
183fn decode_audio_payload(payload: &[u8], encoding: TtsEncodeFormat) -> ZaiResult<bytes::Bytes> {
184    // The frozen TTS contract defines each SSE `data:` payload as the direct
185    // `encode_format` text, not as a JSON envelope. Decode that exact payload;
186    // JSON strings or surrounding whitespace are intentionally rejected.
187    let encoded = std::str::from_utf8(payload)
188        .map_err(|_| stream_protocol_error("TTS SSE audio payload must be UTF-8 encoded text"))?;
189    if encoded.is_empty() {
190        return Err(stream_protocol_error(
191            "TTS SSE audio payload must not be empty",
192        ));
193    }
194    let decoded = match encoding {
195        TtsEncodeFormat::Base64 => base64::engine::general_purpose::STANDARD
196            .decode(encoded)
197            .map_err(|_| stream_protocol_error("invalid base64 TTS SSE audio payload"))?,
198        TtsEncodeFormat::Hex => decode_hex(encoded)?,
199    };
200    Ok(bytes::Bytes::from(decoded))
201}
202
203fn decode_hex(encoded: &str) -> ZaiResult<Vec<u8>> {
204    if !encoded.len().is_multiple_of(2) {
205        return Err(stream_protocol_error(
206            "hex TTS SSE audio payload must have an even length",
207        ));
208    }
209    encoded
210        .as_bytes()
211        .chunks_exact(2)
212        .map(|pair| {
213            let high = hex_nibble(pair[0])?;
214            let low = hex_nibble(pair[1])?;
215            Ok((high << 4) | low)
216        })
217        .collect()
218}
219
220fn hex_nibble(byte: u8) -> ZaiResult<u8> {
221    match byte {
222        b'0'..=b'9' => Ok(byte - b'0'),
223        b'a'..=b'f' => Ok(byte - b'a' + 10),
224        b'A'..=b'F' => Ok(byte - b'A' + 10),
225        _ => Err(stream_protocol_error(
226            "invalid hexadecimal TTS SSE audio payload",
227        )),
228    }
229}
230
231fn stream_protocol_error(message: impl Into<String>) -> ZaiError {
232    ZaiError::ApiError {
233        code: codes::SDK_IO,
234        message: message.into(),
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241    use crate::model::text_to_audio::GlmTts;
242
243    #[test]
244    fn streaming_type_state_forces_pcm_and_scopes_encoding() {
245        let request = TextToAudioRequest::new(GlmTts {})
246            .with_input("hello")
247            .with_response_format(TtsAudioFormat::Wav)
248            .enable_stream();
249        assert!(request.body().is_streaming());
250        assert_eq!(request.body().response_format(), TtsAudioFormat::Pcm);
251        assert_eq!(
252            request.body().encode_format(),
253            Some(TtsEncodeFormat::Base64)
254        );
255
256        let request = request
257            .with_encode_format(TtsEncodeFormat::Hex)
258            .disable_stream();
259        assert!(!request.body().is_streaming());
260        assert_eq!(request.body().response_format(), TtsAudioFormat::Pcm);
261        assert_eq!(request.body().encode_format(), None);
262    }
263
264    #[test]
265    fn base64_and_hex_payloads_decode_to_bytes() {
266        assert_eq!(
267            decode_audio_payload(b"AAEC/w==", TtsEncodeFormat::Base64).unwrap(),
268            bytes::Bytes::from_static(&[0, 1, 2, 255])
269        );
270        assert_eq!(
271            decode_audio_payload(b"000102fF", TtsEncodeFormat::Hex).unwrap(),
272            bytes::Bytes::from_static(&[0, 1, 2, 255])
273        );
274        assert!(decode_audio_payload(b"%%%", TtsEncodeFormat::Base64).is_err());
275        assert!(decode_audio_payload(br#""AAEC/w==""#, TtsEncodeFormat::Base64).is_err());
276        assert!(decode_audio_payload(b" AAEC/w==", TtsEncodeFormat::Base64).is_err());
277        assert!(decode_audio_payload(b"123", TtsEncodeFormat::Hex).is_err());
278        assert!(decode_audio_payload(b"12xz", TtsEncodeFormat::Hex).is_err());
279    }
280}
zai_rs/model/text_to_audio/data.rs

zai_rs/model/text_to_audio/
data.rs
