const OPEN: &str = "<think>";
const CLOSE: &str = "</think>";
#[must_use]
pub fn strip_reasoning(text: &str) -> String {
let mut rest = match (text.find(OPEN), text.find(CLOSE)) {
(open, Some(close)) if open.is_none_or(|open| open > close) => &text[close + CLOSE.len()..],
_ => text,
};
let mut kept = String::with_capacity(rest.len());
let mut joined = false;
while let Some(start) = rest.find(OPEN) {
push_segment(&mut kept, &rest[..start], joined);
joined = true;
rest = match rest[start..].find(CLOSE) {
Some(end) => &rest[start + end + CLOSE.len()..],
None => "",
};
}
push_segment(&mut kept, rest, joined);
kept.trim().to_string()
}
fn push_segment(kept: &mut String, segment: &str, joined: bool) {
if joined && kept.ends_with(' ') {
kept.push_str(segment.trim_start_matches(' '));
} else {
kept.push_str(segment);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn reasoning_is_removed_from_an_answer() {
assert_eq!(strip_reasoning(" Borrow it. "), "Borrow it.");
assert_eq!(
strip_reasoning("<think>why</think>\nBorrow it."),
"Borrow it."
);
assert_eq!(
strip_reasoning("A <think>x</think>B<think>y</think> C"),
"A B C"
);
assert_eq!(strip_reasoning("Answer.<think>never closed"), "Answer.");
assert_eq!(
strip_reasoning("thinking first</think>\nAnswer."),
"Answer."
);
}
#[test]
fn indentation_is_kept() {
let code = "def f():\n return 1";
assert_eq!(strip_reasoning(code), code);
assert_eq!(strip_reasoning(&format!("<think>x</think>\n{code}")), code);
assert_eq!(strip_reasoning(&format!("{code}<think>x</think>")), code);
}
}