rig_core/
audio_generation.rs

1//! Everything related to audio generation (ie, Text To Speech).
2//! Rig abstracts over a number of different providers using the [AudioGenerationModel] trait.
3use crate::markers::{Missing, Provided};
4use crate::{
5    http_client, provider_response,
6    wasm_compat::{WasmCompatSend, WasmCompatSync},
7};
8use serde_json::Value;
9use thiserror::Error;
10
11/// Errors returned by audio generation models.
12///
13/// Inspect provider failures with [`Self::provider_response_body`],
14/// [`Self::provider_response_json`], and [`Self::provider_response_status`].
15///
16/// HTTP audio failures preserve the provider's status and body: a non-success
17/// response surfaces as [`Self::HttpError`], and a provider error envelope
18/// returned with a 2xx status surfaces as [`Self::ProviderResponse`] (for
19/// example the Hyperbolic audio path). Both are read by the helpers.
20#[derive(Debug, Error)]
21#[non_exhaustive]
22pub enum AudioGenerationError {
23    /// Http error (e.g.: connection error, timeout, etc.)
24    #[error("HttpError: {0}")]
25    HttpError(#[from] http_client::Error),
26
27    /// Json error (e.g.: serialization, deserialization)
28    #[error("JsonError: {0}")]
29    JsonError(#[from] serde_json::Error),
30
31    /// Error building the audio generation request
32    #[error("RequestError: {0}")]
33    RequestError(#[from] Box<dyn std::error::Error + Send + Sync + 'static>),
34
35    /// Error parsing the audio generation response
36    #[error("ResponseError: {0}")]
37    ResponseError(String),
38
39    /// Error returned by the audio generation model provider
40    #[error("ProviderError: {0}")]
41    ProviderError(String),
42
43    /// Raw error response preserved from the audio generation model provider
44    #[error("ProviderResponseError: {0}")]
45    ProviderResponse(provider_response::ProviderResponseError),
46}
47
48crate::provider_response::impl_provider_response_helpers!(AudioGenerationError);
49
50pub struct AudioGenerationResponse<T> {
51    pub audio: Vec<u8>,
52    pub response: T,
53}
54
55pub trait AudioGenerationModel: Sized + Clone + WasmCompatSend + WasmCompatSync {
56    type Response: Send + Sync;
57
58    type Client;
59
60    fn make(client: &Self::Client, model: impl Into<String>) -> Self;
61
62    fn audio_generation(
63        &self,
64        request: AudioGenerationRequest,
65    ) -> impl std::future::Future<
66        Output = Result<AudioGenerationResponse<Self::Response>, AudioGenerationError>,
67    > + Send;
68
69    fn audio_generation_request(&self) -> AudioGenerationRequestBuilder<Self, Missing, Missing> {
70        AudioGenerationRequestBuilder::new(self.clone())
71    }
72}
73#[non_exhaustive]
74pub struct AudioGenerationRequest {
75    pub text: String,
76    pub voice: String,
77    pub speed: f32,
78    pub additional_params: Option<Value>,
79}
80
81#[non_exhaustive]
82pub struct AudioGenerationRequestBuilder<M, T = Missing, V = Missing>
83where
84    M: AudioGenerationModel,
85{
86    model: M,
87    text: T,
88    voice: V,
89    speed: f32,
90    additional_params: Option<Value>,
91}
92
93impl<M> AudioGenerationRequestBuilder<M, Missing, Missing>
94where
95    M: AudioGenerationModel,
96{
97    pub fn new(model: M) -> Self {
98        Self {
99            model,
100            text: Missing,
101            voice: Missing,
102            speed: 1.0,
103            additional_params: None,
104        }
105    }
106}
107
108impl<M, T, V> AudioGenerationRequestBuilder<M, T, V>
109where
110    M: AudioGenerationModel,
111{
112    /// Sets the text for the audio generation request
113    pub fn text(self, text: &str) -> AudioGenerationRequestBuilder<M, Provided<String>, V> {
114        AudioGenerationRequestBuilder {
115            model: self.model,
116            text: Provided(text.to_string()),
117            voice: self.voice,
118            speed: self.speed,
119            additional_params: self.additional_params,
120        }
121    }
122
123    /// The voice of the generated audio
124    pub fn voice(self, voice: &str) -> AudioGenerationRequestBuilder<M, T, Provided<String>> {
125        AudioGenerationRequestBuilder {
126            model: self.model,
127            text: self.text,
128            voice: Provided(voice.to_string()),
129            speed: self.speed,
130            additional_params: self.additional_params,
131        }
132    }
133
134    /// The speed of the generated audio
135    pub fn speed(mut self, speed: f32) -> Self {
136        self.speed = speed;
137        self
138    }
139
140    /// Adds additional parameters to the audio generation request.
141    pub fn additional_params(mut self, params: Value) -> Self {
142        self.additional_params = Some(params);
143        self
144    }
145}
146
147impl<M> AudioGenerationRequestBuilder<M, Provided<String>, Provided<String>>
148where
149    M: AudioGenerationModel,
150{
151    pub fn build(self) -> AudioGenerationRequest {
152        AudioGenerationRequest {
153            text: self.text.0,
154            voice: self.voice.0,
155            speed: self.speed,
156            additional_params: self.additional_params,
157        }
158    }
159
160    pub async fn send(self) -> Result<AudioGenerationResponse<M::Response>, AudioGenerationError> {
161        let model = self.model.clone();
162
163        model.audio_generation(self.build()).await
164    }
165}
166
167#[cfg(test)]
168mod provider_response_tests {
169    use super::*;
170    use http::StatusCode;
171
172    #[test]
173    fn audio_generation_error_provider_response_helpers_with_preserved_json_body() {
174        let body = r#"{"error":{"message":"invalid voice"}}"#;
175        let error =
176            AudioGenerationError::ProviderResponse(provider_response::ProviderResponseError {
177                status: None,
178                body: body.to_string(),
179            });
180
181        assert_eq!(error.provider_response_body(), Some(body));
182        assert_eq!(error.provider_response_status(), None);
183        assert_eq!(
184            error.provider_response_json().expect("valid JSON"),
185            Some(serde_json::json!({ "error": { "message": "invalid voice" } }))
186        );
187    }
188
189    #[test]
190    fn audio_generation_error_provider_response_helpers_with_http_non_success() {
191        let body = r#"{"error":{"message":"bad request"}}"#;
192        let error =
193            AudioGenerationError::HttpError(http_client::Error::InvalidStatusCodeWithMessage(
194                StatusCode::BAD_REQUEST,
195                body.to_string(),
196            ));
197
198        assert_eq!(error.provider_response_body(), Some(body));
199        assert_eq!(
200            error.provider_response_status(),
201            Some(StatusCode::BAD_REQUEST)
202        );
203        assert_eq!(
204            error.provider_response_json().expect("valid JSON"),
205            Some(serde_json::json!({ "error": { "message": "bad request" } }))
206        );
207    }
208
209    #[test]
210    fn audio_generation_error_provider_error_is_not_a_provider_response() {
211        let error = AudioGenerationError::ProviderError("internal diagnostic".to_string());
212
213        assert_eq!(error.provider_response_body(), None);
214        assert_eq!(error.provider_response_status(), None);
215        assert_eq!(error.provider_response_json().expect("no body"), None);
216    }
217
218    #[test]
219    fn audio_generation_error_provider_response_helpers_with_unrelated_variant() {
220        let error = AudioGenerationError::ResponseError("parse failed".to_string());
221
222        assert_eq!(error.provider_response_body(), None);
223        assert_eq!(error.provider_response_status(), None);
224        assert_eq!(error.provider_response_json().expect("no body"), None);
225    }
226}
rig_core/audio_generation.rs

rig_core/
audio_generation.rs
