inferencelayer 0.2.3

//! Stop-sequence scanning over the running detokenized text.
//!
//! OpenAI semantics: generation halts at the FIRST occurrence of any `stop` string, and the stop
//! text itself is EXCLUDED from the returned completion. A stop string can straddle a token
//! boundary (the tokenizer never emits it as one piece), so the scan runs on the accumulated
//! decoded TEXT, not on token ids. Truncation is UTF-8-safe by construction: [`str::find`] returns
//! the byte offset at the start of the match, which is always a char boundary.
//!
//! For streaming, [`safe_stream_end`] holds back the last `maxstop − 1` bytes of not-yet-emitted
//! text so a stop string forming across future tokens is never partially streamed before it can be
//! recognized and cut.

/// The earliest stop match in `text`: `(byte_offset, matched_stop)` where `byte_offset` is where
/// the match STARTS, so the caller keeps `&text[..byte_offset]`. Empty stop strings are ignored
/// (they would match at 0 and truncate everything). Ties (two stops at the same offset) resolve to
/// the longer stop, then to the earlier one in `stops`.
pub fn first_stop(text: &str, stops: &[String]) -> Option<(usize, String)> {
    let mut best: Option<(usize, String)> = None;
    for s in stops {
        if s.is_empty() {
            continue;
        }
        if let Some(pos) = text.find(s.as_str()) {
            let better = match &best {
                None => true,
                Some((b, bs)) => pos < *b || (pos == *b && s.len() > bs.len()),
            };
            if better {
                best = Some((pos, s.clone()));
            }
        }
    }
    best
}

/// How many bytes of `text` are safe to STREAM given the pending `stops`: the whole string minus a
/// tail that could still become the prefix of a stop string on the next token. Returns a byte
/// offset on a char boundary (walked back from the raw cut so slicing `&text[..end]` never splits a
/// codepoint). When no stops are configured the whole string is safe.
pub fn safe_stream_end(text: &str, stops: &[String]) -> usize {
    let maxstop = stops.iter().map(String::len).max().unwrap_or(0);
    if maxstop <= 1 {
        return text.len();
    }
    let mut end = text.len().saturating_sub(maxstop - 1);
    while end > 0 && !text.is_char_boundary(end) {
        end -= 1;
    }
    end
}

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

    fn s(v: &[&str]) -> Vec<String> {
        v.iter().map(|x| x.to_string()).collect()
    }

    #[test]
    fn should_find_earliest_stop_and_report_its_start() {
        let (pos, which) = first_stop("hello END world", &s(&["END"])).expect("match");
        assert_eq!(pos, 6);
        assert_eq!(which, "END");
        assert_eq!(&"hello END world"[..pos], "hello ");
    }

    #[test]
    fn should_pick_the_earliest_when_multiple_stops_match() {
        let (pos, which) = first_stop("a STOP b HALT c", &s(&["HALT", "STOP"])).expect("match");
        assert_eq!(pos, 2);
        assert_eq!(which, "STOP");
    }

    #[test]
    fn should_ignore_empty_stop_strings() {
        assert!(first_stop("anything", &s(&[""])).is_none());
    }

    #[test]
    fn should_return_none_when_no_stop_present() {
        assert!(first_stop("nothing here", &s(&["END", "STOP"])).is_none());
    }

    #[test]
    fn should_truncate_on_a_utf8_boundary() {
        // A stop right after a multi-byte codepoint: the cut must land on a char boundary.
        let text = "café STOP";
        let (pos, _) = first_stop(text, &s(&["STOP"])).expect("match");
        assert!(text.is_char_boundary(pos));
        assert_eq!(&text[..pos], "café ");
    }

    #[test]
    fn should_hold_back_potential_stop_prefix_when_streaming() {
        // "STOP" is 4 bytes; the last 3 bytes ("STO") could be a stop-in-progress.
        let end = safe_stream_end("hello STO", &s(&["STOP"]));
        assert_eq!(&"hello STO"[..end], "hello ");
    }

    #[test]
    fn should_stream_everything_when_no_stops() {
        assert_eq!(safe_stream_end("hello world", &[]), "hello world".len());
    }

    #[test]
    fn safe_stream_end_never_splits_a_codepoint() {
        // Hold-back boundary walked back off the middle of a 2-byte 'é'.
        let text = "abcé";
        let end = safe_stream_end(text, &s(&["xé"]));
        assert!(text.is_char_boundary(end));
    }
}