Skip to main content

rig_core/providers/xai/
audio_generation.rs

1use crate::audio_generation::{
2    self, AudioGenerationError, AudioGenerationRequest, AudioGenerationResponse,
3};
4use crate::http_client::{self, HttpClientExt};
5use crate::json_utils::merge_inplace;
6use crate::providers::xai::Client;
7use bytes::Bytes;
8use serde_json::json;
9
10// ================================================================
11// xAI TTS API
12// ================================================================
13pub const TTS_1: &str = "tts-1";
14
15#[derive(Clone)]
16pub struct AudioGenerationModel<T = reqwest::Client> {
17    client: Client<T>,
18    pub model: String,
19}
20
21impl<T> AudioGenerationModel<T> {
22    pub(crate) fn new(client: Client<T>, model: impl Into<String>) -> Self {
23        Self {
24            client,
25            model: model.into(),
26        }
27    }
28}
29
30impl<T> audio_generation::AudioGenerationModel for AudioGenerationModel<T>
31where
32    T: HttpClientExt + Clone + std::fmt::Debug + Default + 'static,
33{
34    type Response = Bytes;
35
36    type Client = Client<T>;
37
38    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
39        Self::new(client.clone(), model)
40    }
41
42    async fn audio_generation(
43        &self,
44        request: AudioGenerationRequest,
45    ) -> Result<AudioGenerationResponse<Self::Response>, AudioGenerationError> {
46        let voice = if request.voice.is_empty() {
47            "eve".to_string()
48        } else {
49            request.voice
50        };
51
52        let mut body = json!({
53            "text": request.text,
54            "voice_id": voice,
55            "language": "en",
56        });
57
58        if let Some(additional_params) = request.additional_params {
59            merge_inplace(&mut body, additional_params);
60        }
61
62        let body = serde_json::to_vec(&body)?;
63
64        let req = self
65            .client
66            .post("/v1/tts")?
67            .body(body)
68            .map_err(http_client::Error::from)?;
69
70        let response = self.client.send(req).await?;
71
72        if !response.status().is_success() {
73            let status = response.status();
74            let text = http_client::text(response).await?;
75
76            return Err(AudioGenerationError::from_http_response(status, text));
77        }
78
79        let bytes: Bytes = response.into_body().await?.into();
80
81        Ok(AudioGenerationResponse {
82            audio: bytes.to_vec(),
83            response: bytes,
84        })
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::audio_generation::AudioGenerationModel as _;
92    use crate::client::audio_generation::AudioGenerationClient;
93
94    #[tokio::test]
95    async fn audio_generation_non_success_preserves_status_and_body() {
96        use crate::test_utils::RecordingHttpClient;
97
98        let body = r#"{"error":"boom","code":"503"}"#;
99        let http_client =
100            RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
101        let client = crate::providers::xai::Client::builder()
102            .api_key("test-key")
103            .http_client(http_client)
104            .build()
105            .expect("build client");
106        let model = client.audio_generation_model(TTS_1);
107
108        let request = model
109            .audio_generation_request()
110            .text("hello")
111            .voice("eve")
112            .build();
113
114        let error = model
115            .audio_generation(request)
116            .await
117            .err()
118            .expect("should fail with non-success status");
119
120        assert!(matches!(error, AudioGenerationError::HttpError(_)));
121        assert_eq!(
122            error.provider_response_status(),
123            Some(http::StatusCode::SERVICE_UNAVAILABLE)
124        );
125        assert_eq!(error.provider_response_body(), Some(body));
126    }
127}