fix(utils): anchor unclosed think-tag regex to string start (#3004)

This commit is contained in:
04cb 2026-04-11 00:47:23 +08:00 committed by Xubin Ren
parent 9bccfa63d2
commit e392c27f7e
2 changed files with 31 additions and 2 deletions

View File

@ -17,10 +17,10 @@ from loguru import logger
def strip_think(text: str) -> str:
"""Remove thinking blocks and any unclosed trailing tag."""
text = re.sub(r"<think>[\s\S]*?</think>", "", text)
text = re.sub(r"<think>[\s\S]*$", "", text)
text = re.sub(r"^\s*<think>[\s\S]*$", "", text)
# Gemma 4 and similar models use <thought>...</thought> blocks
text = re.sub(r"<thought>[\s\S]*?</thought>", "", text)
text = re.sub(r"<thought>[\s\S]*$", "", text)
text = re.sub(r"^\s*<thought>[\s\S]*$", "", text)
return text.strip()

View File

@ -34,3 +34,32 @@ class TestStripThinkTag:
def test_empty_string(self):
assert strip_think("") == ""
class TestStripThinkFalsePositive:
"""Ensure mid-content <think>/<thought> tags are NOT stripped (#3004)."""
def test_backtick_think_tag_preserved(self):
text = "*Think Stripping:* A new utility to strip `<think>` tags from output."
assert strip_think(text) == text
def test_prose_think_tag_preserved(self):
text = "The model emits <think> at the start of its response."
assert strip_think(text) == text
def test_code_block_think_tag_preserved(self):
text = "Example:\n```\ntext = re.sub(r\"<think>[\\s\\S]*\", \"\", text)\n```\nDone."
assert strip_think(text) == text
def test_backtick_thought_tag_preserved(self):
text = "Gemma 4 uses `<thought>` blocks for reasoning."
assert strip_think(text) == text
def test_prefix_unclosed_think_still_stripped(self):
assert strip_think("<think>reasoning without closing") == ""
def test_prefix_unclosed_think_with_whitespace(self):
assert strip_think(" <think>reasoning...") == ""
def test_prefix_unclosed_thought_still_stripped(self):
assert strip_think("<thought>reasoning without closing") == ""