[Kimi K3] Rework skipped-think fix as opt-in force_nonempty_content with streaming coverage (#34187)
Co-authored-by: yhyang201 <yhyang201@gmail.com> Co-authored-by: Yuhao Yang <47235274+yhyang201@users.noreply.github.com>
This commit is contained in:
co-authored by
yhyang201
Yuhao Yang
parent
a6001478f4
commit
27b7a2dc3b
@@ -549,6 +549,7 @@ class KimiK3Detector(BaseReasoningFormatDetector):
|
|||||||
force_reasoning: bool = True,
|
force_reasoning: bool = True,
|
||||||
continue_final_message: bool = False,
|
continue_final_message: bool = False,
|
||||||
previous_content: str = "",
|
previous_content: str = "",
|
||||||
|
force_nonempty_content: bool = False,
|
||||||
):
|
):
|
||||||
# strict-thinking flattens these to single token ids, so the full marker
|
# strict-thinking flattens these to single token ids, so the full marker
|
||||||
# "<|open|>response<|sep|>" is inexpressible. The bare name works: it
|
# "<|open|>response<|sep|>" is inexpressible. The bare name works: it
|
||||||
@@ -574,8 +575,12 @@ class KimiK3Detector(BaseReasoningFormatDetector):
|
|||||||
previous_content=previous_content,
|
previous_content=previous_content,
|
||||||
reasoning_default="thinking",
|
reasoning_default="thinking",
|
||||||
)
|
)
|
||||||
|
# Unlike the base class, K3 cannot use `normal_text == ""` alone:
|
||||||
|
# skipped-think and truncated marker-free reasoning end up identical.
|
||||||
|
self._force_nonempty_content = force_nonempty_content
|
||||||
self._reasoning_done = False
|
self._reasoning_done = False
|
||||||
self._tools_passthrough = False
|
self._tools_passthrough = False
|
||||||
|
self._stream_text = ""
|
||||||
|
|
||||||
def _clean_content(self, text: str) -> str:
|
def _clean_content(self, text: str) -> str:
|
||||||
tools_idx = text.find(TOOLS_OPEN)
|
tools_idx = text.find(TOOLS_OPEN)
|
||||||
@@ -595,6 +600,8 @@ class KimiK3Detector(BaseReasoningFormatDetector):
|
|||||||
in_reasoning = self._in_reasoning or self.think_start_token in text
|
in_reasoning = self._in_reasoning or self.think_start_token in text
|
||||||
if not in_reasoning and self.think_end_token not in text:
|
if not in_reasoning and self.think_end_token not in text:
|
||||||
return StreamingParseResult(normal_text=self._clean_content(text))
|
return StreamingParseResult(normal_text=self._clean_content(text))
|
||||||
|
if self._force_nonempty_content and self._is_skipped_think_answer(text):
|
||||||
|
return StreamingParseResult(normal_text=self._clean_content(text))
|
||||||
|
|
||||||
open_idx = text.find(self.think_start_token)
|
open_idx = text.find(self.think_start_token)
|
||||||
start = open_idx + len(self.think_start_token) if open_idx != -1 else 0
|
start = open_idx + len(self.think_start_token) if open_idx != -1 else 0
|
||||||
@@ -622,8 +629,19 @@ class KimiK3Detector(BaseReasoningFormatDetector):
|
|||||||
reasoning_text=reasoning_text, normal_text=self._clean_content(rest)
|
reasoning_text=reasoning_text, normal_text=self._clean_content(rest)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _is_skipped_think_answer(self, text: str) -> bool:
|
||||||
|
return (
|
||||||
|
self.think_start_token not in text
|
||||||
|
and self.think_end_token not in text
|
||||||
|
and self.think_start_token.removesuffix("<|sep|>") not in text
|
||||||
|
and self.think_end_token.removesuffix("<|sep|>") not in text
|
||||||
|
and (RESPONSE_CLOSE in text or MESSAGE_CLOSE in text)
|
||||||
|
)
|
||||||
|
|
||||||
def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
|
def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
|
||||||
self._buffer += new_text
|
self._buffer += new_text
|
||||||
|
if self._force_nonempty_content:
|
||||||
|
self._stream_text += new_text
|
||||||
|
|
||||||
if not self._in_reasoning and not self._reasoning_done:
|
if not self._in_reasoning and not self._reasoning_done:
|
||||||
open_idx = self._buffer.find(self.think_start_token)
|
open_idx = self._buffer.find(self.think_start_token)
|
||||||
@@ -684,6 +702,23 @@ class KimiK3Detector(BaseReasoningFormatDetector):
|
|||||||
|
|
||||||
return StreamingParseResult(normal_text=self._drain_content())
|
return StreamingParseResult(normal_text=self._drain_content())
|
||||||
|
|
||||||
|
def finish(self) -> StreamingParseResult:
|
||||||
|
if not self._force_nonempty_content:
|
||||||
|
return super().finish()
|
||||||
|
text, self._stream_text = self._stream_text, ""
|
||||||
|
if self._in_reasoning and self._is_skipped_think_answer(text):
|
||||||
|
# _in_reasoning means no channel decision happened mid-stream, so the
|
||||||
|
# answer went out as reasoning; without this gate the re-emit duplicates
|
||||||
|
# answers already streamed as content (RESPONSE_OPEN / force_reasoning=False).
|
||||||
|
self._buffer = ""
|
||||||
|
return StreamingParseResult(normal_text=self._clean_content(text))
|
||||||
|
if self._in_reasoning and not self.stream_reasoning and self._buffer:
|
||||||
|
# super().finish() would emit this buffer as content under
|
||||||
|
# force_nonempty_content — the leak the flag exists to prevent.
|
||||||
|
buffer, self._buffer = self._buffer, ""
|
||||||
|
return StreamingParseResult(reasoning_text=buffer)
|
||||||
|
return StreamingParseResult()
|
||||||
|
|
||||||
def _drain_content(self) -> str:
|
def _drain_content(self) -> str:
|
||||||
buf = self._buffer
|
buf = self._buffer
|
||||||
if not buf:
|
if not buf:
|
||||||
|
|||||||
@@ -190,5 +190,114 @@ def test_reasoning_parser_registration() -> None:
|
|||||||
assert isinstance(ReasoningParser("kimi_k3").detector, KimiK3Detector)
|
assert isinstance(ReasoningParser("kimi_k3").detector, KimiK3Detector)
|
||||||
|
|
||||||
|
|
||||||
|
def _stream_with_finish(detector: KimiK3Detector, chunks: list[str]) -> tuple[str, str]:
|
||||||
|
reasoning, content = _stream(detector, chunks)
|
||||||
|
result = detector.finish()
|
||||||
|
return reasoning + result.reasoning_text, content + result.normal_text
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("text", "reasoning", "content"),
|
||||||
|
[
|
||||||
|
(
|
||||||
|
f"bare answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}",
|
||||||
|
"",
|
||||||
|
"bare answer",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
f"{RESPONSE_OPEN}bare answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}",
|
||||||
|
"",
|
||||||
|
"bare answer",
|
||||||
|
),
|
||||||
|
("still going", "still going", ""),
|
||||||
|
("deep thought<|close|>", "deep thought", ""),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_fnc_non_stream_skipped_think_vs_truncated_reasoning(
|
||||||
|
text: str, reasoning: str, content: str
|
||||||
|
) -> None:
|
||||||
|
detector = KimiK3Detector(force_reasoning=True, force_nonempty_content=True)
|
||||||
|
result = detector.detect_and_parse(text)
|
||||||
|
assert result.reasoning_text == reasoning
|
||||||
|
assert result.normal_text == content
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("chunk_size", [1, 5, 13])
|
||||||
|
def test_fnc_streaming_skipped_think_answer(chunk_size: int) -> None:
|
||||||
|
detector = KimiK3Detector(force_reasoning=True, force_nonempty_content=True)
|
||||||
|
text = f"bare answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}"
|
||||||
|
reasoning, content = _stream_with_finish(detector, _chunks(text, chunk_size))
|
||||||
|
# Streamed as reasoning in real time; finish() re-emits the cleaned
|
||||||
|
# payload as content once the channel close proves skipped-think.
|
||||||
|
assert reasoning == text
|
||||||
|
assert content == "bare answer"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("chunk_size", [1, 5, 13])
|
||||||
|
def test_fnc_streaming_truncated_reasoning_stays_reasoning(chunk_size: int) -> None:
|
||||||
|
detector = KimiK3Detector(force_reasoning=True, force_nonempty_content=True)
|
||||||
|
reasoning, content = _stream_with_finish(
|
||||||
|
detector, _chunks("still going", chunk_size)
|
||||||
|
)
|
||||||
|
assert reasoning == "still going"
|
||||||
|
assert content == ""
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("chunk_size", [5, 13])
|
||||||
|
def test_fnc_streaming_long_think_streams_without_close(chunk_size: int) -> None:
|
||||||
|
detector = KimiK3Detector(force_reasoning=True, force_nonempty_content=True)
|
||||||
|
text = "x" * 20000
|
||||||
|
reasoning = ""
|
||||||
|
for chunk in _chunks(text, chunk_size):
|
||||||
|
# Live chunks flow immediately — no hold-back to starve SSE idle timeouts.
|
||||||
|
reasoning += detector.parse_streaming_increment(chunk).reasoning_text
|
||||||
|
assert reasoning == text
|
||||||
|
result = detector.finish()
|
||||||
|
assert result.reasoning_text == ""
|
||||||
|
assert result.normal_text == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_fnc_streaming_response_open_not_reemitted() -> None:
|
||||||
|
detector = KimiK3Detector(force_reasoning=True, force_nonempty_content=True)
|
||||||
|
text = f"{RESPONSE_OPEN}bare answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}"
|
||||||
|
reasoning, content = _stream_with_finish(detector, _chunks(text, 5))
|
||||||
|
# The channel switch already streamed the answer as content; finish()
|
||||||
|
# must not re-emit it.
|
||||||
|
assert reasoning == ""
|
||||||
|
assert content == "bare answer"
|
||||||
|
|
||||||
|
|
||||||
|
def test_fnc_streaming_force_reasoning_off_not_reemitted() -> None:
|
||||||
|
detector = KimiK3Detector(force_reasoning=False, force_nonempty_content=True)
|
||||||
|
text = f"bare answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}"
|
||||||
|
reasoning, content = _stream_with_finish(detector, _chunks(text, 5))
|
||||||
|
assert reasoning == ""
|
||||||
|
assert content == "bare answer"
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("force_nonempty_content", [False, True])
|
||||||
|
def test_stream_reasoning_off_truncation_flushes_reasoning(
|
||||||
|
force_nonempty_content: bool,
|
||||||
|
) -> None:
|
||||||
|
detector = KimiK3Detector(
|
||||||
|
force_reasoning=True,
|
||||||
|
stream_reasoning=False,
|
||||||
|
force_nonempty_content=force_nonempty_content,
|
||||||
|
)
|
||||||
|
reasoning, content = _stream_with_finish(detector, _chunks("still going", 5))
|
||||||
|
assert reasoning == "still going"
|
||||||
|
assert content == ""
|
||||||
|
|
||||||
|
|
||||||
|
def test_fnc_stream_reasoning_off_skipped_think_reemits_content() -> None:
|
||||||
|
detector = KimiK3Detector(
|
||||||
|
force_reasoning=True, stream_reasoning=False, force_nonempty_content=True
|
||||||
|
)
|
||||||
|
text = f"bare answer{RESPONSE_CLOSE}{MESSAGE_CLOSE}"
|
||||||
|
reasoning, content = _stream_with_finish(detector, _chunks(text, 5))
|
||||||
|
assert reasoning == ""
|
||||||
|
assert content == "bare answer"
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
sys.exit(pytest.main([__file__]))
|
sys.exit(pytest.main([__file__]))
|
||||||
|
|||||||
Reference in New Issue
Block a user