diff --git a/tensorrt_llm/serve/tool_parser/base_tool_parser.py b/tensorrt_llm/serve/tool_parser/base_tool_parser.py index ece736749d2f..5fac1ee979fb 100644 --- a/tensorrt_llm/serve/tool_parser/base_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/base_tool_parser.py @@ -1,7 +1,7 @@ # Adapted from https://github.com/sgl-project/sglang/blob/083629c23564e1a64deaa052f1df5c5d914358d8/python/sglang/srt/function_call/base_format_detector.py import json from abc import ABC, abstractmethod -from typing import Any, Dict, List +from typing import Any, Dict, List, Tuple from partial_json_parser.core.exceptions import MalformedJSON from partial_json_parser.core.options import Allow @@ -110,6 +110,34 @@ def _ends_with_partial_token(self, buffer: str, bot_token: str) -> int: return i return 0 + def _split_leading_normal_text(self, buffer: str, + start_tokens: List[str]) -> Tuple[str, str]: + """Split a buffer at the earliest tool call start token. + + A streaming increment can carry ordinary content and the opening of a + tool call at once. That content precedes the tool call and is not part + of it, so it has to be returned to the client as normal text instead of + being consumed along with the markup. + + The split is verbatim: no whitespace is trimmed. Trimming here would + make the streamed content depend on where increment boundaries happen + to fall, since a prefix that arrives in its own increment is emitted by + the no-tool-call path, which does not trim either. Non-streaming + ``detect_and_parse`` strips the prefix in most parsers, so accumulated + streaming content can carry whitespace that the one-shot path drops. + + Returns a ``(normal_text, remainder)`` pair. ``normal_text`` is empty + when the buffer starts with a tool call or contains none of the tokens. + """ + indices = [ + index for index in (buffer.find(token) for token in start_tokens) + if index != -1 + ] + if not indices: + return "", buffer + split_at = min(indices) + return buffer[:split_at], buffer[split_at:] + def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -151,6 +179,14 @@ def parse_streaming_increment(self, new_text: str, if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) + # This increment may hold ordinary text ahead of the tool call. Hand + # that text back as content and keep only the markup in the buffer, + # otherwise it is consumed along with the tool call and never reaches + # the client. + normal_text, current_text = self._split_leading_normal_text( + current_text, [self.bot_token]) + self._buffer = current_text + flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR try: @@ -165,7 +201,7 @@ def parse_streaming_increment(self, new_text: str, start_idx = 0 if start_idx >= len(current_text): - return StreamingParseResult() + return StreamingParseResult(normal_text=normal_text) (obj, end_idx) = partial_json_loads(current_text[start_idx:], flags) @@ -183,10 +219,10 @@ def parse_streaming_increment(self, new_text: str, current_tool_call = obj except MalformedJSON: - return StreamingParseResult() + return StreamingParseResult(normal_text=normal_text) if not current_tool_call: - return StreamingParseResult() + return StreamingParseResult(normal_text=normal_text) # Case 1: Handle tool name streaming # This happens when we encounter a tool but haven't sent its name yet @@ -284,11 +320,12 @@ def parse_streaming_increment(self, new_text: str, self.prev_tool_call_arr[ self.current_tool_id] = current_tool_call + res.normal_text = normal_text + res.normal_text return res except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult() + return StreamingParseResult(normal_text=normal_text) @abstractmethod def has_tool_call(self, text: str) -> bool: diff --git a/tensorrt_llm/serve/tool_parser/deepseekv31_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv31_parser.py index a0b306a1060e..ea0761f18770 100644 --- a/tensorrt_llm/serve/tool_parser/deepseekv31_parser.py +++ b/tensorrt_llm/serve/tool_parser/deepseekv31_parser.py @@ -112,6 +112,14 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) + # This increment may hold ordinary text ahead of the tool call. Hand that + # text back as content and keep only the markup in the buffer, otherwise + # it is consumed along with the tool call and never reaches the client. + normal_text, current_text = self._split_leading_normal_text( + current_text, [self.bot_token, "<｜tool▁call▁begin｜>"] + ) + self._buffer = current_text + calls: list[ToolCallItem] = [] try: partial_match = re.search( @@ -183,17 +191,17 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami else: self._buffer = "" - result = StreamingParseResult(normal_text="", calls=calls) + result = StreamingParseResult(normal_text=normal_text, calls=calls) self.current_tool_id += 1 self._last_arguments = "" self.current_tool_name_sent = False return result - return StreamingParseResult(normal_text="", calls=calls) + return StreamingParseResult(normal_text=normal_text, calls=calls) except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=normal_text + current_text) def structure_info(self) -> _GetInfoFunc: return lambda name: StructureInfo( diff --git a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py index 3362ab7cb68a..d2160bf33c2c 100644 --- a/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py +++ b/tensorrt_llm/serve/tool_parser/deepseekv32_parser.py @@ -191,6 +191,14 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) + # This increment may hold ordinary text ahead of the tool call. Hand that + # text back as content and keep only the markup in the buffer, otherwise + # it is consumed along with the tool call and never reaches the client. + normal_text, current_text = self._split_leading_normal_text( + current_text, [self.bot_token, "<｜DSML｜invoke"] + ) + self._buffer = current_text + all_calls: list[ToolCallItem] = [] try: # Loop to handle multiple consecutive invoke blocks @@ -288,11 +296,11 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami break # No more invoke blocks found - return StreamingParseResult(normal_text="", calls=all_calls) + return StreamingParseResult(normal_text=normal_text, calls=all_calls) except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=normal_text + current_text) def structure_info(self) -> _GetInfoFunc: return lambda name: StructureInfo( diff --git a/tensorrt_llm/serve/tool_parser/deepseekv3_parser.py b/tensorrt_llm/serve/tool_parser/deepseekv3_parser.py index 8eb49eb81c07..a04c02acb126 100644 --- a/tensorrt_llm/serve/tool_parser/deepseekv3_parser.py +++ b/tensorrt_llm/serve/tool_parser/deepseekv3_parser.py @@ -115,6 +115,14 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) + # This increment may hold ordinary text ahead of the tool call. Hand that + # text back as content and keep only the markup in the buffer, otherwise + # it is consumed along with the tool call and never reaches the client. + normal_text, current_text = self._split_leading_normal_text( + current_text, [self.bot_token, "<｜tool▁call▁begin｜>"] + ) + self._buffer = current_text + calls: list[ToolCallItem] = [] try: partial_match = re.search( @@ -187,17 +195,17 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami else: self._buffer = "" - result = StreamingParseResult(normal_text="", calls=calls) + result = StreamingParseResult(normal_text=normal_text, calls=calls) self.current_tool_id += 1 self._last_arguments = "" self.current_tool_name_sent = False return result - return StreamingParseResult(normal_text="", calls=calls) + return StreamingParseResult(normal_text=normal_text, calls=calls) except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=normal_text + current_text) def structure_info(self) -> _GetInfoFunc: return lambda name: StructureInfo( diff --git a/tensorrt_llm/serve/tool_parser/glm4_parser.py b/tensorrt_llm/serve/tool_parser/glm4_parser.py index 540f3d842309..1a3807e92283 100644 --- a/tensorrt_llm/serve/tool_parser/glm4_parser.py +++ b/tensorrt_llm/serve/tool_parser/glm4_parser.py @@ -338,6 +338,12 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) + # This increment may hold ordinary text ahead of the tool call. Hand that + # text back as content and keep only the markup in the buffer, otherwise + # it is consumed along with the tool call and never reaches the client. + normal_text, current_text = self._split_leading_normal_text(current_text, [self.bot_token]) + self._buffer = current_text + calls: list[ToolCallItem] = [] try: partial_match = re.search( @@ -351,7 +357,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami is_tool_end = partial_match.group(3) if func_name_raw is None or not func_name_raw.strip(): - return StreamingParseResult(normal_text="", calls=[]) + return StreamingParseResult(normal_text=normal_text, calls=[]) func_name = func_name_raw.strip() func_args_raw = func_args_raw.strip() if func_args_raw else "" @@ -442,7 +448,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami self._buffer = current_text[partial_match.end(3) :] - result = StreamingParseResult(normal_text="", calls=calls) + result = StreamingParseResult(normal_text=normal_text, calls=calls) self.current_tool_id += 1 self._last_arguments = "" self.current_tool_name_sent = False @@ -450,11 +456,11 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami self._reset_streaming_state() return result - return StreamingParseResult(normal_text="", calls=calls) + return StreamingParseResult(normal_text=normal_text, calls=calls) except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=normal_text + current_text) def _parse_argument_pairs( self, pairs: List[Tuple[str, str]], func_name: str, tools: List[Tool] diff --git a/tensorrt_llm/serve/tool_parser/kimi_k2_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k2_tool_parser.py index ca2d0a7d7d65..442895a7eaad 100644 --- a/tensorrt_llm/serve/tool_parser/kimi_k2_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/kimi_k2_tool_parser.py @@ -118,6 +118,14 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami if not hasattr(self, "_tool_indices"): self._tool_indices = self._get_tool_indices(tools) + # This increment may hold ordinary text ahead of the tool call. Hand that + # text back as content and keep only the markup in the buffer, otherwise + # it is consumed along with the tool call and never reaches the client. + normal_text, current_text = self._split_leading_normal_text( + current_text, [self.bot_token, self.tool_call_start_token] + ) + self._buffer = current_text + calls: list[ToolCallItem] = [] try: match = self.stream_tool_call_portion_regex.search(current_text) @@ -128,7 +136,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami m = self.tool_call_id_regex.match(function_id) if not m: logger.warning("Unexpected tool_call_id format: %s", function_id) - return StreamingParseResult(normal_text="", calls=calls) + return StreamingParseResult(normal_text=normal_text, calls=calls) function_name = m.group("name") # Initialize state if this is the first tool call @@ -191,7 +199,7 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami else: self._buffer = "" - result = StreamingParseResult(normal_text="", calls=calls) + result = StreamingParseResult(normal_text=normal_text, calls=calls) self.current_tool_id += 1 self._last_arguments = "" self.current_tool_name_sent = False @@ -199,11 +207,11 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami except json.JSONDecodeError: pass - return StreamingParseResult(normal_text="", calls=calls) + return StreamingParseResult(normal_text=normal_text, calls=calls) except Exception as e: logger.error(f"Error in parse_streaming_increment: {e}") - return StreamingParseResult(normal_text=current_text) + return StreamingParseResult(normal_text=normal_text + current_text) def structure_info(self) -> _GetInfoFunc: """Return function that creates StructureInfo for guided generation.""" diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 7b24d4554337..0c1b47bee965 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -4355,3 +4355,154 @@ def test_other_raw_token_parser_keeps_spacing_contract(self) -> None: assert sampling_params.skip_special_tokens is False assert sampling_params.spaces_between_special_tokens is True + + +# ============================================================================ +# Streaming: content sharing an increment with a tool call +# ============================================================================ + +_LEADING_TEXT = "Let me check the weather for you. " + +# Parser class plus a complete tool call in that parser's own format. Every +# tool call starts with the parser's bot_token, so tests can split it there to +# get a delta boundary that never lands inside a token. +_LEADING_TEXT_PARSERS = { + "deepseek_v3": ( + DeepSeekV3Parser, + ("<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n" + '```json\n{"location": "NYC"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>' + ), + ), + "deepseek_v31": ( + DeepSeekV31Parser, + ("<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>get_weather<｜tool▁sep｜>" + '{"location": "NYC"}<｜tool▁call▁end｜><｜tool▁calls▁end｜>'), + ), + "deepseek_v32": ( + DeepSeekV32Parser, + ('<｜DSML｜function_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + " "), + ), + "deepseek_v4": ( + DeepSeekV4Parser, + ('<｜DSML｜tool_calls> <｜DSML｜invoke name="get_weather"> ' + '<｜DSML｜parameter name="location" string="true">NYC ' + " "), + ), + "glm4": ( + Glm4ToolParser, + ("get_weather\nlocation\n" + "NYC\n"), + ), + # Glm47 already preserved the leading text before this fix; it is kept here + # as a guard against regressing it. + "glm47": ( + Glm47ToolParser, + ("get_weather\nlocation\n" + "NYC\n"), + ), + "qwen3": ( + Qwen3ToolParser, + ('\n{"name": "get_weather", "arguments": ' + '{"location": "NYC"}}\n'), + ), + "kimi_k2": ( + KimiK2ToolParser, + ("<|tool_calls_section_begin|><|tool_call_begin|>functions.get_weather:0" + '<|tool_call_argument_begin|>{"location": "NYC"}<|tool_call_end|>' + "<|tool_calls_section_end|>"), + ), +} + + +@pytest.mark.parametrize("parser_cls, tool_call", + list(_LEADING_TEXT_PARSERS.values()), + ids=list(_LEADING_TEXT_PARSERS)) +class TestStreamingLeadingText: + """Content that precedes a tool call must survive streaming. + + Regression tests for issue #17580: the streaming path returned an empty + normal_text as soon as a tool call opened, so anything the model said + before calling the tool was consumed along with the markup and never + reached the client, even though detect_and_parse returns it for the same + input. + """ + + @staticmethod + def _feed(parser, deltas, tools): + """Stream deltas through the parser and accumulate what it emits. + + A trailing empty increment drains the arguments: parsers send the tool + name on the increment that completes the call and the arguments on the + one after it. + """ + normal_text = "" + names = [] + arguments = "" + for delta in [*deltas, ""]: + result = parser.parse_streaming_increment(delta, tools) + normal_text += result.normal_text + for call in result.calls: + if call.name: + names.append(call.name) + arguments += call.parameters + return normal_text, names, arguments + + def _assert_call_intact(self, names, arguments): + assert names == ["get_weather"] + assert json.loads(arguments) == {"location": "NYC"} + + def test_leading_text_in_same_increment(self, sample_tools, parser_cls, + tool_call): + parser = parser_cls() + normal_text, names, arguments = self._feed(parser, + [_LEADING_TEXT + tool_call], + sample_tools) + + assert normal_text == _LEADING_TEXT + self._assert_call_intact(names, arguments) + + def test_leading_text_emitted_once_when_call_spans_increments( + self, sample_tools, parser_cls, tool_call): + parser = parser_cls() + split_at = len(parser.bot_token) + deltas = [_LEADING_TEXT + tool_call[:split_at], tool_call[split_at:]] + + normal_text, names, arguments = self._feed(parser, deltas, sample_tools) + + assert normal_text == _LEADING_TEXT + self._assert_call_intact(names, arguments) + + def test_tool_call_without_leading_text_emits_no_content( + self, sample_tools, parser_cls, tool_call): + parser = parser_cls() + normal_text, names, arguments = self._feed(parser, [tool_call], + sample_tools) + + assert normal_text == "" + self._assert_call_intact(names, arguments) + + def test_content_does_not_depend_on_increment_boundaries( + self, sample_tools, parser_cls, tool_call): + """What the client sees must not change with the delta boundaries. + + This is why the leading segment is emitted verbatim rather than + stripped: a prefix arriving in its own increment leaves the parser + before any tool call is known, so trimming at the split point would + make the streamed content depend on how the text happened to be + chunked. + """ + split_at = len(parser_cls().bot_token) + chunkings = ( + [_LEADING_TEXT + tool_call], + [_LEADING_TEXT, tool_call], + [_LEADING_TEXT + tool_call[:split_at], tool_call[split_at:]], + ) + + emitted = { + self._feed(parser_cls(), deltas, sample_tools)[0] + for deltas in chunkings + } + + assert emitted == {_LEADING_TEXT}