diff --git a/tensorrt_llm/serve/tool_parser/base_tool_parser.py b/tensorrt_llm/serve/tool_parser/base_tool_parser.py index ece736749d2f..0300d46f0382 100644 --- a/tensorrt_llm/serve/tool_parser/base_tool_parser.py +++ b/tensorrt_llm/serve/tool_parser/base_tool_parser.py @@ -110,6 +110,33 @@ def _ends_with_partial_token(self, buffer: str, bot_token: str) -> int: return i return 0 + def _starts_with_leftover_eot_token(self, buffer: str) -> bool: + r"""Check if the buffer opens with the eot_token of a finished call. + + Completing a tool call leaves everything the parser did not consume in + the buffer, starting with that call's eot_token. When the eot_token + itself begins with the tool_call_separator, as in Qwen3 where calls are + separated by "\n" and closed by "\n", that leftover looks + exactly like the separator that introduces the next tool call. + """ + return bool(self.eot_token) and buffer.startswith(self.eot_token) + + def _may_begin_tool_call(self, text: str) -> bool: + """Check whether text could be the opening of a tool call. + + Tells a tool_call_separator that introduces another call apart from one + that merely precedes ordinary prose. A call opens either with the + bot_token or, in the formats this base class streams, with bare JSON. + Text too short to judge counts as a maybe, so the buffer keeps growing + until the answer is certain. + """ + if not text: + return True + if self.bot_token and (text.startswith(self.bot_token) + or self.bot_token.startswith(text)): + return True + return text[0] in "{[" + def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: """ @@ -127,15 +154,65 @@ def parse_streaming_increment(self, new_text: str, For incompatible formats, detectors should override this method with custom logic. """ + pending = self._buffer + new_text + name_sent = self.current_tool_name_sent + result = self._parse_increment_once(new_text, tools) + + # A pass stops at the end of one tool call and leaves the rest of the + # increment in the buffer for the increment after it. Nothing drains + # the buffer once the stream ends, so whatever arrived in the same + # chunk as the closing markup -- trailing content, a further tool call, + # or the arguments of the call that just opened -- would never be + # emitted. Keep parsing while a pass still moves the parser forward. + # + # Forward means the buffer shrank, or the pass sent a tool name and so + # the next one will stream that call's arguments. A pass that does + # neither has nothing left to give, and repeating it would re-parse the + # same bytes on every token of the stream. + while self._buffer and (self._buffer != pending or + (self.current_tool_name_sent + and not name_sent)): + pending = self._buffer + name_sent = self.current_tool_name_sent + step = self._parse_increment_once("", tools) + result.normal_text += step.normal_text + result.calls.extend(step.calls) + + return result + + def _parse_increment_once(self, new_text: str, + tools: List[Tool]) -> StreamingParseResult: + """Run a single parsing pass over the buffer plus the new text.""" # Append new text to buffer self._buffer += new_text + + # Parsing a tool call stops at its closing markup, leaving the + # eot_token at the head of the buffer. Drop it before looking at what + # follows: it is markup rather than content, and while it is still + # there the checks below read it as part of the next tool call. + if self.current_tool_id > 0 and self.eot_token: + if self._starts_with_leftover_eot_token(self._buffer): + self._buffer = self._buffer[len(self.eot_token):] + elif self.eot_token.startswith(self._buffer): + # The end token is still arriving one token at a time. Hold it + # so its opening bytes are not mistaken for content. + return StreamingParseResult() + current_text = self._buffer # The current_text has tool_call if it is the start of a new tool call sequence - # or it is the start of a new tool call after a tool call separator, when there is a previous tool call - if not (self.has_tool_call(current_text) or - (self.current_tool_id > 0 - and current_text.startswith(self.tool_call_separator))): + # or it is the start of a new tool call after a tool call separator, when there is a previous tool call. + # The separator only introduces a call when a call actually follows it; + # a response that resumes with prose on a new line starts the same way, + # and reading that as a call leaves it stuck in the tool call branch + # below, where prose never parses as JSON. + starts_next_tool_call = ( + self.current_tool_id > 0 + and current_text.startswith(self.tool_call_separator) + and self._may_begin_tool_call( + current_text[len(self.tool_call_separator):])) + + if not (self.has_tool_call(current_text) or starts_next_tool_call): # Only clear buffer if we're sure no tool call is starting if not self._ends_with_partial_token(self._buffer, self.bot_token): normal_text = self._buffer @@ -158,8 +235,7 @@ def parse_streaming_increment(self, new_text: str, tool_call_pos = current_text.find(self.bot_token) if tool_call_pos != -1: start_idx = tool_call_pos + len(self.bot_token) - elif self.current_tool_id > 0 and current_text.startswith( - self.tool_call_separator): + elif starts_next_tool_call: start_idx = len(self.tool_call_separator) else: start_idx = 0 @@ -222,6 +298,9 @@ def parse_streaming_increment(self, new_text: str, else: cur_arguments = current_tool_call.get("arguments") res = StreamingParseResult() + argument_diff = None + # Save the ID of the tool that's completing + completing_tool_id = self.current_tool_id if cur_arguments: # Calculate how much of the arguments we've already streamed @@ -233,24 +312,9 @@ def parse_streaming_increment(self, new_text: str, prev_arguments = self.prev_tool_call_arr[ self.current_tool_id].get("arguments") - argument_diff = None - # If the current tool's JSON is complete, send all remaining arguments if is_current_complete: argument_diff = cur_args_json[sent:] - completing_tool_id = ( - self.current_tool_id - ) # Save the ID of the tool that's completing - - # Only remove the processed portion, keep unprocessed content - self._buffer = current_text[start_idx + end_idx:] - - if self.current_tool_id < len(self.prev_tool_call_arr): - self.prev_tool_call_arr[ - self.current_tool_id].clear() - self.current_tool_name_sent = False - self.streamed_args_for_tool[self.current_tool_id] = "" - self.current_tool_id += 1 # If the tool is still being parsed, send incremental changes elif prev_arguments: @@ -260,21 +324,36 @@ def parse_streaming_increment(self, new_text: str, cur_args_json) argument_diff = prefix[sent:] - # Send the argument diff if there's something new - if argument_diff is not None: - # Use the correct tool_index: completing_tool_id for completed tools, current_tool_id for ongoing - tool_index_to_use = (completing_tool_id - if is_current_complete else - self.current_tool_id) - res = StreamingParseResult(calls=[ - ToolCallItem( - tool_index=tool_index_to_use, - parameters=argument_diff, - ) - ], ) - if not is_current_complete: - self.streamed_args_for_tool[ - self.current_tool_id] += argument_diff + # Close the call out whenever its JSON is complete, not only + # when it carried arguments. A call invoked with none still + # ends here, and leaving it open keeps its markup and + # everything after it stuck in the buffer for the rest of the + # stream. + if is_current_complete: + # Only remove the processed portion, keep unprocessed content + self._buffer = current_text[start_idx + end_idx:] + + if self.current_tool_id < len(self.prev_tool_call_arr): + self.prev_tool_call_arr[self.current_tool_id].clear() + self.current_tool_name_sent = False + self.streamed_args_for_tool[self.current_tool_id] = "" + self.current_tool_id += 1 + + # Send the argument diff if there's something new + if argument_diff is not None: + # Use the correct tool_index: completing_tool_id for completed tools, current_tool_id for ongoing + tool_index_to_use = (completing_tool_id + if is_current_complete else + self.current_tool_id) + res = StreamingParseResult(calls=[ + ToolCallItem( + tool_index=tool_index_to_use, + parameters=argument_diff, + ) + ], ) + if not is_current_complete: + self.streamed_args_for_tool[ + self.current_tool_id] += argument_diff # Update prev_tool_call_arr with current state if self.current_tool_id >= 0: diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index 7b24d4554337..05e2944d4ab9 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -97,6 +97,20 @@ def sample_tools(): ] +def _arguments_by_tool_index(calls) -> dict: + """Reassemble each tool call's streamed arguments, keyed by tool index. + + Arguments reach the client as a series of fragments carrying no name, so a + test that wants to check them has to join the fragments per tool index. + """ + joined: dict = {} + for call in calls: + if call.parameters: + joined[call.tool_index] = joined.get(call.tool_index, + "") + call.parameters + return {index: json.loads(text) for index, text in joined.items()} + + # Concrete implementation of BaseToolParser for testing class ConcreteToolParser(BaseToolParser): """Concrete implementation of BaseToolParser for testing abstract methods.""" @@ -312,9 +326,14 @@ def test_parse_streaming_increment_complete_tool(self, sample_tools): '[TOOL_CALLS] {"name":"get_weather","arguments":{"location":"Boston"}}', sample_tools) - # Should have sent tool name (first call) - assert len(result.calls) == 1 + # The whole call arrived at once, so the name and the arguments both + # go out here. Keeping the arguments back for a later increment drops + # them whenever this chunk turns out to be the last one. + assert len(result.calls) == 2 assert result.calls[0].name == "get_weather" + assert result.calls[0].parameters == "" + assert json.loads(result.calls[1].parameters) == {"location": "Boston"} + assert parser._buffer == "" def test_parse_streaming_increment_invalid_tool_name(self, sample_tools): """Test streaming parser handles invalid tool name.""" @@ -765,6 +784,187 @@ def test_parse_streaming_increment_multiple_tools_streaming( assert result.calls[0].parameters == "" assert result.calls[0].tool_index == 1 + # Issue #17740: the eot_token "\n" starts with the + # tool_call_separator "\n", so the leftover buffer of a finished call used + # to be read as the separator announcing the next one. The parser then + # stayed in the tool call branch, where the closing markup never parses as + # JSON, and silently swallowed the rest of the response. + def test_streaming_emits_content_after_completed_tool_call( + self, sample_tools, parser): + """Text streamed after a finished tool call reaches the client.""" + chunks = [ + "\n", + '{"name":"get_weather","arguments":{"location":"NYC"}}', + "\n", + " It is sunny.", + ] + normal_text = "".join( + parser.parse_streaming_increment(chunk, sample_tools).normal_text + for chunk in chunks) + + assert normal_text == " It is sunny." + assert parser._buffer == "" + + def test_streaming_content_after_tool_call_with_split_end_token( + self, sample_tools, parser): + """The end token may arrive across chunks without trapping the buffer.""" + chunks = [ + "\n", + '{"name":"get_weather","arguments":{"location":"NYC"}}', + "\n", + " It is sunny.", + ] + normal_text = "".join( + parser.parse_streaming_increment(chunk, sample_tools).normal_text + for chunk in chunks) + + assert normal_text == " It is sunny." + assert parser._buffer == "" + + def test_streaming_content_after_tool_call_character_by_character( + self, sample_tools, parser): + """Token-by-token streaming keeps the markup out of the content.""" + response = ('\n' + '{"name":"get_weather","arguments":{"location":"NYC"}}\n' + ' It is sunny.') + normal_text = "".join( + parser.parse_streaming_increment(char, sample_tools).normal_text + for char in response) + + assert normal_text == " It is sunny." + assert parser._buffer == "" + + def test_streaming_content_in_same_chunk_as_end_token( + self, sample_tools, parser): + """Content sharing a chunk with the closing markup is still emitted. + + Nothing drains the parser buffer once the stream ends, so a tool call + that completes on the final chunk has to release the rest of that + chunk right away. + """ + chunks = [ + "\n", + '{"name":"get_weather","arguments":{"location":"NYC"}}', + "\n It is sunny.", + ] + normal_text = "".join( + parser.parse_streaming_increment(chunk, sample_tools).normal_text + for chunk in chunks) + + assert normal_text == " It is sunny." + assert parser._buffer == "" + + def test_streaming_whole_response_in_one_chunk(self, sample_tools, parser): + """A single chunk carrying the entire response loses nothing.""" + result = parser.parse_streaming_increment( + '\n{"name":"get_weather","arguments":{"location":"NYC"}}' + '\n It is sunny.', sample_tools) + + assert [call.name for call in result.calls + if call.name] == ["get_weather"] + assert json.loads("".join(call.parameters for call in result.calls + if call.parameters)) == { + "location": "NYC" + } + assert result.normal_text == " It is sunny." + assert parser._buffer == "" + + def test_streaming_two_tool_calls_in_one_chunk(self, sample_tools, parser): + """A second call sharing the chunk with the first is not swallowed.""" + result = parser.parse_streaming_increment( + '\n{"name":"get_weather","arguments":{"location":"NYC"}}' + '\n\n' + '\n{"name":"search_web","arguments":{"query":"AI"}}' + '\n All set.', sample_tools) + + assert [call.name for call in result.calls + if call.name] == ["get_weather", "search_web"] + assert _arguments_by_tool_index(result.calls) == { + 0: { + "location": "NYC" + }, + 1: { + "query": "AI" + }, + } + assert result.normal_text == " All set." + assert parser._buffer == "" + + def test_streaming_content_after_tool_call_on_its_own_line( + self, sample_tools, parser): + """Prose resuming on a new line is content, not another tool call. + + It opens with the tool_call_separator just as a following call would, + so the parser has to look at what comes after the separator. + """ + chunks = [ + "\n", + '{"name":"get_weather","arguments":{"location":"NYC"}}', + "\n", + "\nIt is sunny.", + ] + normal_text = "".join( + parser.parse_streaming_increment(chunk, sample_tools).normal_text + for chunk in chunks) + + assert normal_text == "\nIt is sunny." + assert parser._buffer == "" + + def test_streaming_content_after_zero_argument_tool_call( + self, sample_tools, parser): + """A call invoked with no arguments still closes out. + + The completion bookkeeping used to sit behind a check for arguments to + stream, so a call with none never released the buffer and took the rest + of the response down with it. + """ + chunks = [ + "\n", + '{"name":"get_weather","arguments":{}}', + "\n", + " Nothing to report.", + ] + normal_text = "".join( + parser.parse_streaming_increment(chunk, sample_tools).normal_text + for chunk in chunks) + + assert normal_text == " Nothing to report." + assert parser._buffer == "" + + def test_streaming_content_after_multiple_tool_calls( + self, sample_tools, parser): + """A separator between two calls is still honored before the content.""" + chunks = [ + "\n", + '{"name":"get_weather","arguments":{"location":"NYC"}}', + "\n\n", + "\n", + '{"name":"search_web","arguments":{"query":"AI"}}', + "\n", + " All set.", + ] + normal_text = "" + names = [] + calls = [] + for chunk in chunks: + result = parser.parse_streaming_increment(chunk, sample_tools) + normal_text += result.normal_text + names += [call.name for call in result.calls if call.name] + calls += result.calls + + assert names == ["get_weather", "search_web"] + assert _arguments_by_tool_index(calls) == { + 0: { + "location": "NYC" + }, + 1: { + "query": "AI" + }, + } + assert normal_text == " All set." + assert parser._buffer == "" + def test_structure_info_function(self): """Test structure_info returns correct lambda function.""" parser = Qwen3ToolParser()