aurum_core/remote/
stt_chunk.rs

1//! Time-based remote STT chunk-and-stitch (JOE-2212).
2//!
3//! Long lectures can exceed [`TranscriptLimits::max_segment_chars`] when a vendor
4//! returns a single continuous segment, or truncate on full-file remote paths.
5//! Client-side audio chunking (~210s windows, dual-ref eval band) keeps each
6//! request small and stitches text/segments with time offsets.
7//!
8//! Local whisper is unchanged (handles full-file natively).
9
10use crate::audio::AudioInput;
11use crate::error::{ProviderError, Result};
12use crate::providers::{Segment, TranscriptionOptions, TranscriptionResult};
13use crate::remote::limits::{validate_segments, validate_text_bounds, TranscriptLimits};
14use std::future::Future;
15use std::path::PathBuf;
16use std::sync::Arc;
17
18/// Default remote STT window length (seconds). Matches Plaud dual-ref chunk recipes.
19pub const DEFAULT_REMOTE_STT_CHUNK_SECS: f64 = 210.0;
20
21/// Minimum duration that triggers chunking (must be > 0 and finite).
22pub fn needs_time_chunk(duration_secs: f64, chunk_secs: f64) -> bool {
23    duration_secs.is_finite()
24        && chunk_secs.is_finite()
25        && chunk_secs > 0.0
26        && duration_secs > chunk_secs
27}
28
29/// Inclusive start / exclusive end sample ranges with start offset in seconds.
30#[derive(Debug, Clone, Copy, PartialEq)]
31pub struct ChunkWindow {
32    pub start_sample: usize,
33    pub end_sample: usize,
34    pub offset_secs: f64,
35}
36
37/// Plan non-overlapping PCM windows covering `[0, total_samples)`.
38pub fn plan_chunk_windows(
39    total_samples: usize,
40    sample_rate: u32,
41    chunk_secs: f64,
42) -> Vec<ChunkWindow> {
43    if total_samples == 0 || sample_rate == 0 || !chunk_secs.is_finite() || chunk_secs <= 0.0 {
44        return vec![ChunkWindow {
45            start_sample: 0,
46            end_sample: total_samples,
47            offset_secs: 0.0,
48        }];
49    }
50    let chunk_samples = ((chunk_secs * f64::from(sample_rate)).round() as usize).max(1);
51    if total_samples <= chunk_samples {
52        return vec![ChunkWindow {
53            start_sample: 0,
54            end_sample: total_samples,
55            offset_secs: 0.0,
56        }];
57    }
58    let mut out = Vec::new();
59    let mut start = 0usize;
60    while start < total_samples {
61        let end = (start + chunk_samples).min(total_samples);
62        out.push(ChunkWindow {
63            start_sample: start,
64            end_sample: end,
65            offset_secs: start as f64 / f64::from(sample_rate),
66        });
67        if end == total_samples {
68            break;
69        }
70        start = end;
71    }
72    out
73}
74
75/// Slice mono PCM into a new [`AudioInput`] for one window.
76pub fn slice_audio_window(input: &AudioInput, window: ChunkWindow) -> Result<AudioInput> {
77    let samples = input.samples();
78    let start = window.start_sample.min(samples.len());
79    let end = window.end_sample.min(samples.len()).max(start);
80    if start == end {
81        return Err(ProviderError::Other {
82            message: "remote STT chunk window is empty".into(),
83        }
84        .into());
85    }
86    let slice: Arc<[f32]> = Arc::from(samples[start..end].to_vec());
87    let sr = input.sample_rate();
88    let duration = if sr > 0 {
89        (end - start) as f64 / f64::from(sr)
90    } else {
91        0.0
92    };
93    Ok(AudioInput::from_parts_unchecked(
94        PathBuf::from(format!(
95            "pcm://remote-stt-chunk/{:.3}-{:.3}",
96            window.offset_secs,
97            window.offset_secs + duration
98        )),
99        slice,
100        sr,
101        duration,
102    ))
103}
104
105/// Soft-split an overlong single segment into adjacent pieces under `max_chars`.
106///
107/// Used when a vendor returns one continuous hyp for a chunk that still exceeds
108/// [`TranscriptLimits::max_segment_chars`] (rare for ~210s speech, but fail-closed
109/// without soft-split would block the whole job).
110pub fn soft_split_text_segments(
111    text: &str,
112    start: f64,
113    end: f64,
114    max_chars: usize,
115) -> Vec<Segment> {
116    let max_chars = max_chars.max(1);
117    let chars: Vec<char> = text.chars().collect();
118    if chars.is_empty() {
119        return vec![Segment::from_parts_unchecked(start, end, String::new())];
120    }
121    if chars.len() <= max_chars {
122        return vec![Segment::from_parts_unchecked(start, end, text.to_string())];
123    }
124    let n_parts = chars.len().div_ceil(max_chars);
125    let span = (end - start).max(0.0);
126    let mut segs = Vec::with_capacity(n_parts);
127    for i in 0..n_parts {
128        let c0 = i * max_chars;
129        let c1 = ((i + 1) * max_chars).min(chars.len());
130        let piece: String = chars[c0..c1].iter().collect();
131        let t0 = start + span * (c0 as f64 / chars.len() as f64);
132        let t1 = start + span * (c1 as f64 / chars.len() as f64);
133        segs.push(Segment::from_parts_unchecked(t0, t1.max(t0), piece));
134    }
135    segs
136}
137
138/// Offset segments by `offset_secs` and soft-split any that still exceed limits.
139pub fn normalize_chunk_segments(
140    segments: &[Segment],
141    offset_secs: f64,
142    limits: TranscriptLimits,
143) -> Vec<Segment> {
144    let mut out = Vec::new();
145    for seg in segments {
146        let start = seg.start() + offset_secs;
147        let end = seg.end() + offset_secs;
148        let text = seg.text();
149        if text.chars().count() > limits.max_segment_chars {
150            out.extend(soft_split_text_segments(
151                text,
152                start,
153                end,
154                limits.max_segment_chars,
155            ));
156        } else {
157            out.push(Segment::from_parts_unchecked(start, end, text.to_string()));
158        }
159    }
160    out
161}
162
163/// Join chunk results into one transcript for the full media duration.
164pub fn stitch_chunk_results(
165    parts: &[(f64, TranscriptionResult)],
166    full_duration_secs: f64,
167    provider: &str,
168    limits: TranscriptLimits,
169) -> Result<TranscriptionResult> {
170    if parts.is_empty() {
171        return Err(ProviderError::TranscriptionFailed {
172            reason: "remote STT chunk-and-stitch produced no chunks".into(),
173        }
174        .into());
175    }
176
177    let mut texts: Vec<String> = Vec::with_capacity(parts.len());
178    let mut segments: Vec<Segment> = Vec::new();
179    let mut timestamps_reliable = true;
180    let mut backend_kind = parts[0].1.backend_kind();
181    let model = parts[0].1.model().to_string();
182    let language = parts[0].1.language().map(|s| s.to_string());
183    let provider_name = parts[0].1.provider().to_string();
184
185    for (offset, r) in parts {
186        let t = r.text().trim();
187        if !t.is_empty() {
188            texts.push(t.to_string());
189        }
190        timestamps_reliable &= r.timestamps_reliable();
191        // Prefer ASR label if any chunk is dedicated ASR.
192        if matches!(r.backend_kind(), crate::providers::BackendKind::Asr) {
193            backend_kind = crate::providers::BackendKind::Asr;
194        }
195        segments.extend(normalize_chunk_segments(r.segments(), *offset, limits));
196    }
197
198    let text = join_transcript_parts(&texts);
199    validate_text_bounds(&text, None, limits, provider)?;
200    validate_segments(&segments, full_duration_secs, limits, provider)?;
201
202    let mut result = TranscriptionResult::openrouter(
203        text,
204        segments,
205        language,
206        model,
207        full_duration_secs,
208        timestamps_reliable,
209    );
210    result.set_provider(if provider_name.is_empty() {
211        provider
212    } else {
213        provider_name.as_str()
214    });
215    result.set_backend_kind(backend_kind);
216    result.set_timestamps_reliable(timestamps_reliable);
217    result.validate_segments()?;
218    Ok(result)
219}
220
221fn join_transcript_parts(parts: &[String]) -> String {
222    let mut out = String::new();
223    for p in parts {
224        let p = p.trim();
225        if p.is_empty() {
226            continue;
227        }
228        if !out.is_empty() && !out.ends_with(|c: char| c.is_whitespace()) {
229            out.push(' ');
230        }
231        out.push_str(p);
232    }
233    out
234}
235
236/// Run `one_shot` once, or time-chunk + stitch when audio is longer than `chunk_secs`.
237///
238/// The callback receives **owned** inputs so futures do not borrow across awaits.
239pub async fn transcribe_maybe_chunked<F, Fut>(
240    input: &AudioInput,
241    options: &TranscriptionOptions,
242    provider: &str,
243    chunk_secs: f64,
244    mut one_shot: F,
245) -> Result<TranscriptionResult>
246where
247    F: FnMut(AudioInput, TranscriptionOptions) -> Fut,
248    Fut: Future<Output = Result<TranscriptionResult>>,
249{
250    let duration = input.duration_secs();
251    if !needs_time_chunk(duration, chunk_secs) {
252        return one_shot(input.clone(), options.clone()).await;
253    }
254
255    let windows = plan_chunk_windows(input.len(), input.sample_rate(), chunk_secs);
256    tracing::info!(
257        provider,
258        duration_secs = duration,
259        chunk_secs,
260        chunks = windows.len(),
261        "remote STT time-chunk-and-stitch (JOE-2212)"
262    );
263
264    let mut parts: Vec<(f64, TranscriptionResult)> = Vec::with_capacity(windows.len());
265    for (i, window) in windows.iter().enumerate() {
266        if let Some(flag) = options.cancel.as_ref() {
267            if flag.is_cancelled() {
268                return Err(ProviderError::Cancelled.into());
269            }
270        }
271        let chunk_input = slice_audio_window(input, *window)?;
272        tracing::debug!(
273            provider,
274            chunk = i + 1,
275            of = windows.len(),
276            offset_secs = window.offset_secs,
277            chunk_duration = chunk_input.duration_secs(),
278            "remote STT chunk"
279        );
280        let result = one_shot(chunk_input, options.clone()).await.map_err(|e| {
281            // Surface which chunk failed for operator diagnostics (no secrets).
282            match e {
283                crate::error::TranscriptionError::Provider(
284                    ProviderError::TranscriptionFailed { reason },
285                ) => ProviderError::TranscriptionFailed {
286                    reason: format!(
287                        "chunk {}/{} (offset {:.1}s): {reason}",
288                        i + 1,
289                        windows.len(),
290                        window.offset_secs
291                    ),
292                }
293                .into(),
294                other => other,
295            }
296        })?;
297        parts.push((window.offset_secs, result));
298    }
299
300    stitch_chunk_results(&parts, duration, provider, TranscriptLimits::default())
301}
302
303/// Effective chunk length: `AURUM_REMOTE_STT_CHUNK_SECS` if set and valid, else default.
304pub fn effective_chunk_secs() -> f64 {
305    match std::env::var("AURUM_REMOTE_STT_CHUNK_SECS") {
306        Ok(s) => {
307            let v: f64 = s.trim().parse().unwrap_or(DEFAULT_REMOTE_STT_CHUNK_SECS);
308            if v.is_finite() && v > 0.0 {
309                v
310            } else {
311                DEFAULT_REMOTE_STT_CHUNK_SECS
312            }
313        }
314        Err(_) => DEFAULT_REMOTE_STT_CHUNK_SECS,
315    }
316}
317
318/// Convenience for tests: 16 kHz silence of `duration_secs`.
319#[cfg(test)]
320pub fn silence_input(duration_secs: f64) -> AudioInput {
321    use crate::audio::WHISPER_SAMPLE_RATE;
322    let n = (duration_secs * f64::from(WHISPER_SAMPLE_RATE)).round() as usize;
323    AudioInput::from_pcm_slice(&vec![0.0f32; n.max(1)], WHISPER_SAMPLE_RATE).unwrap()
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329    use crate::audio::WHISPER_SAMPLE_RATE;
330    use crate::providers::BackendKind;
331
332    #[test]
333    fn needs_chunk_threshold() {
334        assert!(!needs_time_chunk(100.0, 210.0));
335        assert!(!needs_time_chunk(210.0, 210.0));
336        assert!(needs_time_chunk(210.1, 210.0));
337        assert!(!needs_time_chunk(f64::NAN, 210.0));
338    }
339
340    #[test]
341    fn plans_four_windows_for_685s() {
342        let total = (685.0 * f64::from(WHISPER_SAMPLE_RATE)) as usize;
343        let windows = plan_chunk_windows(total, WHISPER_SAMPLE_RATE, 210.0);
344        assert_eq!(windows.len(), 4);
345        assert_eq!(windows[0].offset_secs, 0.0);
346        assert!((windows[1].offset_secs - 210.0).abs() < 0.01);
347        assert_eq!(windows.last().unwrap().end_sample, total);
348        // Coverage without gaps
349        for w in windows.windows(2) {
350            assert_eq!(w[0].end_sample, w[1].start_sample);
351        }
352    }
353
354    #[test]
355    fn single_window_when_short() {
356        let total = WHISPER_SAMPLE_RATE as usize * 30;
357        let windows = plan_chunk_windows(total, WHISPER_SAMPLE_RATE, 210.0);
358        assert_eq!(windows.len(), 1);
359        assert_eq!(windows[0].end_sample, total);
360    }
361
362    #[test]
363    fn soft_split_respects_max_chars() {
364        let text = "a".repeat(20);
365        let segs = soft_split_text_segments(&text, 0.0, 10.0, 8);
366        assert!(segs.len() > 1);
367        assert!(segs.iter().all(|s| s.text().chars().count() <= 8));
368        let joined: String = segs.iter().map(|s| s.text()).collect();
369        assert_eq!(joined, text);
370    }
371
372    #[test]
373    fn stitch_offsets_segments() {
374        let mut a = TranscriptionResult::openrouter(
375            "hello".to_string(),
376            vec![Segment::from_parts_unchecked(0.0, 1.0, "hello".to_string())],
377            None,
378            "m".to_string(),
379            1.0,
380            true,
381        );
382        a.set_provider("openai");
383        a.set_backend_kind(BackendKind::Asr);
384        a.set_timestamps_reliable(true);
385
386        let mut b = TranscriptionResult::openrouter(
387            "world".to_string(),
388            vec![Segment::from_parts_unchecked(0.0, 1.0, "world".to_string())],
389            None,
390            "m".to_string(),
391            1.0,
392            true,
393        );
394        b.set_provider("openai");
395        b.set_backend_kind(BackendKind::Asr);
396        b.set_timestamps_reliable(true);
397
398        let stitched = stitch_chunk_results(
399            &[(0.0, a), (210.0, b)],
400            420.0,
401            "openai",
402            TranscriptLimits::default(),
403        )
404        .unwrap();
405        assert_eq!(stitched.text(), "hello world");
406        assert_eq!(stitched.segments().len(), 2);
407        assert!((stitched.segments()[1].start() - 210.0).abs() < 1e-9);
408        assert!(stitched.timestamps_reliable());
409        assert_eq!(stitched.provider(), "openai");
410    }
411
412    #[tokio::test]
413    async fn maybe_chunked_short_is_single_call() {
414        let input = silence_input(5.0);
415        let opts = TranscriptionOptions {
416            model: "whisper-1".into(),
417            language: "en".into(),
418            timestamps: false,
419            cancel: None,
420        };
421        let mut calls = 0u32;
422        let out = transcribe_maybe_chunked(&input, &opts, "t", 210.0, |inp, _| {
423            calls += 1;
424            let mut r = TranscriptionResult::openrouter(
425                "ok".to_string(),
426                vec![Segment::from_parts_unchecked(
427                    0.0,
428                    inp.duration_secs(),
429                    "ok".to_string(),
430                )],
431                None,
432                "whisper-1".to_string(),
433                inp.duration_secs(),
434                false,
435            );
436            r.set_provider("t");
437            async move { Ok(r) }
438        })
439        .await
440        .unwrap();
441        assert_eq!(calls, 1);
442        assert_eq!(out.text(), "ok");
443    }
444
445    #[tokio::test]
446    async fn maybe_chunked_long_invokes_multiple() {
447        let input = silence_input(500.0);
448        let opts = TranscriptionOptions {
449            model: "whisper-1".into(),
450            language: "en".into(),
451            timestamps: false,
452            cancel: None,
453        };
454        let mut calls = 0u32;
455        let out = transcribe_maybe_chunked(&input, &opts, "t", 210.0, |inp, _| {
456            calls += 1;
457            let label = format!("c{calls}");
458            let mut r = TranscriptionResult::openrouter(
459                label.clone(),
460                vec![Segment::from_parts_unchecked(
461                    0.0,
462                    inp.duration_secs(),
463                    label,
464                )],
465                None,
466                "whisper-1".into(),
467                inp.duration_secs(),
468                false,
469            );
470            r.set_provider("t");
471            async move { Ok(r) }
472        })
473        .await
474        .unwrap();
475        assert_eq!(calls, 3); // 210+210+80
476        assert!(out.text().contains(' '));
477        assert_eq!(out.segments().len(), 3);
478        assert!((out.duration_secs() - 500.0).abs() < 0.02);
479    }
480
481    #[tokio::test]
482    async fn maybe_chunked_honours_cancel() {
483        let input = silence_input(500.0);
484        let flag = crate::cancel::CancelFlag::new();
485        flag.cancel();
486        let opts = TranscriptionOptions {
487            model: "whisper-1".into(),
488            language: "en".into(),
489            timestamps: false,
490            cancel: Some(flag),
491        };
492        let err = transcribe_maybe_chunked(&input, &opts, "t", 210.0, |_inp, _| async {
493            unreachable!("should cancel before first shot")
494        })
495        .await
496        .unwrap_err();
497        assert!(err.to_string().to_lowercase().contains("cancel"));
498    }
499}
aurum_core/remote/stt_chunk.rs

aurum_core/remote/
stt_chunk.rs
