diff --git a/nanobot/utils/helpers.py b/nanobot/utils/helpers.py
index c545e9ca..32824aaf 100644
--- a/nanobot/utils/helpers.py
+++ b/nanobot/utils/helpers.py
@@ -65,6 +65,27 @@ def _estimate_tools_tokens(
return token_count
+def _tag_regex(tags: tuple[str, ...]) -> str:
+ return rf"(?:{'|'.join(re.escape(tag) for tag in tags)})"
+
+
+_THINKING_TAGS = ("think", "thinking", "thought")
+_INLINE_SELF_CLOSING_PRESERVED_TAGS = frozenset({"think", "thought"})
+_INLINE_SELF_CLOSING_THINKING_TAGS = tuple(
+ tag for tag in _THINKING_TAGS if tag not in _INLINE_SELF_CLOSING_PRESERVED_TAGS
+)
+_THINKING_TAG = _tag_regex(_THINKING_TAGS)
+_INLINE_SELF_CLOSING_THINKING_TAG = _tag_regex(_INLINE_SELF_CLOSING_THINKING_TAGS)
+_THINKING_TAG_PREFIX = "|".join(
+ sorted(
+ {re.escape(tag[:i]) for tag in _THINKING_TAGS for i in range(1, len(tag) + 1)},
+ key=len,
+ reverse=True,
+ )
+)
+_PARTIAL_THINKING_TAG = rf"?(?:{_THINKING_TAG_PREFIX})>?"
+
+
def strip_think(text: str) -> str:
"""Remove thinking blocks, unclosed trailing tags, and tokenizer-level
template leaks occasionally emitted by some models (notably Gemma 4's
@@ -92,38 +113,28 @@ def strip_think(text: str) -> str:
assistant discusses the tokens themselves.
"""
# Well-formed blocks first.
- text = re.sub(r"[\s\S]*?", "", text)
- text = re.sub(r"^\s*[\s\S]*$", "", text)
- text = re.sub(r"[\s\S]*?", "", text)
- text = re.sub(r"^\s*[\s\S]*$", "", text)
- text = re.sub(r"[\s\S]*?", "", text)
- text = re.sub(r"^\s*[\s\S]*$", "", text)
+ text = re.sub(rf"<(?P{_THINKING_TAG})>[\s\S]*?(?P=tag)>", "", text)
+ text = re.sub(rf"^\s*<{_THINKING_TAG}>[\s\S]*$", "", text)
# Self-closing `` is an empty marker, not user-visible text.
- text = re.sub(r"^\s*\s*", "", text)
- text = re.sub(r"\s*\s*$", "", text)
+ text = re.sub(rf"^\s*<{_INLINE_SELF_CLOSING_THINKING_TAG}/>\s*", "", text)
+ text = re.sub(rf"\s*<{_INLINE_SELF_CLOSING_THINKING_TAG}/>\s*$", "", text)
# Malformed opening tags: `` / `/` — we can't use `\w` here because in Python's default
# Unicode regex mode it matches CJK characters too, which would defeat
# the primary fix for `/])", "", text)
- text = re.sub(r"/])", "", text)
- text = re.sub(r"/])", "", text)
+ text = re.sub(rf"<{_THINKING_TAG}(?![A-Za-z0-9_\-:>/])", "", text)
# Edge-only orphan closing tags (start or end of text).
- text = re.sub(r"^\s*\s*", "", text)
- text = re.sub(r"\s*\s*$", "", text)
- text = re.sub(r"^\s*\s*", "", text)
- text = re.sub(r"\s*\s*$", "", text)
- text = re.sub(r"^\s*\s*", "", text)
- text = re.sub(r"\s*\s*$", "", text)
+ text = re.sub(rf"^\s*{_THINKING_TAG}>\s*", "", text)
+ text = re.sub(rf"\s*{_THINKING_TAG}>\s*$", "", text)
# Edge-only channel markers (harmony / Gemma 4 variant leaks).
text = re.sub(r"^\s*<\|?channel\|?>\s*", "", text)
# Stream chunks may end in the middle of a control tag. Strip only known
# control-token prefixes at the very end.
partial_control_tag = (
- r"?(?:t|th|thi|thin|think|thinki|thinkin|thinking|tho|thou|thoug|though|thought)>?"
- r"|<\|?(?:c|ch|cha|chan|chann|channe|channel)(?:\|?>?)?"
+ rf"{_PARTIAL_THINKING_TAG}|"
+ r"<\|?(?:c|ch|cha|chan|chann|channe|channel)(?:\|?>?)?"
)
text = re.sub(rf"(?:{partial_control_tag})$", "", text)
text = re.sub(r"^\s*<\|?$", "", text)
@@ -134,15 +145,12 @@ def strip_reasoning_tags(text: object) -> str:
"""Remove wrapper tags from text that is already known to be reasoning."""
if not isinstance(text, str):
return ""
- partial_reasoning_tag = (
- r"?(?:t|th|thi|thin|think|thinki|thinkin|thinking|tho|thou|thoug|though|thought)>?"
- )
- text = re.sub(rf"^\s*(?:{partial_reasoning_tag})$", "", text)
- text = re.sub(r"^\s*<(?:think|thinking|thought)/>\s*", "", text)
- text = re.sub(r"\s*<(?:think|thinking|thought)/>\s*$", "", text)
- text = re.sub(r"^\s*<(?:think|thinking|thought)>\s*", "", text)
- text = re.sub(r"\s*(?:think|thinking|thought)>\s*$", "", text)
- text = re.sub(rf"\s*(?:{partial_reasoning_tag})$", "", text)
+ text = re.sub(rf"^\s*(?:{_PARTIAL_THINKING_TAG})$", "", text)
+ text = re.sub(rf"^\s*<{_THINKING_TAG}/>\s*", "", text)
+ text = re.sub(rf"\s*<{_THINKING_TAG}/>\s*$", "", text)
+ text = re.sub(rf"^\s*<{_THINKING_TAG}>\s*", "", text)
+ text = re.sub(rf"\s*{_THINKING_TAG}>\s*$", "", text)
+ text = re.sub(rf"\s*(?:{_PARTIAL_THINKING_TAG})$", "", text)
return text.strip()
@@ -154,12 +162,8 @@ def extract_think(text: str) -> tuple[str | None, str]:
text but not surfaced — :func:`strip_think` handles that case.
"""
parts: list[str] = []
- for m in re.finditer(r"([\s\S]*?)", text):
- parts.append(m.group(1).strip())
- for m in re.finditer(r"([\s\S]*?)", text):
- parts.append(m.group(1).strip())
- for m in re.finditer(r"([\s\S]*?)", text):
- parts.append(m.group(1).strip())
+ for m in re.finditer(rf"<(?P{_THINKING_TAG})>([\s\S]*?)(?P=tag)>", text):
+ parts.append(m.group(2).strip())
thinking = "\n\n".join(parts) if parts else None
return thinking, strip_think(text)