path: root/crates/asist-core/src/text.rs

//! Preparing the text that will be read aloud.
//!
//! Two different problems, both on the latency critical path:
//!
//! 1. **Splitting into sentences while the model writes.** Waiting for the
//!    whole answer before synthesizing adds the LLM time to the TTS time.
//!    Cutting per sentence, the assistant starts speaking while still thinking.
//! 2. **Removing what is not pronounced.** The synthesizer reads asterisks and
//!    hashes literally, so markdown that slips out of the model must be
//!    cleaned first.

/// Splits a text stream into speakable sentences.
///
/// The first sentence goes out as soon as it is minimally decent and the next
/// ones wait until they have more body: the start dominates perceived
/// latency, but once the voice is playing, longer sentences get better
/// intonation.
#[derive(Debug)]
pub struct SentenceSplitter {
    buffer: String,
    emitted: usize,
    /// Minimum length of the first sentence.
    first_min: usize,
    /// Minimum length of the following ones.
    rest_min: usize,
    /// Length after which the text is cut even without punctuation, so a
    /// model that does not punctuate does not leave the assistant mute.
    hard_max: usize,
}

impl Default for SentenceSplitter {
    fn default() -> Self {
        Self {
            buffer: String::new(),
            emitted: 0,
            first_min: 12,
            rest_min: 40,
            hard_max: 240,
        }
    }
}

impl SentenceSplitter {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_limits(first_min: usize, rest_min: usize, hard_max: usize) -> Self {
        Self {
            first_min,
            rest_min,
            hard_max,
            ..Self::default()
        }
    }

    pub fn emitted(&self) -> usize {
        self.emitted
    }

    /// Adds newly arrived text and returns the sentences already complete.
    pub fn push(&mut self, delta: &str) -> Vec<String> {
        self.buffer.push_str(delta);
        let mut out = Vec::new();
        while let Some(sentence) = self.take_ready() {
            out.push(sentence);
        }
        out
    }

    /// Delivers whatever is left when the answer ends.
    pub fn flush(&mut self) -> Option<String> {
        let rest = clean_for_speech(&std::mem::take(&mut self.buffer));
        if !is_speakable(&rest) {
            return None;
        }
        self.emitted += 1;
        Some(rest)
    }

    fn min_len(&self) -> usize {
        if self.emitted == 0 {
            self.first_min
        } else {
            self.rest_min
        }
    }

    fn take_ready(&mut self) -> Option<String> {
        let cut = self.boundary()?;
        let head: String = self.buffer.drain(..cut).collect();
        let head = clean_for_speech(&head);
        if !is_speakable(&head) {
            // It was only punctuation or markdown: dropped without spending a
            // synthesis turn, but the cut has been consumed.
            return self.take_ready();
        }
        self.emitted += 1;
        Some(head)
    }

    /// Byte index to cut at, if any.
    fn boundary(&self) -> Option<usize> {
        let min = self.min_len();
        let mut last_soft = None;

        for (i, c) in self.buffer.char_indices() {
            let end = i + c.len_utf8();
            if end < min {
                continue;
            }
            if is_terminator(c) {
                // The dot of an abbreviation or a decimal does not end a sentence.
                if c == '.' && !ends_sentence(&self.buffer, i) {
                    continue;
                }
                // It only closes if a space follows, or if nothing is left:
                // otherwise it would cut in the middle of «3.14».
                match self.buffer[end..].chars().next() {
                    Some(next) if next.is_whitespace() => return Some(end),
                    None => {}
                    Some(_) => continue,
                }
            }
            // A comma or semicolon counts as an emergency cut if the sentence
            // has already run too long.
            if matches!(c, ',' | ';' | ':') {
                last_soft = Some(end);
            }
            if end >= self.hard_max {
                return Some(last_soft.unwrap_or(end));
            }
        }
        None
    }
}

/// Is there anything to pronounce?
///
/// A fragment with only punctuation (what is left after cleaning a «**» or a
/// stray ellipsis) costs a whole synthesis and sounds like nothing, so it
/// never leaves the splitter.
pub fn is_speakable(text: &str) -> bool {
    text.chars().any(char::is_alphanumeric)
}

fn is_terminator(c: char) -> bool {
    matches!(c, '.' | '!' | '?' | '…' | '\n')
}

/// Does the dot at `idx` really end a sentence?
///
/// Rules out decimals («3.14»), half-written ellipses and the most common
/// Spanish abbreviations, which would otherwise split the sentence right
/// before the name.
fn ends_sentence(text: &str, idx: usize) -> bool {
    let before = &text[..idx];
    let after = &text[idx + 1..];

    if after.starts_with(|c: char| c.is_ascii_digit())
        && before.ends_with(|c: char| c.is_ascii_digit())
    {
        return false;
    }
    if after.starts_with('.') || before.ends_with('.') {
        return false;
    }
    let word = before
        .rsplit(|c: char| c.is_whitespace())
        .next()
        .unwrap_or("")
        .to_lowercase();
    const ABBREVIATIONS: &[&str] = &[
        "sr", "sra", "srta", "dr", "dra", "ud", "uds", "etc", "ej", "p.ej", "av", "núm", "num",
        "pág", "pag", "vol", "art", "ap", "aprox", "ee.uu", "d", "dña",
    ];
    !ABBREVIATIONS.contains(&word.as_str())
}

/// Gets the text ready for the synthesizer.
///
/// It only removes what is not pronounced; it neither rewrites nor
/// summarizes, because what is heard must be what the model said.
pub fn clean_for_speech(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    let mut chars = text.chars().peekable();
    let mut at_line_start = true;

    while let Some(c) = chars.next() {
        match c {
            // Emphasis and code: the delimiters would be read aloud.
            '*' | '_' | '`' | '~' => continue,
            // Headings and bullets, only at the start of a line: a hash or a
            // dash in the middle of a sentence does mean something.
            '#' if at_line_start => {
                while chars.peek() == Some(&'#') {
                    chars.next();
                }
                while chars.peek().is_some_and(|c| *c == ' ') {
                    chars.next();
                }
                continue;
            }
            '-' | '•' | '–' if at_line_start && chars.peek() == Some(&' ') => {
                chars.next();
                continue;
            }
            '>' if at_line_start => {
                while chars.peek().is_some_and(|c| *c == ' ') {
                    chars.next();
                }
                continue;
            }
            // Line breaks are spoken as pauses.
            '\n' | '\r' | '\t' => {
                at_line_start = c == '\n';
                if !out.ends_with(' ') && !out.is_empty() {
                    out.push(' ');
                }
                continue;
            }
            _ => {}
        }
        at_line_start = false;
        out.push(c);
    }

    // Collapses the spaces the cleanup leaves behind.
    let mut collapsed = String::with_capacity(out.len());
    let mut space = false;
    for c in out.chars() {
        if c == ' ' {
            space = true;
            continue;
        }
        if space && !collapsed.is_empty() {
            collapsed.push(' ');
        }
        space = false;
        collapsed.push(c);
    }
    collapsed.trim().to_string()
}

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

    fn split_all(chunks: &[&str]) -> Vec<String> {
        let mut splitter = SentenceSplitter::new();
        let mut out: Vec<String> = chunks.iter().flat_map(|c| splitter.push(c)).collect();
        out.extend(splitter.flush());
        out
    }

    #[test]
    fn first_sentence_comes_out_before_the_rest() {
        // The first one only needs to be short; the second waits for body.
        let out = split_all(&["Claro que sí. ", "Vale. ", "Aún no.", ""]);
        assert_eq!(out[0], "Claro que sí.");
        assert_eq!(out[1], "Vale. Aún no.");
    }

    #[test]
    fn chunks_as_deltas_arrive() {
        let mut splitter = SentenceSplitter::new();
        assert!(splitter.push("Hola, ¿qué ").is_empty());