Lumotia/crates/ai-formatting/src/pipeline.rs

use kon_core::constants::SMART_PARAGRAPH_GAP_SECS;
use kon_core::types::Segment;
use kon_llm::LlmEngine;

use crate::{llm_client, rule_based, to_plain_text::to_plain_text};

/// Post-processing options for a transcription pipeline run.
pub struct PostProcessOptions {
    pub remove_fillers: bool,
    pub british_english: bool,
    pub anti_hallucination: bool,
    pub format_mode: FormatMode,
    /// Custom vocabulary terms loaded from the user's dictionary. Injected
    /// into the LLM cleanup prompt so the model knows how to spell them.
    pub dictionary_terms: Vec<String>,
}

/// How aggressively to format the transcript text.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FormatMode {
    Raw,
    Clean,
    Smart,
}

impl FormatMode {
    pub fn parse(s: &str) -> Self {
        match s {
            "Clean" => Self::Clean,
            "Smart" => Self::Smart,
            _ => Self::Raw,
        }
    }
}

/// Apply all post-processing steps to a list of segments.
/// Modifies segments in place. Composed from individual pure functions.
pub fn post_process_segments(
    segments: &mut Vec<Segment>,
    options: &PostProcessOptions,
    llm: Option<&LlmEngine>,
) {
    if options.anti_hallucination {
        segments.retain(|seg| !rule_based::is_hallucination(&seg.text));
    }

    for seg in segments.iter_mut() {
        if options.remove_fillers {
            seg.text = rule_based::remove_fillers(&seg.text);
        }
        if options.british_english {
            seg.text = rule_based::to_british_english(&seg.text);
        }
        if options.format_mode != FormatMode::Raw {
            seg.text = rule_based::collapse_repetitions(&seg.text);
            seg.text = rule_based::format_text(&seg.text);
        }
    }

    if options.format_mode == FormatMode::Smart && segments.len() > 1 {
        for i in (1..segments.len()).rev() {
            let gap = segments[i].start - segments[i - 1].end;
            if gap > SMART_PARAGRAPH_GAP_SECS {
                segments[i].text = format!("\n\n{}", segments[i].text.trim_start());
            }
        }
    }

    if let Some(engine) = llm {
        if engine.is_loaded() && options.format_mode != FormatMode::Raw {
            // Plain-text pre-formatter (brief item #29): collapse
            // segments into a single natural-language string before
            // the LLM call. Whitespace normalisation + empty-filter
            // live in `to_plain_text`; the pipeline's job here is
            // deciding whether to invoke the LLM at all.
            let joined = to_plain_text(segments);

            if !joined.is_empty() {
                // Pipeline-internal cleanup (used by file-based + live
                // transcribe paths) runs with the Default preset. The
                // named-preset UX (B.1 #15) flows through the explicit
                // cleanup_transcript_text_cmd path instead, where the
                // frontend decides which preset the user has selected.
                match llm_client::cleanup_text(
                    engine,
                    &joined,
                    &options.dictionary_terms,
                    llm_client::LlmPromptPreset::Default,
                ) {
                    Ok(cleaned) if !cleaned.trim().is_empty() => {
                        replace_segments_with_cleaned(segments, cleaned.trim());
                    }
                    Ok(_) => {}
                    Err(err) => eprintln!(
                        "[ai-formatting] LLM cleanup failed, keeping rule-based output: {err}"
                    ),
                }
            }
        }
    }
}

fn replace_segments_with_cleaned(segments: &mut Vec<Segment>, cleaned: &str) {
    if segments.is_empty() || cleaned.trim().is_empty() {
        return;
    }

    let start = segments.first().map(|segment| segment.start).unwrap_or(0.0);
    let end = segments.last().map(|segment| segment.end).unwrap_or(start);
    segments.clear();
    segments.push(Segment {
        start,
        end,
        text: cleaned.to_string(),
    });
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_segments() -> Vec<Segment> {
        vec![
            Segment {
                start: 0.0,
                end: 1.0,
                text: "um hello world".into(),
            },
            Segment {
                start: 1.0,
                end: 2.0,
                text: "[blank_audio]".into(),
            },
            Segment {
                start: 5.0,
                end: 6.0,
                text: "organize the color scheme".into(),
            },
        ]
    }

    #[test]
    fn dictionary_terms_stored_on_options() {
        let options = PostProcessOptions {
            remove_fillers: false,
            british_english: false,
            anti_hallucination: false,
            format_mode: FormatMode::Raw,
            dictionary_terms: vec!["Wren".to_string(), "CORBEL".to_string()],
        };
        assert_eq!(options.dictionary_terms.len(), 2);
        assert_eq!(options.dictionary_terms[0], "Wren");
    }

    #[test]
    fn post_process_applies_all_filters() {
        let mut segments = make_segments();
        let options = PostProcessOptions {
            remove_fillers: true,
            british_english: true,
            anti_hallucination: true,
            format_mode: FormatMode::Clean,
            dictionary_terms: vec![],
        };

        post_process_segments(&mut segments, &options, None);

        assert_eq!(segments.len(), 2);
        let lower0 = segments[0].text.to_lowercase();
        let lower1 = segments[1].text.to_lowercase();
        assert!(!lower0.contains("um"));
        assert!(lower1.contains("organise"));
        assert!(lower1.contains("colour"));
    }

    #[test]
    fn post_process_adds_paragraph_breaks_on_long_pauses() {
        let mut segments = make_segments();
        let options = PostProcessOptions {
            remove_fillers: false,
            british_english: false,
            anti_hallucination: false,
            format_mode: FormatMode::Smart,
            dictionary_terms: vec![],
        };

        post_process_segments(&mut segments, &options, None);

        assert!(segments[2].text.starts_with("\n\n"));
    }

    #[test]
    fn post_process_collapses_repeated_phrases_in_clean_modes() {
        let mut segments = vec![Segment {
            start: 0.0,
            end: 1.0,
            text: "I need I need to go to the shops".into(),
        }];
        let options = PostProcessOptions {
            remove_fillers: false,
            british_english: false,
            anti_hallucination: false,
            format_mode: FormatMode::Clean,
            dictionary_terms: vec![],
        };

        post_process_segments(&mut segments, &options, None);

        assert_eq!(segments[0].text, "I need to go to the shops");
    }
}