Skip to main content

ferrin_core/
transcription.rs

1//! Transcription: [`transcribe`] for complete audio (bytes or URL) and
2//! [`stream_transcribe`] for live audio.
3//!
4//! Design: `docs/01-architecture/11-other-modalities.md` ยง4.
5
6use 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
45/// Media type used when detection fails.
46const DEFAULT_AUDIO_MEDIA_TYPE: &str = "audio/wav";
47
48/// Audio to transcribe.
49#[derive(Debug, Clone, PartialEq, Eq)]
50pub enum AudioInput {
51    /// Audio bytes.
52    Bytes(Bytes),
53    /// A URL fetched with the download function.
54    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/// Result of [`transcribe`] and of [`StreamTranscribeResult::consume`].
76#[derive(Debug, Clone, PartialEq)]
77pub struct TranscribeResult {
78    /// The transcript.
79    pub text: String,
80    /// Timed segments.
81    pub segments: Vec<TranscriptionSegment>,
82    /// Detected language.
83    pub language: Option<String>,
84    /// Audio duration in seconds.
85    pub duration_in_seconds: Option<f64>,
86    /// Adapter warnings.
87    pub warnings: Vec<Warning>,
88    /// Request metadata.
89    pub request: RequestMetadata,
90    /// Response metadata of the calls made.
91    pub responses: Vec<ResponseMetadata>,
92    /// Provider-specific metadata.
93    pub provider_metadata: Option<ProviderMetadata>,
94}
95
96/// Transcribes complete audio.
97#[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
111/// Builder returned by [`transcribe`]; `.await` runs the call.
112pub 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    /// Overrides the detected media type of the audio.
134    #[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    /// Sets the function used to fetch [`AudioInput::Url`].
141    #[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
159/// Fetches audio given as a URL.
160async fn fetch_audio(
161    input: AudioInput,
162    download: Option<Arc<dyn DownloadFn>>,
163    cancellation: &tokio_util::sync::CancellationToken,
164) -> Result<(Bytes, Option<MediaType>), Error> {
165    match input {
166        AudioInput::Bytes(bytes) => Ok((bytes, None)),
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, file.media_type)),
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, downloaded_media_type) =
201                fetch_audio(builder.audio, builder.download, &token).await?;
202            let media_type = builder
203                .media_type
204                .or(downloaded_media_type)
205                .or_else(|| detect_media_type_for(&audio, "audio"))
206                .unwrap_or_else(|| MediaType::new(DEFAULT_AUDIO_MEDIA_TYPE));
207            let headers = base.request_headers();
208            let result = retry(&base.retry_policy, &token, |_| {
209                let options = TranscriptionOptions {
210                    audio: audio.clone(),
211                    media_type: media_type.clone(),
212                    provider_options: base.provider_options.clone(),
213                    headers: headers.clone(),
214                    cancellation: token.child_token(),
215                };
216                let model = &model;
217                async move { model.do_generate(options).await.map_err(Error::from) }
218            })
219            .await?;
220            spans::log_warnings(&result.warnings, &identity);
221            if result.text.is_empty() {
222                return Err(Error::NoTranscriptGenerated {
223                    responses: vec![result.response],
224                });
225            }
226            Ok(TranscribeResult {
227                text: result.text,
228                segments: result.segments,
229                language: result.language,
230                duration_in_seconds: result.duration_in_seconds,
231                warnings: result.warnings,
232                request: result.request,
233                responses: vec![result.response],
234                provider_metadata: result.provider_metadata,
235            })
236        }
237        .instrument(span)
238    })
239    .await
240}
241
242/// Transcribes a live audio stream. The model must support streaming
243/// transcription (`supports_stream`).
244#[must_use]
245pub fn stream_transcribe(
246    model: impl Into<TranscriptionModelRef>,
247    audio: impl Stream<Item = Bytes> + Send + 'static,
248    input_audio_format: AudioFormat,
249) -> StreamTranscribe {
250    StreamTranscribe {
251        model: model.into(),
252        audio: Box::pin(audio),
253        input_audio_format,
254        include_raw_chunks: false,
255        base: ModalityOptions::default(),
256    }
257}
258
259/// Builder returned by [`stream_transcribe`]; `.await` opens the stream.
260pub struct StreamTranscribe {
261    model: TranscriptionModelRef,
262    audio: BoxStream<'static, Bytes>,
263    input_audio_format: AudioFormat,
264    include_raw_chunks: bool,
265    base: ModalityOptions,
266}
267
268impl fmt::Debug for StreamTranscribe {
269    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
270        f.debug_struct("StreamTranscribe")
271            .field("model", &self.model)
272            .field("input_audio_format", &self.input_audio_format)
273            .field("include_raw_chunks", &self.include_raw_chunks)
274            .field("base", &self.base)
275            .finish_non_exhaustive()
276    }
277}
278
279impl StreamTranscribe {
280    /// Forwards raw provider chunks as [`TranscriptionStreamPart::Raw`].
281    #[must_use]
282    pub fn include_raw_chunks(mut self) -> Self {
283        self.include_raw_chunks = true;
284        self
285    }
286}
287
288impl_modality_builder!(@no_retry StreamTranscribe);
289
290impl IntoFuture for StreamTranscribe {
291    type Output = Result<StreamTranscribeResult, Error>;
292    type IntoFuture = BoxFuture<'static, Self::Output>;
293
294    fn into_future(self) -> Self::IntoFuture {
295        Box::pin(async move {
296            let model = resolve_model(&self.model, ProviderRegistry::transcription_model)?;
297            let identity = ModelIdentity::new(model.provider().clone(), model.model_id().clone());
298            if !model.supports_stream() {
299                return Err(Error::from(ProviderError::unsupported(format!(
300                    "streaming transcription (model `{}` of provider `{}`)",
301                    identity.model_id, identity.provider
302                ))));
303            }
304            let deadline = StreamDeadline::new(&self.base.cancellation, self.base.timeout);
305            let result = deadline
306                .run(async {
307                    model
308                        .do_stream(TranscriptionStreamOptions {
309                            audio: self.audio,
310                            input_audio_format: self.input_audio_format,
311                            provider_options: self.base.provider_options.clone(),
312                            headers: self.base.request_headers(),
313                            include_raw_chunks: self.include_raw_chunks,
314                            cancellation: deadline.cancellation.clone(),
315                        })
316                        .await
317                        .map_err(Error::from)
318                })
319                .await?;
320            let log_identity = identity.clone();
321            let parts = result.stream.inspect(move |part| {
322                if let TranscriptionStreamPart::StreamStart { warnings } = part {
323                    spans::log_warnings(warnings, &log_identity);
324                }
325            });
326            Ok(StreamTranscribeResult {
327                request: result.request,
328                response: result.response,
329                parts: deadline.wrap(
330                    Box::pin(parts),
331                    |error| TranscriptionStreamPart::Error { error },
332                    |part| {
333                        matches!(
334                            part,
335                            TranscriptionStreamPart::Finish { .. }
336                                | TranscriptionStreamPart::Error { .. }
337                        )
338                    },
339                ),
340            })
341        })
342    }
343}
344
345/// Result of [`stream_transcribe`]: the part stream plus request and
346/// response metadata known when the stream opened.
347pub struct StreamTranscribeResult {
348    /// Request metadata.
349    pub request: RequestMetadata,
350    /// Response metadata known at stream start.
351    pub response: ResponseMetadata,
352    parts: BoxStream<'static, TranscriptionStreamPart>,
353}
354
355impl fmt::Debug for StreamTranscribeResult {
356    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
357        f.debug_struct("StreamTranscribeResult")
358            .field("request", &self.request)
359            .field("response", &self.response)
360            .finish_non_exhaustive()
361    }
362}
363
364impl StreamTranscribeResult {
365    /// The stream of provider parts.
366    pub fn parts(&mut self) -> &mut BoxStream<'static, TranscriptionStreamPart> {
367        &mut self.parts
368    }
369
370    /// Takes the stream of provider parts.
371    #[must_use]
372    pub fn into_parts(self) -> BoxStream<'static, TranscriptionStreamPart> {
373        self.parts
374    }
375
376    /// Transcript deltas; error parts end the stream with [`Error::Stream`].
377    pub fn text_stream(self) -> impl Stream<Item = Result<String, Error>> + Send {
378        self.parts
379            .map(|part| match part {
380                TranscriptionStreamPart::TranscriptDelta { delta, .. } => Some(Ok(delta)),
381                TranscriptionStreamPart::Error { error } => Some(Err(Error::stream(error))),
382                _ => None,
383            })
384            .filter_map(std::future::ready)
385    }
386
387    /// Drains the stream and returns the final transcript.
388    ///
389    /// # Errors
390    ///
391    /// Returns [`Error::Stream`] when the provider emitted an error part
392    /// and [`Error::NoTranscriptGenerated`] when the stream ended without
393    /// a non-empty transcript.
394    pub async fn consume(mut self) -> Result<TranscribeResult, Error> {
395        let mut warnings: Vec<Warning> = Vec::new();
396        let mut response = self.response.clone();
397        while let Some(part) = self.parts.next().await {
398            match part {
399                TranscriptionStreamPart::StreamStart { warnings: started } => {
400                    warnings.extend(started)
401                }
402                TranscriptionStreamPart::ResponseMetadata {
403                    timestamp,
404                    model_id,
405                    headers,
406                    body,
407                } => {
408                    if timestamp.is_some() {
409                        response.timestamp = timestamp;
410                    }
411                    if model_id.is_some() {
412                        response.model_id = model_id;
413                    }
414                    if headers.is_some() {
415                        response.headers = headers;
416                    }
417                    if body.is_some() {
418                        response.body = body;
419                    }
420                }
421                TranscriptionStreamPart::Finish {
422                    text,
423                    segments,
424                    language,
425                    duration_in_seconds,
426                    provider_metadata,
427                } => {
428                    if text.is_empty() {
429                        return Err(Error::NoTranscriptGenerated {
430                            responses: vec![response],
431                        });
432                    }
433                    return Ok(TranscribeResult {
434                        text,
435                        segments,
436                        language,
437                        duration_in_seconds,
438                        warnings,
439                        request: self.request,
440                        responses: vec![response],
441                        provider_metadata,
442                    });
443                }
444                TranscriptionStreamPart::Error { error } => return Err(Error::stream(error)),
445                #[allow(
446                    unreachable_patterns,
447                    reason = "TranscriptionStreamPart is non-exhaustive"
448                )]
449                _ => {}
450            }
451        }
452        Err(Error::NoTranscriptGenerated {
453            responses: vec![response],
454        })
455    }
456}
457
458/// Converts a static list of parts into a stream (used by adapters and
459/// tests to build simple transcription streams).
460#[must_use]
461pub fn transcription_parts_stream(
462    parts: Vec<TranscriptionStreamPart>,
463) -> BoxStream<'static, TranscriptionStreamPart> {
464    Box::pin(stream::iter(parts))
465}