Skip to main content

rig_core/providers/huggingface/
transcription.rs

1use 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    /// Name of the model (e.g.: gpt-3.5-turbo-1106)
39    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        // A 200 OK body that is not a valid `TranscriptionResponse` (no `text`
157        // field) falls through the untagged `ApiResponse` to its `Err(Value)`
158        // variant, which the provider routes through `from_http_response`.
159        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}