1use std::fmt;
7use std::future::IntoFuture;
8use std::sync::Arc;
9
10use bytes::Bytes;
11use ferrin_provider_util::media_type::detect_media_type_for;
12use ferrin_spec::AudioFormat;
13use ferrin_spec::BoxFuture;
14use ferrin_spec::BoxStream;
15use ferrin_spec::MediaType;
16use ferrin_spec::ProviderMetadata;
17use ferrin_spec::RequestMetadata;
18use ferrin_spec::ResponseMetadata;
19use ferrin_spec::TranscriptionModelRef;
20use ferrin_spec::Warning;
21use ferrin_spec::error::ProviderError;
22use ferrin_spec::transcription_model::TranscriptionOptions;
23pub use ferrin_spec::transcription_model::TranscriptionSegment;
24use ferrin_spec::transcription_model::TranscriptionStreamOptions;
25pub use ferrin_spec::transcription_model::TranscriptionStreamPart;
26use futures_core::Stream;
27use futures_util::StreamExt;
28use futures_util::stream;
29use tracing::Instrument;
30use url::Url;
31
32use crate::error::Error;
33use crate::modality::ModalityOptions;
34use crate::modality::impl_modality_builder;
35use crate::modality_stream::StreamDeadline;
36use crate::prompt::DefaultDownloader;
37use crate::prompt::DownloadFn;
38use crate::prompt::DownloadRequest;
39use crate::registry::ProviderRegistry;
40use crate::registry::default::resolve_model;
41use crate::retry::retry;
42use crate::telemetry::ModelIdentity;
43use crate::telemetry::spans;
44
45const DEFAULT_AUDIO_MEDIA_TYPE: &str = "audio/wav";
47
48#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum AudioInput {
51 Bytes(Bytes),
53 Url(Url),
55}
56
57impl From<Bytes> for AudioInput {
58 fn from(bytes: Bytes) -> Self {
59 Self::Bytes(bytes)
60 }
61}
62
63impl From<Vec<u8>> for AudioInput {
64 fn from(bytes: Vec<u8>) -> Self {
65 Self::Bytes(Bytes::from(bytes))
66 }
67}
68
69impl From<Url> for AudioInput {
70 fn from(url: Url) -> Self {
71 Self::Url(url)
72 }
73}
74
75#[derive(Debug, Clone, PartialEq)]
77pub struct TranscribeResult {
78 pub text: String,
80 pub segments: Vec<TranscriptionSegment>,
82 pub language: Option<String>,
84 pub duration_in_seconds: Option<f64>,
86 pub warnings: Vec<Warning>,
88 pub request: RequestMetadata,
90 pub responses: Vec<ResponseMetadata>,
92 pub provider_metadata: Option<ProviderMetadata>,
94}
95
96#[must_use]
98pub fn transcribe(
99 model: impl Into<TranscriptionModelRef>,
100 audio: impl Into<AudioInput>,
101) -> Transcribe {
102 Transcribe {
103 model: model.into(),
104 audio: audio.into(),
105 media_type: None,
106 download: None,
107 base: ModalityOptions::default(),
108 }
109}
110
111pub struct Transcribe {
113 model: TranscriptionModelRef,
114 audio: AudioInput,
115 media_type: Option<MediaType>,
116 download: Option<Arc<dyn DownloadFn>>,
117 base: ModalityOptions,
118}
119
120impl fmt::Debug for Transcribe {
121 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
122 f.debug_struct("Transcribe")
123 .field("model", &self.model)
124 .field("audio", &self.audio)
125 .field("media_type", &self.media_type)
126 .field("has_download", &self.download.is_some())
127 .field("base", &self.base)
128 .finish()
129 }
130}
131
132impl Transcribe {
133 #[must_use]
135 pub fn media_type(mut self, media_type: impl Into<MediaType>) -> Self {
136 self.media_type = Some(media_type.into());
137 self
138 }
139
140 #[must_use]
142 pub fn download(mut self, download: Arc<dyn DownloadFn>) -> Self {
143 self.download = Some(download);
144 self
145 }
146}
147
148impl_modality_builder!(Transcribe);
149
150impl IntoFuture for Transcribe {
151 type Output = Result<TranscribeResult, Error>;
152 type IntoFuture = BoxFuture<'static, Self::Output>;
153
154 fn into_future(self) -> Self::IntoFuture {
155 Box::pin(run(self))
156 }
157}
158
159async fn fetch_audio(
161 input: AudioInput,
162 download: Option<Arc<dyn DownloadFn>>,
163 cancellation: &tokio_util::sync::CancellationToken,
164) -> Result<Bytes, Error> {
165 match input {
166 AudioInput::Bytes(bytes) => Ok(bytes),
167 AudioInput::Url(url) => {
168 let downloader: Arc<dyn DownloadFn> = match download {
169 Some(download) => download,
170 None => Arc::new(DefaultDownloader::try_default()?),
171 };
172 let mut downloaded = downloader
173 .download(
174 vec![DownloadRequest {
175 url: url.clone(),
176 is_url_supported_by_model: false,
177 }],
178 cancellation.clone(),
179 )
180 .await?;
181 match downloaded.pop().flatten() {
182 Some(file) => Ok(file.data),
183 None => Err(Error::download(
184 url,
185 None,
186 Some("the download function returned no data".into()),
187 )),
188 }
189 }
190 }
191}
192
193async fn run(builder: Transcribe) -> Result<TranscribeResult, Error> {
194 let model = resolve_model(&builder.model, ProviderRegistry::transcription_model)?;
195 let identity = ModelIdentity::new(model.provider().clone(), model.model_id().clone());
196 let span = spans::modality_span("transcription", &identity);
197 let base = builder.base.clone();
198 base.run(|base, token| {
199 async move {
200 let audio = fetch_audio(builder.audio, builder.download, &token).await?;
201 let media_type = builder
202 .media_type
203 .or_else(|| detect_media_type_for(&audio, "audio"))
204 .unwrap_or_else(|| MediaType::new(DEFAULT_AUDIO_MEDIA_TYPE));
205 let headers = base.request_headers();
206 let result = retry(&base.retry_policy, &token, |_| {
207 let options = TranscriptionOptions {
208 audio: audio.clone(),
209 media_type: media_type.clone(),
210 provider_options: base.provider_options.clone(),
211 headers: headers.clone(),
212 cancellation: token.child_token(),
213 };
214 let model = &model;
215 async move { model.do_generate(options).await.map_err(Error::from) }
216 })
217 .await?;
218 spans::log_warnings(&result.warnings, &identity);
219 if result.text.is_empty() {
220 return Err(Error::NoTranscriptGenerated {
221 responses: vec![result.response],
222 });
223 }
224 Ok(TranscribeResult {
225 text: result.text,
226 segments: result.segments,
227 language: result.language,
228 duration_in_seconds: result.duration_in_seconds,
229 warnings: result.warnings,
230 request: result.request,
231 responses: vec![result.response],
232 provider_metadata: result.provider_metadata,
233 })
234 }
235 .instrument(span)
236 })
237 .await
238}
239
240#[must_use]
243pub fn stream_transcribe(
244 model: impl Into<TranscriptionModelRef>,
245 audio: impl Stream<Item = Bytes> + Send + 'static,
246 input_audio_format: AudioFormat,
247) -> StreamTranscribe {
248 StreamTranscribe {
249 model: model.into(),
250 audio: Box::pin(audio),
251 input_audio_format,
252 include_raw_chunks: false,
253 base: ModalityOptions::default(),
254 }
255}
256
257pub struct StreamTranscribe {
259 model: TranscriptionModelRef,
260 audio: BoxStream<'static, Bytes>,
261 input_audio_format: AudioFormat,
262 include_raw_chunks: bool,
263 base: ModalityOptions,
264}
265
266impl fmt::Debug for StreamTranscribe {
267 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
268 f.debug_struct("StreamTranscribe")
269 .field("model", &self.model)
270 .field("input_audio_format", &self.input_audio_format)
271 .field("include_raw_chunks", &self.include_raw_chunks)
272 .field("base", &self.base)
273 .finish_non_exhaustive()
274 }
275}
276
277impl StreamTranscribe {
278 #[must_use]
280 pub fn include_raw_chunks(mut self) -> Self {
281 self.include_raw_chunks = true;
282 self
283 }
284}
285
286impl_modality_builder!(@no_retry StreamTranscribe);
287
288impl IntoFuture for StreamTranscribe {
289 type Output = Result<StreamTranscribeResult, Error>;
290 type IntoFuture = BoxFuture<'static, Self::Output>;
291
292 fn into_future(self) -> Self::IntoFuture {
293 Box::pin(async move {
294 let model = resolve_model(&self.model, ProviderRegistry::transcription_model)?;
295 let identity = ModelIdentity::new(model.provider().clone(), model.model_id().clone());
296 if !model.supports_stream() {
297 return Err(Error::from(ProviderError::unsupported(format!(
298 "streaming transcription (model `{}` of provider `{}`)",
299 identity.model_id, identity.provider
300 ))));
301 }
302 let started_at = chrono::Utc::now();
303 let deadline = StreamDeadline::new(&self.base.cancellation, self.base.timeout);
304 let result = deadline
305 .run(async {
306 model
307 .do_stream(TranscriptionStreamOptions {
308 audio: self.audio,
309 input_audio_format: self.input_audio_format,
310 provider_options: self.base.provider_options.clone(),
311 headers: self.base.request_headers(),
312 include_raw_chunks: self.include_raw_chunks,
313 cancellation: deadline.cancellation.clone(),
314 })
315 .await
316 .map_err(Error::from)
317 })
318 .await?;
319 let mut response = result.response;
320 response.timestamp.get_or_insert(started_at);
321 response
322 .model_id
323 .get_or_insert_with(|| identity.model_id.clone());
324 let log_identity = identity.clone();
325 let parts = result.stream.inspect(move |part| {
326 if let TranscriptionStreamPart::StreamStart { warnings } = part {
327 spans::log_warnings(warnings, &log_identity);
328 }
329 });
330 Ok(StreamTranscribeResult {
331 request: result.request,
332 response,
333 parts: deadline.wrap(
334 Box::pin(parts),
335 |error| TranscriptionStreamPart::Error { error },
336 |part| {
337 matches!(
338 part,
339 TranscriptionStreamPart::Finish { .. }
340 | TranscriptionStreamPart::Error { .. }
341 )
342 },
343 ),
344 })
345 })
346 }
347}
348
349pub struct StreamTranscribeResult {
352 pub request: RequestMetadata,
354 pub response: ResponseMetadata,
356 parts: BoxStream<'static, TranscriptionStreamPart>,
357}
358
359impl fmt::Debug for StreamTranscribeResult {
360 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
361 f.debug_struct("StreamTranscribeResult")
362 .field("request", &self.request)
363 .field("response", &self.response)
364 .finish_non_exhaustive()
365 }
366}
367
368impl StreamTranscribeResult {
369 pub fn parts(&mut self) -> &mut BoxStream<'static, TranscriptionStreamPart> {
371 &mut self.parts
372 }
373
374 #[must_use]
376 pub fn into_parts(self) -> BoxStream<'static, TranscriptionStreamPart> {
377 self.parts
378 }
379
380 pub fn text_stream(self) -> impl Stream<Item = Result<String, Error>> + Send {
382 self.parts
383 .map(|part| match part {
384 TranscriptionStreamPart::TranscriptDelta { delta, .. } => Some(Ok(delta)),
385 TranscriptionStreamPart::Error { error } => Some(Err(Error::stream(error))),
386 _ => None,
387 })
388 .filter_map(std::future::ready)
389 }
390
391 pub async fn consume(mut self) -> Result<TranscribeResult, Error> {
399 let mut warnings: Vec<Warning> = Vec::new();
400 let mut response = self.response.clone();
401 while let Some(part) = self.parts.next().await {
402 match part {
403 TranscriptionStreamPart::StreamStart { warnings: started } => {
404 warnings.extend(started)
405 }
406 TranscriptionStreamPart::ResponseMetadata {
407 timestamp,
408 model_id,
409 headers,
410 body,
411 } => {
412 if timestamp.is_some() {
413 response.timestamp = timestamp;
414 }
415 if model_id.is_some() {
416 response.model_id = model_id;
417 }
418 if headers.is_some() {
419 response.headers = headers;
420 }
421 if body.is_some() {
422 response.body = body;
423 }
424 }
425 TranscriptionStreamPart::Finish {
426 text,
427 segments,
428 language,
429 duration_in_seconds,
430 provider_metadata,
431 } => {
432 if text.is_empty() {
433 return Err(Error::NoTranscriptGenerated {
434 responses: vec![response],
435 });
436 }
437 return Ok(TranscribeResult {
438 text,
439 segments,
440 language,
441 duration_in_seconds,
442 warnings,
443 request: self.request,
444 responses: vec![response],
445 provider_metadata,
446 });
447 }
448 TranscriptionStreamPart::Error { error } => return Err(Error::stream(error)),
449 #[allow(
450 unreachable_patterns,
451 reason = "TranscriptionStreamPart is non-exhaustive"
452 )]
453 _ => {}
454 }
455 }
456 Err(Error::NoTranscriptGenerated {
457 responses: vec![response],
458 })
459 }
460}
461
462#[must_use]
465pub fn transcription_parts_stream(
466 parts: Vec<TranscriptionStreamPart>,
467) -> BoxStream<'static, TranscriptionStreamPart> {
468 Box::pin(stream::iter(parts))
469}