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, 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/// Transcribes a live audio stream. The model must support streaming
241/// transcription (`supports_stream`).
242#[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
257/// Builder returned by [`stream_transcribe`]; `.await` opens the stream.
258pub 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    /// Forwards raw provider chunks as [`TranscriptionStreamPart::Raw`].
279    #[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
349/// Result of [`stream_transcribe`]: the part stream plus request and
350/// response metadata known when the stream opened.
351pub struct StreamTranscribeResult {
352    /// Request metadata.
353    pub request: RequestMetadata,
354    /// Response metadata known at stream start.
355    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    /// The stream of provider parts.
370    pub fn parts(&mut self) -> &mut BoxStream<'static, TranscriptionStreamPart> {
371        &mut self.parts
372    }
373
374    /// Takes the stream of provider parts.
375    #[must_use]
376    pub fn into_parts(self) -> BoxStream<'static, TranscriptionStreamPart> {
377        self.parts
378    }
379
380    /// Transcript deltas; error parts end the stream with [`Error::Stream`].
381    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    /// Drains the stream and returns the final transcript.
392    ///
393    /// # Errors
394    ///
395    /// Returns [`Error::Stream`] when the provider emitted an error part
396    /// and [`Error::NoTranscriptGenerated`] when the stream ended without
397    /// a non-empty transcript.
398    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/// Converts a static list of parts into a stream (used by adapters and
463/// tests to build simple transcription streams).
464#[must_use]
465pub fn transcription_parts_stream(
466    parts: Vec<TranscriptionStreamPart>,
467) -> BoxStream<'static, TranscriptionStreamPart> {
468    Box::pin(stream::iter(parts))
469}