Skip to main content

rig_core/providers/openai/
audio_generation.rs

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