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, 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#[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
259pub 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 #[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
345pub struct StreamTranscribeResult {
348 pub request: RequestMetadata,
350 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 pub fn parts(&mut self) -> &mut BoxStream<'static, TranscriptionStreamPart> {
367 &mut self.parts
368 }
369
370 #[must_use]
372 pub fn into_parts(self) -> BoxStream<'static, TranscriptionStreamPart> {
373 self.parts
374 }
375
376 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 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#[must_use]
461pub fn transcription_parts_stream(
462 parts: Vec<TranscriptionStreamPart>,
463) -> BoxStream<'static, TranscriptionStreamPart> {
464 Box::pin(stream::iter(parts))
465}