diff --git a/tensorrt_llm/llmapi/reasoning_parser.py b/tensorrt_llm/llmapi/reasoning_parser.py index ccc81e2e05e2..c04456bdf6d3 100644 --- a/tensorrt_llm/llmapi/reasoning_parser.py +++ b/tensorrt_llm/llmapi/reasoning_parser.py @@ -210,6 +210,31 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult: raise RuntimeError( "Unreachable code reached in `DeepSeekR1Parser.parse_delta`") + def finish(self) -> ReasoningParserResult: + """Flush text withheld by `parse_delta` when the stream ends. + + `parse_delta` holds back a trailing fragment that could still grow + into a `` / `` tag. If the stream ends while such a + fragment is buffered - because the response was truncated mid-tag, or + simply ends with a literal `<` - the fragment is ordinary model output + and must still be emitted, otherwise it is silently dropped. The + buffered text is attributed to the block it was withheld in: inside a + reasoning block it is reasoning content, otherwise it is visible + content. + + A buffer holding exactly a complete tag is a delimiter rather than + model output, so it is discarded as it would have been had more text + followed. + """ + remaining = self._buffer + self._buffer = "" + if not remaining or remaining in (self.reasoning_start, + self.reasoning_end): + return ReasoningParserResult() + if self.in_reasoning: + return ReasoningParserResult(reasoning_content=remaining) + return ReasoningParserResult(content=remaining) + @register_reasoning_parser("deepseek_v4") class DeepSeekV4ReasoningParser(BaseReasoningParser): diff --git a/tests/unittest/llmapi/test_reasoning_parser.py b/tests/unittest/llmapi/test_reasoning_parser.py index 11c70bf2b334..b2a8f78457a5 100644 --- a/tests/unittest/llmapi/test_reasoning_parser.py +++ b/tests/unittest/llmapi/test_reasoning_parser.py @@ -62,6 +62,79 @@ def test_deepseek_r1_reasoning_parser_stream(delta_texts: list, content: list, assert result.reasoning_content == reasoning_context[i] +@pytest.mark.parametrize( + ("parser_key", "text"), + [ + # `finish()` flushes as reasoning: the stream starts inside the + # reasoning block, so the withheld `<` is reasoning output. + ("deepseek-r1", "a <"), + # `finish()` flushes as content: no reasoning block was entered. + ("qwen3", "a<"), + # `finish()` discards: the buffer holds exactly a delimiter. Passes on + # `main` too - it guards against a fix that leaks the tag instead. + ("deepseek-r1", f"a{R1_END}"), + ]) +def test_deepseek_r1_reasoning_parser_stream_matches_non_stream( + parser_key: str, text: str) -> None: + """Streaming char-by-char then finishing must match a non-streaming parse. + + One `(parser_key, text)` pair per branch of `finish()`. This is the + contract the missing flush violated, so it subsumes example-based tests + of the individual branches. + """ + expected = ReasoningParserFactory.create_reasoning_parser(parser_key).parse( + text) + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + parser_key) + content, reasoning_context = "", "" + for char in text: + result = reasoning_parser.parse_delta(char) + content += result.content + reasoning_context += result.reasoning_content + result = reasoning_parser.finish() + content += result.content + reasoning_context += result.reasoning_content + assert content == expected.content + assert reasoning_context == expected.reasoning_content + + +def test_deepseek_r1_reasoning_parser_finish_flushes_partial_tag() -> None: + """A partial tag arriving alongside text must still be flushed. + + Such a delta fills `_buffer` through the `rfind` branch of `parse_delta`, + which one-character-at-a-time streaming never reaches - and it is the + shape a real stream delivers. + """ + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + "deepseek-r1") + assert reasoning_parser.parse_delta("a None: + """`finish()` must reach whichever parser the thinking flag selected. + + `DeepSeekV4ReasoningParser` delegates to two different targets: + `DeepSeekR1Parser`, which now flushes, and `IdentityReasoningParser`, + which withholds nothing to flush. + """ + reasoning_parser = ReasoningParserFactory.create_reasoning_parser( + "deepseek_v4", chat_template_kwargs) + reasoning_parser.parse_delta("a <") + result = reasoning_parser.finish() + assert result.reasoning_content == flushed + assert result.content == "" + + @pytest.mark.parametrize("chat_template_kwargs", [{ "thinking": True }, {