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

//! Plain-text pre-formatter for LLM cleanup.
//!
//! Brief item #29: before sending transcription segments to the LLM,
//! join them into a single natural-language string with timestamps
//! stripped and whitespace normalised. Source: Scriberr PR #288 —
//! feeding raw Whisper JSON (with its timestamps and per-segment
//! structure) degraded cleanup quality materially; plain-text input
//! raised it back.
//!
//! `Segment.text` in Kon already holds just the spoken text (the
//! `start`/`end` f64 fields carry the timing), so "timestamp
//! stripping" falls out of using the text field alone. The work here
//! is the whitespace pass and empty-segment filter, plus a single
//! public function the pipeline can depend on.

use kon_core::types::Segment;

/// Join transcription segments into a single plain-text string
/// suitable for feeding to an LLM cleanup prompt.
///
/// Rules:
/// - each segment's text is whitespace-normalised (any run of
///   whitespace — spaces, tabs, newlines, non-breaking spaces —
///   collapses to a single ASCII space),
/// - segments that are empty or whitespace-only are dropped,
/// - the remaining segments are joined with a single ASCII space,
/// - the final string is whitespace-normalised again (so a segment
///   ending in a space and the next beginning with one do not produce
///   a double space) and trimmed of leading/trailing whitespace.
///
/// Pure function. No panics. Returns an empty string if every segment
/// filters out.
pub fn to_plain_text(segments: &[Segment]) -> String {
    let joined = segments
        .iter()
        .map(|s| normalise_whitespace(&s.text))
        .map(|s| s.trim().to_string())
        .filter(|s| !s.is_empty())
        .collect::<Vec<_>>()
        .join(" ");
    normalise_whitespace(&joined).trim().to_string()
}

/// Collapse any run of unicode whitespace into a single ASCII space.
/// Kept private; the module's contract is `to_plain_text`.
fn normalise_whitespace(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut prev_was_space = false;
    for ch in s.chars() {
        if ch.is_whitespace() {
            if !prev_was_space {
                out.push(' ');
                prev_was_space = true;
            }
        } else {
            out.push(ch);
            prev_was_space = false;
        }
    }
    out
}

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

    fn seg(text: &str) -> Segment {
        Segment {
            start: 0.0,
            end: 1.0,
            text: text.into(),
        }
    }

    #[test]
    fn empty_input_is_empty_output() {
        assert_eq!(to_plain_text(&[]), "");
    }

    #[test]
    fn single_segment_returns_its_text_trimmed() {
        let out = to_plain_text(&[seg("  hello world  ")]);
        assert_eq!(out, "hello world");
    }

    #[test]
    fn multiple_segments_are_joined_with_single_space() {
        let out = to_plain_text(&[seg("the cat"), seg("sat on the mat")]);
        assert_eq!(out, "the cat sat on the mat");
    }

    #[test]
    fn empty_and_whitespace_segments_are_filtered() {
        let out = to_plain_text(&[
            seg("hello"),
            seg(""),
            seg("   "),
            seg("\n\t "),
            seg("world"),
        ]);
        assert_eq!(out, "hello world");
    }

    #[test]
    fn internal_whitespace_runs_collapse_to_single_space() {
        let out = to_plain_text(&[seg("hello\t\t   \nworld")]);
        assert_eq!(out, "hello world");
    }

    #[test]
    fn join_boundary_does_not_produce_double_spaces() {
        // First segment ends with whitespace, next starts with it —
        // naive join would produce "foo  bar".
        let out = to_plain_text(&[seg("foo "), seg(" bar")]);
        assert_eq!(out, "foo bar");
    }

    #[test]
    fn non_breaking_space_is_treated_as_whitespace() {
        // \u{00A0} is NBSP — char::is_whitespace returns true for it.
        // LLM cleanup should not see NBSP leaked in.
        let out = to_plain_text(&[seg("hello\u{00A0}world")]);
        assert_eq!(out, "hello world");
    }

    #[test]
    fn newlines_inside_segments_collapse() {
        let out = to_plain_text(&[seg("line one\nline two\n\nline three")]);
        assert_eq!(out, "line one line two line three");
    }

    #[test]
    fn idempotent_on_already_normalised_text() {
        // If the pipeline ever calls us twice, the second call must
        // not mangle the result.
        let once = to_plain_text(&[seg("hello world"), seg("foo bar")]);
        let twice = to_plain_text(&[seg(&once)]);
        assert_eq!(once, twice);
    }

    #[test]
    fn only_empty_segments_yields_empty_string() {
        let out = to_plain_text(&[seg(""), seg("   "), seg("\t")]);
        assert_eq!(out, "");
    }

    #[test]
    fn no_panic_on_pathological_whitespace_runs() {
        // A segment that is 10k spaces long normalises in linear time
        // without panicking on capacity guesses.
        let big_spaces = " ".repeat(10_000);
        let out = to_plain_text(&[seg(&format!("a{big_spaces}b"))]);
        assert_eq!(out, "a b");
    }
}