Skip to main content

rig_core/providers/mistral/
transcription.rs

1//! Implements Mistral (basic) transcription API
2use bytes::Bytes;
3use serde::Deserialize;
4
5use crate::http_client::multipart::Part;
6use crate::http_client::{HttpClientExt, MultipartForm};
7use crate::providers::mistral::Client;
8use crate::transcription::{self, TranscriptionError};
9use crate::wasm_compat::WasmCompatSend;
10
11// ================================================================
12// Mistral Transcription API
13// ================================================================
14
15/// Voxtral Mini model (latest version)
16pub const VOXTRAL_MINI: &str = "voxtral-mini-latest";
17/// Voxtral Small model (latest version)
18pub const VOXTRAL_SMALL: &str = "voxtral-small-latest";
19
20/// Request usage statistics
21#[derive(Debug, Deserialize)]
22pub struct TranscriptionUsage {
23    pub prompt_audio_seconds: Option<i32>,
24    pub prompt_tokens: i32,
25    pub total_tokens: i32,
26    pub completion_tokens: i32,
27    pub prompt_tokens_details: Option<serde_json::Value>,
28}
29
30impl std::fmt::Display for TranscriptionUsage {
31    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        writeln!(f, "Usage:")?;
33        writeln!(f, "  prompt_tokens:     {}", self.prompt_tokens)?;
34        writeln!(f, "  completion_tokens: {}", self.completion_tokens)?;
35        writeln!(f, "  total_tokens:      {}", self.total_tokens)?;
36        if let Some(details) = &self.prompt_tokens_details {
37            writeln!(f, "  prompt_token_details: {:?}", details)?;
38        } else {
39            writeln!(f, "  prompt_token_details: N/A")?;
40        }
41        if let Some(secs) = self.prompt_audio_seconds {
42            write!(f, "  audio_seconds:     {secs}")?;
43        } else {
44            write!(f, "  audio_seconds:     N/A")?;
45        }
46        Ok(())
47    }
48}
49
50/// Diarization information, tells when each speaker started and ended talking plus what they said.
51#[derive(Debug, Deserialize)]
52pub struct SegmentChunk {
53    /// Start time in seconds
54    pub start: f32,
55    /// End time in seconds
56    pub end: f32,
57    /// Segment transcribed text
58    pub text: String,
59    pub score: Option<f32>,
60    /// Speaker identification.
61    pub speaker_id: Option<String>,
62    #[serde(rename = "type")]
63    pub segment_type: String,
64}
65
66#[derive(Debug, Deserialize)]
67pub struct MistralTranscriptionResponse {
68    /// Audio language
69    pub language: Option<String>,
70    /// Model name (e.g. voxtra-mini-latest)
71    pub model: String,
72    /// An array of transcript segments, each containing a portion of the transcribed text along with its start and end times in seconds and speaker id (if diarization was enabled).
73    pub segments: Vec<SegmentChunk>,
74    /// Audio Transcription
75    pub text: String,
76    /// Request token usage statistics
77    pub usage: TranscriptionUsage,
78}
79
80impl TryFrom<MistralTranscriptionResponse>
81    for transcription::TranscriptionResponse<MistralTranscriptionResponse>
82{
83    type Error = TranscriptionError;
84
85    fn try_from(value: MistralTranscriptionResponse) -> Result<Self, Self::Error> {
86        Ok(transcription::TranscriptionResponse {
87            text: value.text.clone(),
88            response: value,
89        })
90    }
91}
92
93#[derive(Clone)]
94pub struct TranscriptionModel<T = reqwest::Client> {
95    client: Client<T>,
96    pub model: String,
97}
98
99impl<T> transcription::TranscriptionModel for TranscriptionModel<T>
100where
101    T: HttpClientExt + Clone + std::fmt::Debug + Default + WasmCompatSend + 'static,
102{
103    type Response = MistralTranscriptionResponse;
104    type Client = Client<T>;
105
106    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
107        Self::new(client.clone(), model)
108    }
109
110    async fn transcription(
111        &self,
112        request: transcription::TranscriptionRequest,
113    ) -> Result<transcription::TranscriptionResponse<Self::Response>, TranscriptionError> {
114        let data = request.data;
115
116        let mut body = MultipartForm::new()
117            .text("model", self.model.clone())
118            .part(Part::bytes("file", data).filename(request.filename.clone()));
119
120        if let Some(language) = request.language {
121            body = body.text("language", language);
122        }
123
124        if let Some(ref temperature) = request.temperature {
125            body = body.text("temperature", temperature.to_string());
126        }
127
128        if let Some(ref additional_params) = request.additional_params {
129            for (key, value) in additional_params.as_object().ok_or_else(|| {
130                TranscriptionError::RequestError(
131                    "Additional Parameters to Mistral Transcription should be a map".into(),
132                )
133            })? {
134                body = body.text(key.to_owned(), value.to_string());
135            }
136        }
137
138        let req = self
139            .client
140            .post("/v1/audio/transcriptions")?
141            .body(body)
142            .map_err(|e| TranscriptionError::RequestError(e.into()))?;
143
144        let response = self
145            .client
146            .send_multipart::<Bytes>(req)
147            .await
148            .map_err(TranscriptionError::HttpError)?;
149
150        let status = response.status();
151        let response_bytes = response.into_body().await?;
152
153        if status.is_success() {
154            let response_body: MistralTranscriptionResponse =
155                serde_json::from_slice(&response_bytes)?;
156
157            tracing::info!(target: "rig", "Mistral transcription token usage: {}", &response_body.usage);
158
159            Ok(transcription::TranscriptionResponse::try_from(
160                response_body,
161            )?)
162        } else {
163            Err(TranscriptionError::from_http_response(
164                status,
165                String::from_utf8_lossy(&response_bytes),
166            ))
167        }
168    }
169}
170
171impl<T> TranscriptionModel<T> {
172    pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
173        Self {
174            client,
175            model: model.into(),
176        }
177    }
178}
179
180#[cfg(test)]
181mod test {
182    use super::*;
183    use crate::transcription::TranscriptionResponse;
184
185    #[test]
186    fn test_mistral_transcription_response_deserialize() {
187        let json = r#" {
188          "model": "voxtral-mini-latest",
189          "text": "The sun was setting slowly, casting long shadows across the empty field.",
190          "language": null,
191          "segments": [
192            {
193              "text": "The sun was setting slowly, casting long shadows across the empty field.",
194              "start": 0.2,
195              "end": 4.6,
196              "speaker_id": "speaker_1",
197              "type": "transcription_segment"
198            }
199          ],
200          "usage": {
201            "prompt_audio_seconds": 5,
202            "prompt_tokens": 5,
203            "total_tokens": 404,
204            "completion_tokens": 24,
205            "prompt_tokens_details": {
206              "cached_tokens": 368
207            }
208          },
209          "finish_reason": null
210            }"#;
211
212        let response: MistralTranscriptionResponse =
213            serde_json::from_str(json).expect("should deserialize");
214
215        assert_eq!(response.language, None);
216        assert_eq!(response.model, VOXTRAL_MINI);
217        assert_eq!(response.segments.len(), 1);
218
219        let seg0 = &response.segments[0];
220        assert_eq!(seg0.start, 0.2);
221        assert_eq!(seg0.end, 4.6);
222        assert_eq!(seg0.score, None);
223        assert_eq!(seg0.speaker_id, Some("speaker_1".to_string()));
224        assert_eq!(seg0.segment_type, "transcription_segment");
225
226        assert_eq!(response.usage.prompt_audio_seconds, Some(5));
227        assert_eq!(response.usage.prompt_tokens, 5);
228        assert_eq!(response.usage.total_tokens, 404);
229        let usage_token_details = response.usage.prompt_tokens_details.unwrap();
230        let cached_token = usage_token_details.get("cached_tokens").unwrap();
231
232        assert_eq!(cached_token.to_string().parse::<i32>().unwrap(), 368);
233    }
234
235    #[test]
236    fn test_response_conversion() {
237        let mistral_response = MistralTranscriptionResponse {
238            language: Some("en".to_string()),
239            model: VOXTRAL_MINI.to_string(),
240            segments: vec![SegmentChunk {
241                start: 0.0,
242                end: 1.0,
243                text: "Lorem Ipsum is simply dummy text of the printing and typesetting industry."
244                    .into(),
245                score: None,
246                speaker_id: None,
247                segment_type: "speech".to_string(),
248            }],
249            text: "Lorem Ipsum is simply dummy text of the printing and typesetting industry."
250                .to_string(),
251            usage: TranscriptionUsage {
252                prompt_audio_seconds: Some(1),
253                prompt_tokens: 10,
254                total_tokens: 20,
255                completion_tokens: 10,
256                prompt_tokens_details: None,
257            },
258        };
259
260        let response: TranscriptionResponse<MistralTranscriptionResponse> = mistral_response
261            .try_into()
262            .expect("conversion should succeed");
263
264        assert_eq!(
265            response.text,
266            "Lorem Ipsum is simply dummy text of the printing and typesetting industry."
267        );
268        assert_eq!(response.response.model, VOXTRAL_MINI);
269        assert_eq!(response.response.language, Some("en".to_string()));
270    }
271
272    #[tokio::test]
273    async fn transcription_non_success_preserves_status_and_body() {
274        use crate::client::transcription::TranscriptionClient;
275        use crate::test_utils::RecordingHttpClient;
276        use crate::transcription::{TranscriptionError, TranscriptionModel as _};
277
278        let body = r#"{"error":{"message":"boom"}}"#;
279        let http_client =
280            RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
281        let client = Client::builder()
282            .api_key("test-key")
283            .http_client(http_client)
284            .build()
285            .expect("build client");
286        let model = client.transcription_model(VOXTRAL_MINI);
287
288        let error = match model
289            .transcription_request()
290            .data(vec![0u8; 16])
291            .send()
292            .await
293        {
294            Err(error) => error,
295            Ok(_) => panic!("transcription should fail with non-success status"),
296        };
297
298        assert!(matches!(error, TranscriptionError::HttpError(_)));
299        assert_eq!(
300            error.provider_response_status(),
301            Some(http::StatusCode::SERVICE_UNAVAILABLE)
302        );
303        assert_eq!(error.provider_response_body(), Some(body));
304    }
305}