aimux_core/
transcription_model.rs

1//! The `TranscriptionModel` trait — the provider-facing interface for
2//! speech-to-text.
3//!
4//! Aligned with Vercel AI SDK `TranscriptionModelV4`
5//! (`reference/ai/packages/provider/src/transcription-model/v4/`).
6//!
7//! This is the only non-chat model type with an optional streaming method,
8//! [`TranscriptionModel::do_stream`]. The default implementation returns
9//! [`AiMuxError::UnsupportedFunctionality`]; providers override it and users
10//! enter through [`stream_transcribe`].
11
12use std::pin::Pin;
13
14use async_trait::async_trait;
15use futures::Stream;
16use serde::{Deserialize, Serialize};
17use ts_rs::TS;
18
19use crate::error::AiMuxError;
20use crate::shared::{SharedHeaders, SharedProviderMetadata, SharedProviderOptions, Warning};
21use crate::{AbortSignal, retry, timeout};
22
23/// Audio input: raw bytes or a base64-encoded string.
24#[derive(Debug, Clone, Serialize, Deserialize, TS)]
25#[ts(export)]
26pub enum AudioInput {
27    /// Raw binary bytes.
28    Binary(Vec<u8>),
29    /// A base64-encoded string.
30    Base64(String),
31}
32
33/// A chunk of audio in a streaming transcription request.
34#[derive(Debug, Clone, Serialize, Deserialize, TS)]
35#[ts(export)]
36pub enum AudioChunk {
37    /// Raw binary bytes.
38    Binary(Vec<u8>),
39    /// A base64-encoded string.
40    Base64(String),
41}
42
43/// A transcript segment with timing information.
44#[derive(Debug, Clone, Serialize, Deserialize, TS)]
45#[ts(export)]
46pub struct TranscriptionSegment {
47    /// The text content of this segment.
48    pub text: String,
49    /// The start time of this segment in seconds.
50    pub start_second: f64,
51    /// The end time of this segment in seconds.
52    pub end_second: f64,
53}
54
55/// Options passed to [`TranscriptionModel::do_generate`].
56///
57/// Aligned with V4 `TranscriptionModelV4CallOptions`.
58#[derive(Debug, Clone, Serialize, Deserialize, TS)]
59#[ts(export)]
60pub struct TranscriptionCallOptions {
61    /// Audio data to transcribe (raw bytes or a base64-encoded string).
62    pub audio: AudioInput,
63
64    /// The IANA media type of the audio data, e.g. `"audio/mp3"`.
65    pub media_type: String,
66
67    /// Additional provider-specific options, keyed by provider name.
68    pub provider_options: Option<SharedProviderOptions>,
69
70    /// Abort signal for cancelling the operation.
71    #[serde(skip)]
72    #[ts(skip)]
73    pub abort_signal: Option<AbortSignal>,
74
75    /// Per-call retry override. `None` uses the model default.
76    pub max_retries: Option<u32>,
77
78    /// Per-call operation timeout.
79    pub timeout: Option<crate::options::TimeoutConfiguration>,
80
81    /// Additional HTTP headers to send with the request.
82    pub headers: Option<SharedHeaders>,
83}
84
85impl TranscriptionCallOptions {
86    /// Create options with the given audio and media type.
87    pub fn new(audio: AudioInput, media_type: impl Into<String>) -> Self {
88        Self {
89            audio,
90            media_type: media_type.into(),
91            provider_options: None,
92            abort_signal: None,
93            max_retries: None,
94            timeout: None,
95            headers: None,
96        }
97    }
98}
99
100/// The result of [`TranscriptionModel::do_generate`].
101///
102/// Aligned with V4 `TranscriptionModelV4Result`.
103#[derive(Debug, Clone, Serialize, Deserialize, TS)]
104#[ts(export)]
105pub struct TranscriptionResult {
106    /// The complete transcribed text from the audio.
107    pub text: String,
108
109    /// Transcript segments with timing information.
110    pub segments: Vec<TranscriptionSegment>,
111
112    /// The detected language (ISO 639-1 code, e.g. `"en"`), if detected.
113    pub language: Option<String>,
114
115    /// The total duration of the audio in seconds, if determined.
116    pub duration_in_seconds: Option<f64>,
117
118    /// Warnings for the call, e.g. unsupported settings.
119    pub warnings: Vec<Warning>,
120
121    /// Optional request information for telemetry and debugging.
122    pub request: Option<TranscriptionRequest>,
123
124    /// Response information for telemetry and debugging.
125    pub response: TranscriptionResponse,
126
127    /// Additional provider-specific metadata, keyed by provider name.
128    pub provider_metadata: Option<SharedProviderMetadata>,
129}
130
131/// Optional request information for a transcription call.
132#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
133#[ts(export)]
134pub struct TranscriptionRequest {
135    /// Raw request HTTP body that was sent (JSON stringified).
136    pub body: Option<String>,
137}
138
139/// Response information for a transcription call.
140#[derive(Debug, Clone, Default, Serialize, Deserialize, TS)]
141#[ts(export)]
142pub struct TranscriptionResponse {
143    /// Timestamp for the start of the generated response (ISO 8601 string).
144    pub timestamp: Option<String>,
145    /// The ID of the model that was used to generate the response.
146    pub model_id: Option<String>,
147    /// Response headers.
148    pub headers: Option<SharedHeaders>,
149    /// Response body (opaque JSON).
150    pub body: Option<serde_json::Value>,
151}
152
153/// The input audio format for a streaming transcription request.
154#[derive(Debug, Clone, Serialize, Deserialize, TS)]
155#[ts(export)]
156pub struct InputAudioFormat {
157    /// Audio format type, e.g. `"audio/pcm"`, `"audio/pcmu"`, `"audio/pcma"`.
158    pub format_type: String,
159    /// Sample rate in Hz. Only applicable for formats that require a rate.
160    pub rate: Option<u32>,
161}
162
163/// Options passed to [`TranscriptionModel::do_stream`].
164///
165/// Aligned with V4 `TranscriptionModelV4StreamOptions`.
166pub struct TranscriptionStreamOptions {
167    /// Audio chunks to transcribe (raw bytes or base64-encoded strings).
168    pub audio: Pin<Box<dyn Stream<Item = AudioChunk> + Send>>,
169
170    /// The input audio format for the raw audio chunks.
171    pub input_audio_format: InputAudioFormat,
172
173    /// Additional provider-specific options, keyed by provider name.
174    pub provider_options: Option<SharedProviderOptions>,
175
176    /// Abort signal for cancelling the operation.
177    pub abort_signal: Option<AbortSignal>,
178
179    /// Additional HTTP headers for HTTP/WebSocket-based providers.
180    pub headers: Option<SharedHeaders>,
181
182    /// When `true`, providers should include raw provider chunks in the
183    /// stream.
184    pub include_raw_chunks: bool,
185
186    /// Timeout configuration (RFC-0028): `first_chunk_ms` bounds connect +
187    /// session establishment, `chunk_ms` the gap between events, `total_ms`
188    /// the whole stream. `None` = no timeouts (waits are unbounded apart from
189    /// abort).
190    pub timeout: Option<crate::options::TimeoutConfiguration>,
191}
192
193impl std::fmt::Debug for TranscriptionStreamOptions {
194    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
195        f.debug_struct("TranscriptionStreamOptions")
196            .field("input_audio_format", &self.input_audio_format)
197            .field("provider_options", &self.provider_options)
198            .field("abort_signal", &self.abort_signal)
199            .field("headers", &self.headers)
200            .field("include_raw_chunks", &self.include_raw_chunks)
201            .field("timeout", &self.timeout)
202            .field("audio", &"<stream>")
203            .finish()
204    }
205}
206
207/// A single chunk in the stream returned by `do_stream`.
208///
209/// Aligned with V4 `TranscriptionModelV4StreamPart`.
210#[derive(Debug, Serialize, Deserialize, TS)]
211#[ts(export)]
212pub enum TranscriptionStreamPart {
213    /// Stream start event, carrying warnings.
214    StreamStart { warnings: Vec<Warning> },
215    /// Append-only transcript delta.
216    TranscriptDelta {
217        id: Option<String>,
218        delta: String,
219        provider_metadata: Option<SharedProviderMetadata>,
220    },
221    /// Non-final transcript text (may be revised by later parts).
222    TranscriptPartial {
223        id: Option<String>,
224        text: String,
225        start_second: Option<f64>,
226        duration_in_seconds: Option<f64>,
227        channel_index: Option<u32>,
228        provider_metadata: Option<SharedProviderMetadata>,
229    },
230    /// Final transcript text for a provider-defined segment or utterance.
231    TranscriptFinal {
232        id: Option<String>,
233        text: String,
234        start_second: Option<f64>,
235        end_second: Option<f64>,
236        channel_index: Option<u32>,
237        provider_metadata: Option<SharedProviderMetadata>,
238    },
239    /// Response metadata, emitted once available.
240    ResponseMetadata {
241        timestamp: Option<String>,
242        model_id: Option<String>,
243        headers: Option<SharedHeaders>,
244        body: Option<serde_json::Value>,
245    },
246    /// Metadata available after the stream finishes.
247    Finish {
248        text: String,
249        segments: Vec<TranscriptionSegment>,
250        language: Option<String>,
251        duration_in_seconds: Option<f64>,
252        provider_metadata: Option<SharedProviderMetadata>,
253    },
254    /// Raw provider chunk (only when `include_raw_chunks` is `true`).
255    Raw { raw_value: serde_json::Value },
256    /// An error occurred mid-stream.
257    Error { error: serde_json::Value },
258}
259
260/// The result of [`TranscriptionModel::do_stream`].
261///
262/// Aligned with V4 `TranscriptionModelV4StreamResult`.
263pub struct TranscriptionStreamResult {
264    /// The stream of [`TranscriptionStreamPart`] items.
265    pub stream: Pin<Box<dyn Stream<Item = Result<TranscriptionStreamPart, AiMuxError>> + Send>>,
266    /// Optional request information for debugging.
267    pub request: Option<TranscriptionRequest>,
268    /// Optional response information.
269    pub response: Option<TranscriptionResponse>,
270}
271
272impl std::fmt::Debug for TranscriptionStreamResult {
273    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
274        f.debug_struct("TranscriptionStreamResult")
275            .field("request", &self.request)
276            .field("response", &self.response)
277            .field("stream", &"<stream>")
278            .finish()
279    }
280}
281
282/// The unified transcription (STT) model trait (provider-facing).
283///
284/// Aligned with V4 `TranscriptionModelV4`.
285#[async_trait]
286pub trait TranscriptionModel: Send + Sync {
287    /// Specification version (always `"v4"`).
288    fn specification_version(&self) -> &'static str {
289        "v4"
290    }
291
292    /// Provider name, e.g. `"openai"`.
293    fn provider(&self) -> &str;
294
295    /// Provider-specific model ID, e.g. `"whisper-1"`.
296    fn model_id(&self) -> &str;
297
298    fn retry_config(&self) -> crate::retry::RetryConfig {
299        crate::retry::RetryConfig::default()
300    }
301
302    /// Generate a transcript.
303    ///
304    /// Naming: the `do_` prefix prevents accidental direct usage by users.
305    async fn do_generate(
306        &self,
307        options: &TranscriptionCallOptions,
308    ) -> Result<TranscriptionResult, AiMuxError>;
309
310    /// Stream a transcript for live audio.
311    ///
312    /// Default implementation returns [`AiMuxError::UnsupportedFunctionality`]; providers
313    /// override it as needed. This mirrors the optional `doStream?` in the TS
314    /// spec.
315    async fn do_stream(
316        &self,
317        _options: TranscriptionStreamOptions,
318    ) -> Result<TranscriptionStreamResult, AiMuxError> {
319        Err(AiMuxError::UnsupportedFunctionality(format!(
320            "transcription streaming is not supported by provider `{}`",
321            self.provider()
322        )))
323    }
324}
325
326/// User-facing non-streaming transcription with Core-owned retry and timeout.
327///
328/// # Errors
329///
330/// Returns the provider failure, retry exhaustion, timeout, or caller abort.
331pub async fn transcribe(
332    model: &dyn TranscriptionModel,
333    options: TranscriptionCallOptions,
334) -> Result<TranscriptionResult, AiMuxError> {
335    let timeout = timeout::OperationTimeout::new(options.timeout.unwrap_or_default())?;
336    let abort_signal = options.abort_signal.clone();
337    let retries = retry::prepare_retries(
338        options.max_retries,
339        model.retry_config(),
340        abort_signal.clone(),
341    );
342    timeout::run(
343        retries.retry(|| model.do_generate(&options)),
344        abort_signal.as_ref(),
345        timeout,
346    )
347    .await
348}
349
350/// Start user-facing live transcription.
351///
352/// The live audio stream cannot be replayed, so session setup is attempted
353/// once. The provider session applies the supplied abort signal and streaming
354/// timeouts throughout connect, send, and receive.
355///
356/// # Errors
357///
358/// Returns the provider's session-establishment error.
359pub async fn stream_transcribe(
360    model: &dyn TranscriptionModel,
361    options: TranscriptionStreamOptions,
362) -> Result<TranscriptionStreamResult, AiMuxError> {
363    model.do_stream(options).await
364}
aimux_core/transcription_model.rs

aimux_core/
transcription_model.rs
