rig_core/providers/mistral/
transcription.rs1use 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
11pub const VOXTRAL_MINI: &str = "voxtral-mini-latest";
17pub const VOXTRAL_SMALL: &str = "voxtral-small-latest";
19
20#[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#[derive(Debug, Deserialize)]
52pub struct SegmentChunk {
53 pub start: f32,
55 pub end: f32,
57 pub text: String,
59 pub score: Option<f32>,
60 pub speaker_id: Option<String>,
62 #[serde(rename = "type")]
63 pub segment_type: String,
64}
65
66#[derive(Debug, Deserialize)]
67pub struct MistralTranscriptionResponse {
68 pub language: Option<String>,
70 pub model: String,
72 pub segments: Vec<SegmentChunk>,
74 pub text: String,
76 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}