newt_core/

reasoning.rs

//! Separating a model's **reasoning** from its **content**.
//!
//! Thinking models (Nemotron with `detailed thinking on`, DeepSeek-R1, Qwen3, …)
//! emit their chain-of-thought *inline* in the content stream, wrapped in
//! `<think>…</think>` tags — distinct from the Ollama/OpenAI *separate* reasoning
//! field the loop already detects. Left in place, that reasoning leaks into the
//! reply shown to the user and into the content the agentic loop parses (issue
//! #385). This module strips it.
//!
//! It's a **split**, not just a strip: [`split_reasoning`] returns
//! `(content, reasoning)` so the reasoning can be *captured* (the `plan_mode`
//! technique, `docs/design/thinking-effort-and-plan-mode.md`) rather than thrown
//! away — today every caller discards the reasoning half, matching prior behavior.
//!
//! Two surfaces:
//! - [`split_reasoning`] — batch, for a fully-assembled reply (the non-streaming
//!   completion paths).
//! - [`ThinkFilter`] — incremental, for the streaming path, where a `<think>` block
//!   spans token boundaries (`<thi` / `nk>…</thi` / `nk>`) and must be suppressed
//!   live without ever printing a partial tag.

const OPEN: &str = "<think>";
const CLOSE: &str = "</think>";

/// Split inline `<think>…</think>` reasoning out of `content`, returning
/// `(clean_content, reasoning)`.
///
/// Handles the shapes seen in the wild:
/// - paired `<think>R</think>A` (one or many, anywhere) → reasoning removed;
/// - a **lone leading** `R</think>A` (some templates start mid-thought, emitting
///   only the closer) → everything up to the first `</think>` is reasoning;
/// - an **unterminated** `A<think>R` (a truncated thinking block) → everything from
///   `<think>` on is reasoning.
///
/// **No-op fast path:** content with no `<think>`/`</think>` is returned unchanged
/// (not even trimmed), so ordinary replies are bit-for-bit untouched.
#[must_use]
pub fn split_reasoning(content: &str) -> (String, Option<String>) {
    // Lone leading closer (no opener) — only meaningful when there is no `<think>`.
    if !content.contains(OPEN) {
        return match content.find(CLOSE) {
            Some(close) => {
                let reasoning = content[..close].trim();
                let clean = content[close + CLOSE.len()..].trim();
                (clean.to_string(), non_empty(reasoning))
            }
            None => (content.to_string(), None), // nothing to strip — untouched
        };
    }

    let mut clean = String::new();
    let mut reasoning = String::new();
    let mut rest = content;
    while let Some(open) = rest.find(OPEN) {
        clean.push_str(&rest[..open]);
        let after = &rest[open + OPEN.len()..];
        match after.find(CLOSE) {
            Some(close) => {
                reasoning.push_str(&after[..close]);
                rest = &after[close + CLOSE.len()..];
            }
            None => {
                // Unterminated <think>: the remainder is all reasoning.
                reasoning.push_str(after);
                rest = "";
                break;
            }
        }
    }
    clean.push_str(rest);
    (clean.trim().to_string(), non_empty(reasoning.trim()))
}

fn non_empty(s: &str) -> Option<String> {
    if s.is_empty() {
        None
    } else {
        Some(s.to_string())
    }
}

/// An incremental filter that suppresses `<think>…</think>` spans from a *streamed*
/// content sequence, emitting only the clean (non-reasoning) text — even when a tag
/// is split across feeds.
///
/// Feed each streamed token through [`feed`](Self::feed) and print/accumulate what
/// it returns; call [`finish`](Self::finish) once at the end. A trailing partial tag
/// (`…<thi`) is held back rather than emitted, so a tag is never printed in pieces.
#[derive(Debug, Default)]
pub struct ThinkFilter {
    inside: bool,
    /// `inside` was *assumed* (lone-leading-closer mode) rather than entered via
    /// a literal `<think>`, so buffered content is only provisionally reasoning:
    /// if no `</think>` ever arrives it is flushed as clean (the model wasn't
    /// actually emitting reasoning), so a reply is never lost. See
    /// [`with_leading_reasoning`](Self::with_leading_reasoning).
    implicit: bool,
    buf: String,
}

impl ThinkFilter {
    /// A fresh filter (outside any think block).
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// A filter that begins *inside* a reasoning block — for models that stream
    /// the **lone-leading-closer** shape (`reasoning</think>answer`, with no
    /// opening `<think>`), e.g. Nemotron with `detailed thinking on`. Everything
    /// up to the first `</think>` is treated as reasoning, matching
    /// [`split_reasoning`]'s batch handling of the same shape (#528).
    ///
    /// The leading block is held *provisionally*: if a `</think>` never arrives
    /// (the model wasn't actually reasoning — e.g. thinking turned off), the
    /// buffered text is flushed as clean by [`finish`](Self::finish), so a reply
    /// is never dropped. Gate via [`emits_leading_reasoning`].
    #[must_use]
    pub fn with_leading_reasoning() -> Self {
        Self {
            inside: true,
            implicit: true,
            buf: String::new(),
        }
    }

    /// Feed one streamed token; returns the clean text to emit now (may be empty).
    pub fn feed(&mut self, token: &str) -> String {
        self.feed_split(token).0
    }

    /// Feed one streamed token; returns `(clean, reasoning)` for this feed — the
    /// clean text to emit now AND the reasoning text suppressed from it (either
    /// may be empty). Lets a caller *render* the live reasoning (the cargo-style
    /// thinking spinner) instead of discarding it, while `feed` keeps the
    /// suppress-only behavior for everyone else.
    pub fn feed_split(&mut self, token: &str) -> (String, String) {
        self.buf.push_str(token);
        let mut clean = String::new();
        let mut reasoning = String::new();
        loop {
            if self.inside {
                match self.buf.find(CLOSE) {
                    Some(i) => {
                        reasoning.push_str(&self.buf[..i]);
                        self.buf.drain(..i + CLOSE.len());
                        self.inside = false; // continue: there may be clean text after
                        self.implicit = false; // the closer resolved a leading block
                    }
                    None => {
                        if self.implicit {
                            // Assumed leading reasoning: keep everything buffered
                            // until a closer proves it was reasoning. finish()
                            // flushes it as clean if none ever arrives, so a
                            // non-thinking reply is not lost.
                            break;
                        }
                        // Suppress reasoning, but keep a trailing partial `</think>`.
                        let cut = safe_len(&self.buf, CLOSE);
                        reasoning.push_str(&self.buf[..cut]);
                        self.buf.drain(..cut);
                        break;
                    }
                }
            } else {
                match self.buf.find(OPEN) {
                    Some(i) => {
                        clean.push_str(&self.buf[..i]);
                        self.buf.drain(..i + OPEN.len());
                        self.inside = true;
                    }
                    None => {
                        // Emit all but a trailing partial `<think>`.
                        let cut = safe_len(&self.buf, OPEN);
                        clean.push_str(&self.buf[..cut]);
                        self.buf.drain(..cut);
                        break;
                    }
                }
            }
        }
        (clean, reasoning)
    }

    /// Flush at end of stream: emits any buffered clean tail (an unterminated
    /// `<think>` leaves its reasoning suppressed).
    pub fn finish(&mut self) -> String {
        // A real unterminated `<think>` suppresses its reasoning; an *assumed*
        // (implicit) leading block that never closed turned out to be the reply
        // itself, so flush it as clean — no data loss for a non-thinking model.
        let out = if self.inside && !self.implicit {
            String::new()
        } else {
            std::mem::take(&mut self.buf)
        };
        self.buf.clear();
        self.inside = false;
        self.implicit = false;
        out
    }
}

/// Whether a model streams its chain-of-thought as a **lone leading closer**
/// (`reasoning</think>answer`, with no opening `<think>`) — the shape Nemotron
/// (`detailed thinking on`), DeepSeek-R1, and Qwen3 emit. Such models need
/// [`ThinkFilter::with_leading_reasoning`] so the streamed reasoning (and the
/// bare `</think>`) is suppressed rather than printed into the reply.
///
/// Name-based stopgap until a per-model capability card carries this flag
/// (#384); matched case-insensitively against the known families.
#[must_use]
pub fn emits_leading_reasoning(model: &str) -> bool {
    let m = model.to_ascii_lowercase();
    ["nemotron", "deepseek-r1", "qwen3"]
        .iter()
        .any(|fam| m.contains(fam))
}

/// The length of `buf` that is safe to commit (emit *or* discard) without splitting
/// a possible `tag` — i.e. `buf` minus its longest suffix that is a proper prefix of
/// `tag`. (`"a<thi"`, tag `"<think>"` → `1`, holding back `"<thi"`.)
fn safe_len(buf: &str, tag: &str) -> usize {
    let max = buf.len().min(tag.len() - 1);
    for k in (1..=max).rev() {
        let start = buf.len() - k;
        if buf.is_char_boundary(start) && tag.starts_with(&buf[start..]) {
            return start;
        }
    }
    buf.len()
}

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

    #[test]
    fn no_tags_is_untouched() {
        let (c, r) = split_reasoning("just a normal reply");
        assert_eq!(c, "just a normal reply");
        assert!(r.is_none());
    }

    #[test]
    fn strips_a_leading_paired_block() {
        let (c, r) = split_reasoning("<think>let me think</think>\n\nThe answer is 42.");
        assert_eq!(c, "The answer is 42.");
        assert_eq!(r.as_deref(), Some("let me think"));
    }

    #[test]
    fn strips_multiple_and_embedded_blocks() {
        let (c, r) = split_reasoning("A<think>r1</think>B<think>r2</think>C");
        assert_eq!(c, "ABC");
        assert_eq!(r.as_deref(), Some("r1r2"));
    }

    #[test]
    fn lone_leading_closer_is_reasoning() {
        let (c, r) = split_reasoning("reasoning with no opener</think>the answer");
        assert_eq!(c, "the answer");
        assert_eq!(r.as_deref(), Some("reasoning with no opener"));
    }

    #[test]
    fn unterminated_open_swallows_the_tail() {
        let (c, r) = split_reasoning("partial answer<think>cut off mid-thought");
        assert_eq!(c, "partial answer");
        assert_eq!(r.as_deref(), Some("cut off mid-thought"));
    }

    #[test]
    fn all_reasoning_yields_empty_content() {
        let (c, r) = split_reasoning("<think>only thinking, no answer</think>");
        assert_eq!(c, "");
        assert_eq!(r.as_deref(), Some("only thinking, no answer"));
    }

    /// Drive the streaming filter with an arbitrary token split and assert it equals
    /// the batch result's clean half.
    fn stream(content: &str, tokens: &[&str]) -> String {
        let mut f = ThinkFilter::new();
        let mut out = String::new();
        for t in tokens {
            out.push_str(&f.feed(t));
        }
        out.push_str(&f.finish());
        let _ = content;
        out
    }

    #[test]
    fn streaming_suppresses_a_block_split_across_tokens() {
        // The <think> tags are deliberately shredded across token boundaries.
        let tokens = [
            "Here",
            " <thi",
            "nk>my ",
            "reason",
            "ing</thi",
            "nk> is the ",
            "answer",
        ];
        assert_eq!(stream("", &tokens), "Here  is the answer");
    }

    #[test]
    fn streaming_char_by_char_matches_batch() {
        let content = "<think>step one\nstep two</think>Final line.";
        let toks: Vec<String> = content.chars().map(|c| c.to_string()).collect();
        let refs: Vec<&str> = toks.iter().map(String::as_str).collect();
        assert_eq!(stream(content, &refs), "Final line.");
    }

    #[test]
    fn streaming_plain_text_is_verbatim() {
        assert_eq!(stream("", &["no ", "tags ", "here"]), "no tags here");
    }

    #[test]
    fn feed_split_captures_reasoning_across_token_boundaries() {
        let tokens = [
            "Here",
            " <thi",
            "nk>my ",
            "reason",
            "ing</thi",
            "nk> is the ",
            "answer",
        ];
        let mut f = ThinkFilter::new();
        let (mut clean, mut reasoning) = (String::new(), String::new());
        for t in tokens {
            let (c, r) = f.feed_split(t);
            clean.push_str(&c);
            reasoning.push_str(&r);
        }
        clean.push_str(&f.finish());
        assert_eq!(clean, "Here  is the answer");
        assert_eq!(reasoning, "my reasoning");
        // `feed` still suppresses (delegates to feed_split's clean half).
        let mut g = ThinkFilter::new();
        assert_eq!(g.feed("a<think>b"), "a");
        assert_eq!(g.feed("c</think>d"), "d");
    }

    #[test]
    fn streaming_unterminated_block_is_dropped() {
        assert_eq!(
            stream("", &["answer ", "<think>still ", "thinking"]),
            "answer "
        );
    }

    /// Drive the leading-reasoning filter; returns `(clean, reasoning)`.
    fn stream_leading(tokens: &[&str]) -> (String, String) {
        let mut f = ThinkFilter::with_leading_reasoning();
        let (mut clean, mut reasoning) = (String::new(), String::new());
        for t in tokens {
            let (c, r) = f.feed_split(t);
            clean.push_str(&c);
            reasoning.push_str(&r);
        }
        clean.push_str(&f.finish());
        (clean, reasoning)
    }

    #[test]
    fn leading_reasoning_strips_lone_closer() {
        // The Nemotron shape: reasoning, a bare `</think>`, then the answer —
        // no opening `<think>`. Matches the batch `lone_leading_closer_is_reasoning`.
        let (clean, reasoning) = stream_leading(&["reasoning", "</think>", "the answer"]);
        assert_eq!(clean, "the answer");
        assert_eq!(reasoning, "reasoning");
    }

    #[test]
    fn leading_reasoning_handles_closer_split_across_tokens() {
        let (clean, reasoning) = stream_leading(&["reason", "ing</thi", "nk>the ", "answer"]);
        assert_eq!(clean, "the answer");
        assert_eq!(reasoning, "reasoning");
    }

    #[test]
    fn leading_reasoning_without_closer_flushes_as_clean() {
        // A flagged model that emits no `</think>` (e.g. thinking turned off)
        // must still get its full reply — provisional reasoning is flushed as
        // clean, never lost.
        let (clean, reasoning) = stream_leading(&["a plain ", "reply with ", "no tags"]);
        assert_eq!(clean, "a plain reply with no tags");
        assert_eq!(reasoning, "");
    }

    #[test]
    fn leading_reasoning_tolerates_a_paired_block() {
        // Even if a flagged model does emit an opener, the literal tag is folded
        // into the (discarded) reasoning and the answer stays clean.
        let (clean, _r) = stream_leading(&["<think>r</think>", "answer"]);
        assert_eq!(clean, "answer");
    }

    #[test]
    fn emits_leading_reasoning_matches_known_families() {
        assert!(emits_leading_reasoning("nemotron-3-nano:30b"));
        assert!(emits_leading_reasoning("nemotron3:33b"));
        assert!(emits_leading_reasoning("deepseek-r1:7b"));
        assert!(emits_leading_reasoning("qwen3:8b"));
        assert!(!emits_leading_reasoning("llama3:8b"));
        assert!(!emits_leading_reasoning("qwen2.5-coder:7b"));
        assert!(!emits_leading_reasoning("gpt-4o"));
    }
}