rig_core/providers/openai/
audio_generation.rs1use 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}