aurum_core/cleanup/
mod.rs

1//! Post-transcription text cleanup ( "flow").
2//!
3//! Two backends:
4//! - **rules** (default): pure on-device regex/heuristics — no network, no model
5//! - **openrouter**: LLM rewrite via chat completions (optional, explicit)
6//!
7//! Cleanup is **opt-in** at the CLI/library boundary so raw ASR stays available.
8
9pub mod openrouter;
10pub mod rules;
11
12use crate::error::{Result, UserError};
13use crate::providers::TranscriptionResult;
14use crate::runtime::OpContext;
15use async_trait::async_trait;
16use serde::{Deserialize, Serialize};
17
18/// How aggressively / in what shape to clean transcript text.
19#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
20#[serde(rename_all = "snake_case")]
21pub enum CleanupStyle {
22    /// No cleanup (identity).
23    #[default]
24    Raw,
25    /// Collapse whitespace, drop fillers, light punctuation.
26    Clean,
27    /// Clean + bullet list.
28    Bullets,
29    /// Clean + expand contractions / slightly more formal tone (rules) or LLM polish.
30    Professional,
31    /// Extractive (rules) or abstractive (LLM) short summary.
32    Summary,
33}
34
35impl CleanupStyle {
36    pub fn parse(s: &str) -> Result<Self> {
37        match s.trim().to_ascii_lowercase().as_str() {
38            "raw" | "none" | "off" => Ok(Self::Raw),
39            "clean" => Ok(Self::Clean),
40            "bullets" | "bullet" => Ok(Self::Bullets),
41            "professional" | "pro" => Ok(Self::Professional),
42            "summary" | "sum" => Ok(Self::Summary),
43            other => Err(UserError::Other {
44                message: format!(
45 "unknown cleanup style '{other}'\n Hint: use one of: raw, clean, bullets, professional, summary"
46 ),
47            }
48            .into()),
49        }
50    }
51
52    pub fn as_str(self) -> &'static str {
53        match self {
54            Self::Raw => "raw",
55            Self::Clean => "clean",
56            Self::Bullets => "bullets",
57            Self::Professional => "professional",
58            Self::Summary => "summary",
59        }
60    }
61
62    /// Styles that rewrite structure so ASR segment timings no longer match text.
63    pub fn is_structural(self) -> bool {
64        matches!(self, Self::Bullets | Self::Summary)
65    }
66}
67
68/// Where cleanup runs.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
70#[serde(rename_all = "snake_case")]
71pub enum CleanupProviderKind {
72    /// On-device rules only (default).
73    #[default]
74    Rules,
75    /// OpenRouter chat completion.
76    OpenRouter,
77}
78
79impl CleanupProviderKind {
80    pub fn parse(s: &str) -> Result<Self> {
81        match s.trim().to_ascii_lowercase().as_str() {
82            "rules" | "local" | "on-device" | "ondevice" => Ok(Self::Rules),
83            "openrouter" | "remote" | "llm" => Ok(Self::OpenRouter),
84            other => Err(UserError::Other {
85                message: format!(
86                    "unknown cleanup provider '{other}'\n Hint: use one of: rules, openrouter"
87                ),
88            }
89            .into()),
90        }
91    }
92
93    pub fn as_str(self) -> &'static str {
94        match self {
95            Self::Rules => "rules",
96            Self::OpenRouter => "openrouter",
97        }
98    }
99}
100
101/// How to treat ASR segments after cleaning the full transcript text.
102#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
103#[serde(rename_all = "snake_case")]
104pub enum SegmentCleanupPolicy {
105    /// Leave segments unchanged (best for timed SRT with `clean` / `professional`).
106    Keep,
107    /// Drop all segments (default for structural styles: bullets / summary).
108    Clear,
109    /// Run the same cleanup style on each segment's text (keeps timings).
110    PerSegment,
111    /// Pick a sensible default from the style (see [`default_for_style`](Self::default_for_style)).
112    #[default]
113    Auto,
114}
115
116impl SegmentCleanupPolicy {
117    pub fn parse(s: &str) -> Result<Self> {
118        match s.trim().to_ascii_lowercase().as_str() {
119            "keep" | "none" => Ok(Self::Keep),
120            "clear" | "drop" | "empty" => Ok(Self::Clear),
121            "per-segment" | "per_segment" | "each" | "segments" => Ok(Self::PerSegment),
122            "auto" | "default" => Ok(Self::Auto),
123            other => Err(UserError::Other {
124                message: format!(
125 "unknown segment cleanup policy '{other}'\n Hint: use one of: auto, keep, clear, per-segment"
126 ),
127            }
128            .into()),
129        }
130    }
131
132    pub fn as_str(self) -> &'static str {
133        match self {
134            Self::Keep => "keep",
135            Self::Clear => "clear",
136            Self::PerSegment => "per-segment",
137            Self::Auto => "auto",
138        }
139    }
140
141    /// Resolve `Auto` against a cleanup style.
142    pub fn resolve(self, style: CleanupStyle) -> Self {
143        match self {
144            Self::Auto => Self::default_for_style(style),
145            other => other,
146        }
147    }
148
149    /// Structural styles clear segments; light styles keep them.
150    pub fn default_for_style(style: CleanupStyle) -> Self {
151        if style.is_structural() {
152            Self::Clear
153        } else {
154            Self::Keep
155        }
156    }
157}
158
159/// Result of a cleanup pass over bare text.
160#[derive(Debug, Clone, Serialize, Deserialize)]
161pub struct CleanupResult {
162    pub text: String,
163    pub style: CleanupStyle,
164    pub provider: CleanupProviderKind,
165    /// Original text before cleanup (for debugging / JSON).
166    pub original_text: String,
167}
168
169/// Structured report for a full-result cleanup transaction.
170#[derive(Debug, Clone, Serialize, Deserialize)]
171pub struct CleanupReport {
172    pub style: CleanupStyle,
173    pub provider: CleanupProviderKind,
174    pub segment_policy: SegmentCleanupPolicy,
175    /// Fields that changed relative to the input result.
176    pub changed_fields: Vec<String>,
177    pub warnings: Vec<String>,
178    /// Segments dropped because rewritten text was empty (per-segment policy).
179    pub dropped_segments: usize,
180    /// Whether timings were cleared under the applied policy.
181    pub segments_cleared: bool,
182}
183
184/// Hard bounds for per-segment cleanup (failure semantics live here; remote
185/// batching belongs to a follow-up issue).
186pub const MAX_PER_SEGMENT_COUNT: usize = 2_000;
187pub const MAX_PER_SEGMENT_CHARS: usize = 8_000;
188
189/// Pluggable cleanup backend.
190#[async_trait]
191pub trait TextCleanup: Send + Sync {
192    fn name(&self) -> &'static str;
193    fn kind(&self) -> CleanupProviderKind;
194
195    async fn cleanup(&self, text: &str, style: CleanupStyle) -> Result<CleanupResult>;
196
197    /// Clean many segment texts transactionally (JOE-1832).
198    ///
199    /// Default: sequential single-text calls. Remote backends may batch with
200    /// stable indices and must return only after all work succeeds.
201    async fn cleanup_segments(&self, texts: &[&str], style: CleanupStyle) -> Result<Vec<String>> {
202        let mut out = Vec::with_capacity(texts.len());
203        for t in texts {
204            out.push(self.cleanup(t, style).await?.text);
205        }
206        Ok(out)
207    }
208}
209
210/// Apply cleanup to a full transcription result.
211///
212/// Updates `text` and cleanup metadata. Segment handling follows `segments` policy.
213pub async fn apply_cleanup(
214    result: &mut TranscriptionResult,
215    cleanup: &dyn TextCleanup,
216    style: CleanupStyle,
217) -> Result<CleanupResult> {
218    let (out, _report) =
219        apply_cleanup_with_segments(result, cleanup, style, SegmentCleanupPolicy::Auto).await?;
220    Ok(out)
221}
222
223/// Like [`apply_cleanup`] with explicit segment policy.
224///
225/// **Transactional:** the caller's `TranscriptionResult` is mutated only after
226/// the full operation validates. An injected failure on segment N leaves the
227/// original result byte-for-byte unchanged.
228pub async fn apply_cleanup_with_segments(
229    result: &mut TranscriptionResult,
230    cleanup: &dyn TextCleanup,
231    style: CleanupStyle,
232    segments: SegmentCleanupPolicy,
233) -> Result<(CleanupResult, CleanupReport)> {
234    apply_cleanup_with_segments_op(result, cleanup, style, segments, None).await
235}
236
237/// Like [`apply_cleanup_with_segments`] with an optional [`OpContext`] for cancel,
238/// deadline, and progress (JOE-1831).
239pub async fn apply_cleanup_with_segments_op(
240    result: &mut TranscriptionResult,
241    cleanup: &dyn TextCleanup,
242    style: CleanupStyle,
243    segments: SegmentCleanupPolicy,
244    op: Option<&OpContext>,
245) -> Result<(CleanupResult, CleanupReport)> {
246    let owned_op;
247    let op = match op {
248        Some(o) => o,
249        None => {
250            owned_op = OpContext::new();
251            &owned_op
252        }
253    };
254    let policy = segments.resolve(style);
255    op.emit(
256        "cleanup",
257        format!("style={} policy={policy:?}", style.as_str()),
258    );
259    op.check()?;
260
261    if matches!(style, CleanupStyle::Raw) {
262        // Raw remains the identity default — no metadata pollution.
263        result.set_cleanup_style(CleanupStyle::Raw);
264        result.set_cleanup_provider(None);
265        result.set_original_text(None);
266        result.set_original_segments(None);
267        result.set_cleanup_segment_policy(None);
268        let out = CleanupResult {
269            text: result.text().to_string(),
270            style,
271            provider: cleanup.kind(),
272            original_text: result.text().to_string(),
273        };
274        let report = CleanupReport {
275            style,
276            provider: cleanup.kind(),
277            segment_policy: policy,
278            changed_fields: vec![],
279            warnings: vec![],
280            dropped_segments: 0,
281            segments_cleared: false,
282        };
283        return Ok((out, report));
284    }
285
286    // Bound per-segment work before any mutation.
287    if matches!(policy, SegmentCleanupPolicy::PerSegment) {
288        if result.segments().len() > MAX_PER_SEGMENT_COUNT {
289            return Err(UserError::Other {
290                message: format!(
291                    "per-segment cleanup refused: {} segments exceeds limit of {MAX_PER_SEGMENT_COUNT}",
292                    result.segments().len()
293                ),
294            }
295            .into());
296        }
297        for (i, seg) in result.segments().iter().enumerate() {
298            let n = seg.text().chars().count();
299            if n > MAX_PER_SEGMENT_CHARS {
300                return Err(UserError::Other {
301                    message: format!(
302                        "per-segment cleanup refused: segment {i} has {n} chars \
303                         (limit {MAX_PER_SEGMENT_CHARS})"
304                    ),
305                }
306                .into());
307            }
308        }
309    }
310
311    // Snapshot inputs; mutate only after full success.
312    let original_text = result.text().to_string();
313    let original_segments = result.segments().to_vec();
314    let mut warnings = Vec::new();
315    let mut dropped_segments = 0usize;
316    let mut segments_cleared = false;
317
318    op.emit("cleanup", "full_text");
319    let out = cleanup.cleanup(&original_text, style).await?;
320    op.check()?;
321
322    let proposed_segments = match policy {
323        SegmentCleanupPolicy::Keep | SegmentCleanupPolicy::Auto => {
324            // Keep timings; TXT may diverge from SRT segment text for light styles.
325            if matches!(style, CleanupStyle::Clean | CleanupStyle::Professional) {
326                warnings.push(
327                    "segment timings kept; segment text is pre-cleanup ASR while \
328                     `text` is cleaned (JSON exposes original_text)"
329                        .into(),
330                );
331            }
332            original_segments.clone()
333        }
334        SegmentCleanupPolicy::Clear => {
335            segments_cleared = true;
336            warnings.push(
337                "segments cleared under structural/explicit clear policy; \
338                 use original_segments for raw ASR timings"
339                    .into(),
340            );
341            Vec::new()
342        }
343        SegmentCleanupPolicy::PerSegment => {
344            op.emit("cleanup", "per_segment");
345            // Transactional: no host mutation until every segment cleans successfully
346            // (JOE-1832). OpenRouter batches remotely with stable ids.
347            let refs: Vec<&str> = original_segments.iter().map(|s| s.text()).collect();
348            let cleaned_texts =
349                cleanup
350                    .cleanup_segments(&refs, style)
351                    .await
352                    .map_err(|e| UserError::Other {
353                        message: format!("per-segment cleanup failed (transaction aborted): {e}"),
354                    })?;
355            if cleaned_texts.len() != original_segments.len() {
356                return Err(UserError::Other {
357                    message: format!(
358                        "per-segment cleanup returned {} texts for {} segments",
359                        cleaned_texts.len(),
360                        original_segments.len()
361                    ),
362                }
363                .into());
364            }
365            let mut cleaned = Vec::with_capacity(original_segments.len());
366            for (seg, text) in original_segments.iter().zip(cleaned_texts) {
367                let mut seg = seg.clone();
368                seg.set_text(text);
369                if seg.text().trim().is_empty() {
370                    dropped_segments += 1;
371                } else {
372                    cleaned.push(seg);
373                }
374            }
375            cleaned
376        }
377    };
378
379    // Commit proposed state only now.
380    op.emit("cleanup", "commit");
381    let mut changed_fields = vec![
382        "text".into(),
383        "cleanup_style".into(),
384        "cleanup_provider".into(),
385    ];
386    result.set_original_text(Some(original_text.clone()));
387    result.set_original_segments(Some(original_segments));
388    result.set_text(out.text.clone());
389    result.set_cleanup_style(out.style);
390    result.set_cleanup_provider(Some(out.provider));
391    result.set_cleanup_segment_policy(Some(policy));
392    if result.segments() != proposed_segments {
393        changed_fields.push("segments".into());
394    }
395    result.set_segments(proposed_segments);
396    changed_fields.push("original_text".into());
397    changed_fields.push("original_segments".into());
398
399    let report = CleanupReport {
400        style: out.style,
401        provider: out.provider,
402        segment_policy: policy,
403        changed_fields,
404        warnings,
405        dropped_segments,
406        segments_cleared,
407    };
408    Ok((out, report))
409}
410
411/// Clean a bare string (stdin / text file path) without a full transcription result.
412pub async fn cleanup_text(
413    text: &str,
414    cleanup: &dyn TextCleanup,
415    style: CleanupStyle,
416) -> Result<CleanupResult> {
417    if matches!(style, CleanupStyle::Raw) {
418        return Ok(CleanupResult {
419            text: text.trim().to_string(),
420            style,
421            provider: cleanup.kind(),
422            original_text: text.to_string(),
423        });
424    }
425    cleanup.cleanup(text, style).await
426}
427
428pub use openrouter::OpenRouterCleanup;
429pub use rules::RulesCleanup;
430
431#[cfg(test)]
432mod tests {
433    use super::*;
434    use crate::providers::Segment;
435
436    #[tokio::test]
437    async fn rules_clean_strips_fillers() {
438        let c = RulesCleanup::new();
439        let out = c
440            .cleanup(
441                "um, hello there, you know, this is a test",
442                CleanupStyle::Clean,
443            )
444            .await
445            .unwrap();
446        let lower = out.text.to_ascii_lowercase();
447        assert!(!lower.contains(" um"));
448        assert!(lower.contains("hello"));
449        assert!(lower.contains("test"));
450    }
451
452    #[tokio::test]
453    async fn rules_bullets() {
454        let c = RulesCleanup::new();
455        let out = c
456            .cleanup(
457                "First point here. Second point there. Third idea now.",
458                CleanupStyle::Bullets,
459            )
460            .await
461            .unwrap();
462        assert!(out.text.contains('•'));
463        assert!(out.text.lines().count() >= 2);
464    }
465
466    #[tokio::test]
467    async fn raw_trims_only() {
468        let c = RulesCleanup::new();
469        let out = c.cleanup(" keep this ", CleanupStyle::Raw).await.unwrap();
470        assert_eq!(out.text, "keep this");
471    }
472
473    #[tokio::test]
474    async fn structural_auto_clears_segments() {
475        let c = RulesCleanup::new();
476        let mut result = TranscriptionResult::local(
477            "One. Two. Three.".into(),
478            vec![
479                Segment::from_parts_unchecked(0.0, 1.0, "One.".to_string()),
480                Segment::from_parts_unchecked(1.0, 2.0, "Two.".to_string()),
481            ],
482            Some("en".into()),
483            "tiny".into(),
484            2.0,
485        );
486        let (_out, report) = apply_cleanup_with_segments(
487            &mut result,
488            &c,
489            CleanupStyle::Bullets,
490            SegmentCleanupPolicy::Auto,
491        )
492        .await
493        .unwrap();
494        assert!(result.segments().is_empty());
495        assert!(result.text().contains('•'));
496        assert_eq!(result.cleanup_style(), CleanupStyle::Bullets);
497        assert_eq!(result.cleanup_provider(), Some(CleanupProviderKind::Rules));
498        assert!(report.segments_cleared);
499        assert!(result.original_segments().unwrap().len() == 2);
500        assert_eq!(result.original_text(), Some("One. Two. Three."));
501    }
502
503    #[tokio::test]
504    async fn keep_preserves_segments() {
505        let c = RulesCleanup::new();
506        let mut result = TranscriptionResult::local(
507            "um hello there".into(),
508            vec![Segment::from_parts_unchecked(
509                0.0,
510                1.0,
511                "um hello there".to_string(),
512            )],
513            None,
514            "tiny".into(),
515            1.0,
516        );
517        apply_cleanup_with_segments(
518            &mut result,
519            &c,
520            CleanupStyle::Clean,
521            SegmentCleanupPolicy::Keep,
522        )
523        .await
524        .unwrap();
525        assert_eq!(result.segments().len(), 1);
526        assert_eq!(result.segments()[0].text(), "um hello there");
527        assert!(result.original_text().is_some());
528    }
529
530    /// Failing backend used only for transactional guarantees.
531    struct FailAfterN {
532        n: std::sync::atomic::AtomicUsize,
533        fail_at: usize,
534        inner: RulesCleanup,
535    }
536
537    #[async_trait]
538    impl TextCleanup for FailAfterN {
539        fn name(&self) -> &'static str {
540            "fail-after-n"
541        }
542        fn kind(&self) -> CleanupProviderKind {
543            CleanupProviderKind::Rules
544        }
545        async fn cleanup(&self, text: &str, style: CleanupStyle) -> Result<CleanupResult> {
546            let i = self.n.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
547            if i >= self.fail_at {
548                return Err(UserError::Other {
549                    message: format!("injected failure at call {i}"),
550                }
551                .into());
552            }
553            self.inner.cleanup(text, style).await
554        }
555    }
556
557    #[tokio::test]
558    async fn per_segment_failure_leaves_result_unchanged() {
559        let backend = FailAfterN {
560            n: std::sync::atomic::AtomicUsize::new(0),
561            // Full-transcript cleanup succeeds (call 0); first segment rewrite fails (call 1).
562            fail_at: 1,
563            inner: RulesCleanup::new(),
564        };
565        let mut result = TranscriptionResult::local(
566            "um one. um two.".into(),
567            vec![
568                Segment::from_parts_unchecked(0.0, 1.0, "um one.".to_string()),
569                Segment::from_parts_unchecked(1.0, 2.0, "um two.".to_string()),
570            ],
571            None,
572            "tiny".into(),
573            2.0,
574        );
575        let before = serde_json::to_string(&result).unwrap();
576        let err = apply_cleanup_with_segments(
577            &mut result,
578            &backend,
579            CleanupStyle::Clean,
580            SegmentCleanupPolicy::PerSegment,
581        )
582        .await
583        .unwrap_err();
584        assert!(err.to_string().contains("segment"), "{err}");
585        let after = serde_json::to_string(&result).unwrap();
586        assert_eq!(before, after, "result must be unchanged after failure");
587    }
588
589    #[tokio::test]
590    async fn cleanup_text_standalone() {
591        let c = RulesCleanup::new();
592        let out = cleanup_text("um, hi there", &c, CleanupStyle::Clean)
593            .await
594            .unwrap();
595        assert!(!out.text.to_ascii_lowercase().contains("um"));
596    }
597}
aurum_core/cleanup/mod.rs

aurum_core/cleanup/
mod.rs
