rig_core/providers/huggingface/
transcription.rs1use crate::http_client::HttpClientExt;
2use crate::providers::huggingface::Client;
3use crate::providers::huggingface::completion::ApiResponse;
4use crate::transcription;
5use crate::transcription::TranscriptionError;
6use crate::wasm_compat::WasmCompatSync;
7use base64::Engine;
8use base64::prelude::BASE64_STANDARD;
9use serde::Deserialize;
10use serde_json::json;
11
12pub const WHISPER_LARGE_V3: &str = "openai/whisper-large-v3";
13pub const WHISPER_LARGE_V3_TURBO: &str = "openai/whisper-large-v3-turbo";
14
15pub const WHISPER_SMALL: &str = "openai/whisper-small";
16
17#[derive(Debug, Deserialize)]
18pub struct TranscriptionResponse {
19 pub text: String,
20}
21
22impl TryFrom<TranscriptionResponse>
23 for transcription::TranscriptionResponse<TranscriptionResponse>
24{
25 type Error = TranscriptionError;
26
27 fn try_from(value: TranscriptionResponse) -> Result<Self, Self::Error> {
28 Ok(transcription::TranscriptionResponse {
29 text: value.text.clone(),
30 response: value,
31 })
32 }
33}
34
35#[derive(Clone)]
36pub struct TranscriptionModel<T = reqwest::Client> {
37 client: Client<T>,
38 pub model: String,
40}
41
42impl<T> TranscriptionModel<T> {
43 pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
44 Self {
45 client,
46 model: model.into(),
47 }
48 }
49}
50impl<T> transcription::TranscriptionModel for TranscriptionModel<T>
51where
52 T: HttpClientExt + Clone + WasmCompatSync + 'static,
53{
54 type Response = TranscriptionResponse;
55
56 type Client = Client<T>;
57
58 fn make(client: &Self::Client, model: impl Into<String>) -> Self {
59 TranscriptionModel::new(client.clone(), model)
60 }
61
62 async fn transcription(
63 &self,
64 request: transcription::TranscriptionRequest,
65 ) -> Result<transcription::TranscriptionResponse<Self::Response>, TranscriptionError> {
66 let data = request.data;
67 let data = BASE64_STANDARD.encode(data);
68
69 let request = json!({
70 "inputs": data
71 });
72
73 let route = self
74 .client
75 .subprovider()
76 .transcription_endpoint(&self.model)?;
77
78 let request = serde_json::to_vec(&request)?;
79
80 let req = self
81 .client
82 .post(&route)?
83 .header("Content-Type", "application/json")
84 .body(request)
85 .map_err(|e| TranscriptionError::HttpError(e.into()))?;
86
87 let response = self.client.send(req).await?;
88 let status = response.status();
89 let body: Vec<u8> = response.into_body().await?;
90
91 if !status.is_success() {
92 return Err(TranscriptionError::from_http_response(
93 status,
94 String::from_utf8_lossy(&body),
95 ));
96 }
97
98 match serde_json::from_slice::<ApiResponse<TranscriptionResponse>>(&body)? {
99 ApiResponse::Ok(response) => response.try_into(),
100 ApiResponse::Err(err) => {
101 let message = err
102 .get("error")
103 .and_then(|e| {
104 e.as_str()
105 .or_else(|| e.get("message").and_then(|m| m.as_str()))
106 })
107 .or_else(|| err.get("message").and_then(|m| m.as_str()))
108 .unwrap_or_default();
109 tracing::warn!(message = %message, "provider returned an error response");
110 Err(TranscriptionError::from_http_response(
111 status,
112 String::from_utf8_lossy(&body),
113 ))
114 }
115 }
116 }
117}
118
119#[cfg(test)]
120mod tests {
121 use super::*;
122 use crate::client::transcription::TranscriptionClient;
123 use crate::test_utils::RecordingHttpClient;
124 use crate::transcription::TranscriptionModel as _;
125
126 #[tokio::test]
127 async fn transcription_non_success_preserves_status_and_body() {
128 let body = r#"{"error":{"message":"boom"}}"#;
129 let http_client =
130 RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
131 let client = Client::builder()
132 .api_key("test-key")
133 .http_client(http_client)
134 .build()
135 .expect("build client");
136 let model = client.transcription_model(WHISPER_LARGE_V3);
137
138 let request = model.transcription_request().data(vec![0u8; 16]).build();
139
140 let error = model
141 .transcription(request)
142 .await
143 .err()
144 .expect("should fail with non-success status");
145
146 assert!(matches!(error, TranscriptionError::HttpError(_)));
147 assert_eq!(
148 error.provider_response_status(),
149 Some(http::StatusCode::SERVICE_UNAVAILABLE)
150 );
151 assert_eq!(error.provider_response_body(), Some(body));
152 }
153
154 #[tokio::test]
155 async fn transcription_2xx_error_envelope_preserves_status_and_body() {
156 let body = r#"{"error":"Model openai/whisper-large-v3 is currently loading"}"#;
160 let http_client = RecordingHttpClient::new(body);
161 let client = Client::builder()
162 .api_key("test-key")
163 .http_client(http_client)
164 .build()
165 .expect("build client");
166 let model = client.transcription_model(WHISPER_LARGE_V3);
167
168 let request = model.transcription_request().data(vec![0u8; 16]).build();
169
170 let error = model
171 .transcription(request)
172 .await
173 .err()
174 .expect("should fail with provider error envelope");
175
176 match &error {
177 TranscriptionError::ProviderResponse(stored) => {
178 assert_eq!(stored.body, body);
179 assert_eq!(stored.status, Some(http::StatusCode::OK));
180 }
181 other => panic!("expected ProviderResponse, got {other:?}"),
182 }
183 }
184}