Fixing nested JSON args parsing for tool-calls in streaming

nemotron_toolcall_parser_streaming.py

@@ -2,7 +2,7 @@ import json
 from collections.abc import Sequence
 from random import choices
 from string import ascii_letters, digits
-from typing import Union
+from typing import Optional, Union
 
 import partial_json_parser
 import regex as re
@@ -16,8 +16,6 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
                                               FunctionCall, ToolCall)
 from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
     ToolParser, ToolParserManager)
-from vllm.entrypoints.openai.tool_parsers.utils import (
-    extract_intermediate_diff)
 from vllm.logger import init_logger
 from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
 
@@ -61,6 +59,7 @@ class NemotronToolParser(ToolParser):
         self.current_tool_name_sent: bool = False
         self.streamed_args_for_tool: list[str] = [
         ]  # map what has been streamed for each tool so far to a list
+        self.tool_args_emitted: list[bool] = []
         self.bot_token = "<TOOLCALL>"
         self.bot_token_id = self.vocab.get(self.bot_token)
         logger.info(f"Nemotron Tool Parser: bot_token: {self.bot_token}, bot_token_id: {self.bot_token_id}")
@@ -75,6 +74,116 @@ class NemotronToolParser(ToolParser):
         # a forthcoming <TOOLCALL> or </TOOLCALL> tag in streaming.
         self._pending_tag_buffer: str = ""
 
+    @staticmethod
+    def _strip_trailing_auto_closers(chunk: str) -> str:
+        """
+        Remove parser auto-completed closing braces/brackets plus trailing whitespace.
+        These should be flushed only when a tool call completes to avoid duplicate
+        argument fragments.
+        """
+        idx = len(chunk)
+        while idx > 0 and chunk[idx - 1] in " \t\r\n}]":
+            idx -= 1
+        # Remove trailing non-escaped double quotes (partial JSON auto-closes strings)
+        while idx > 0 and chunk[idx - 1] == '"':
+            # keep escaped quotes (\"), only strip bare ones
+            if idx - 2 >= 0 and chunk[idx - 2] == '\\':
+                break
+            idx -= 1
+        return chunk[:idx]
+
+    @staticmethod
+    def _common_prefix_len(left: str, right: str) -> int:
+        """
+        Return the length of the shared prefix between left and right strings.
+        """
+        max_len = min(len(left), len(right))
+        idx = 0
+        while idx < max_len and left[idx] == right[idx]:
+            idx += 1
+        return idx
+
+    def _compute_arguments_delta(self, cur_arguments_json: str,
+                                 end_of_call: bool) -> str:
+        """
+        Determine the incremental suffix to stream for the current tool call.
+        Ensures we only emit monotonic chunks by trimming our tracked prefix to
+        the longest common prefix with the latest JSON snapshot.
+        """
+        tool_idx = self.current_tool_id
+        if tool_idx < 0 or tool_idx >= len(self.streamed_args_for_tool):
+            return ""
+
+        streamed_prefix = self.streamed_args_for_tool[tool_idx]
+        had_any = (self.tool_args_emitted[tool_idx]
+                   if tool_idx < len(self.tool_args_emitted) else False)
+
+        lcp_len = self._common_prefix_len(cur_arguments_json,
+                                          streamed_prefix)
+        if lcp_len != len(streamed_prefix):
+            streamed_prefix = streamed_prefix[:lcp_len]
+            self.streamed_args_for_tool[tool_idx] = streamed_prefix
+
+        if (not had_any and not end_of_call and lcp_len == 0
+                and cur_arguments_json.endswith('": ""}')
+                and '": ""' in cur_arguments_json):
+            closing_pos = cur_arguments_json.rfind('": ""}')
+            if closing_pos != -1:
+                arguments_delta = cur_arguments_json[:closing_pos + 4]
+            else:
+                arguments_delta = cur_arguments_json
+        else:
+            arguments_delta = cur_arguments_json[lcp_len:]
+
+        if not arguments_delta:
+            return ""
+
+        if not end_of_call:
+            arguments_delta = self._strip_trailing_auto_closers(
+                arguments_delta)
+
+        if (not had_any and not end_of_call and arguments_delta
+                and arguments_delta.endswith('}')):
+            arguments_delta = arguments_delta[:-1]
+            if arguments_delta.endswith('"'):
+                arguments_delta = arguments_delta[:-1]
+
+        return arguments_delta
+
+    def _visible_delta_outside_tool(self, delta_text: str,
+                                    start_token: Optional[str],
+                                    end_token: Optional[str]) -> str:
+        """
+        Consume characters that could begin a tool tag. Only suppress the exact
+        <TOOLCALL> / </TOOLCALL> sequences, and let everything else (e.g. </think>)
+        pass through untouched.
+        """
+        if not delta_text:
+            return delta_text
+
+        visible: list[str] = []
+        for ch in delta_text:
+            if self._pending_tag_buffer or ch == '<':
+                self._pending_tag_buffer += ch
+
+                if start_token and start_token.startswith(self._pending_tag_buffer):
+                    if self._pending_tag_buffer == start_token:
+                        self._pending_tag_buffer = ""
+                    continue
+
+                if end_token and end_token.startswith(self._pending_tag_buffer):
+                    if self._pending_tag_buffer == end_token:
+                        self._pending_tag_buffer = ""
+                    continue
+
+                # Not a tool tag; flush buffered characters as normal content.
+                visible.append(self._pending_tag_buffer)
+                self._pending_tag_buffer = ""
+            else:
+                visible.append(ch)
+
+        return "".join(visible)
+
     def adjust_request(
             self, request: ChatCompletionRequest) -> ChatCompletionRequest:
         if not isinstance(
@@ -175,64 +284,13 @@ class NemotronToolParser(ToolParser):
         # if candidates tool call tokens are in the tokens generated so far, that
         # means we're parsing as tool calls now. Suppress streaming if we are
         # currently generating any prefix of the start or end tag.
+        visible_delta_text = delta_text
         try:
             start_token = self.bot_token
             end_token = f"</{self.bot_token[1:]}" if self.bot_token.startswith('<') else None
 
-            # Handle potential start of tool call tags by buffering partial sequences
-            if delta_text == '<' and not self._pending_tag_buffer:
-                # Start buffering a potential tag
-                self._pending_tag_buffer = '<'
-                return None
-            
-            # If we have a pending tag buffer, accumulate and decide
-            if self._pending_tag_buffer:
-                # Accumulate the current token into the buffer
-                self._pending_tag_buffer += delta_text
-                
-                # Extract just the alphabetic part after '<'
-                alpha_part = ""
-                for i in range(1, len(self._pending_tag_buffer)):
-                    if self._pending_tag_buffer[i].isalpha():
-                        alpha_part += self._pending_tag_buffer[i].upper()
-                    else:
-                        break
-                
-                
-                # Check if we have a complete opening tag '<TOOLCALL>'
-                if '<TOOLCALL>' in self._pending_tag_buffer:
-                    # We have the complete opening tag - stop buffering and let normal processing take over
-                    buffered_content = self._pending_tag_buffer
-                    self._pending_tag_buffer = ""
-                    
-                    # Update the text variables to include the buffered content
-                    updated_current_text = previous_text + buffered_content
-                    updated_delta_text = buffered_content  # The entire buffered content is the delta
-                    
-                    # Continue processing with the complete tool call content
-                    current_text = updated_current_text
-                    delta_text = updated_delta_text
-                    # Fall through to normal processing
-                elif self._pending_tag_buffer.startswith('</'):
-                    # End tag pattern - keep buffering until we see if it's a valid end tag
-                    return None
-                elif alpha_part and "TOOLCALL".startswith(alpha_part) and len(alpha_part) < 8:
-                    # Could be building to TOOLCALL and haven't completed it yet - keep buffering
-                    return None
-                elif len(alpha_part) > 0 and not "TOOLCALL".startswith(alpha_part):
-                    # Alphabetic content that definitely won't become TOOLCALL - flush as content
-                    content_to_flush = self._pending_tag_buffer
-                    self._pending_tag_buffer = ""
-                    return DeltaMessage(content=content_to_flush)
-                else:
-                    # Keep buffering - not enough info yet
-                    return None
-
-            # Suppress ANY partial prefix of the start/end tag to avoid leaking tag characters.
-            if any(current_text.endswith(start_token[:k]) for k in range(1, len(start_token))):
-                return None
-            if end_token and any(current_text.endswith(end_token[:k]) for k in range(1, len(end_token))):
-                return None
+            visible_delta_text = self._visible_delta_outside_tool(
+                delta_text, start_token, end_token)
         except Exception:
             # Fallback to conservative checks in case of any issues
             if current_text.endswith('<') or current_text.endswith('<T') or current_text.endswith('<TO') or current_text.endswith('<TOOL') or current_text.endswith('<TOOLCALL'):
@@ -241,12 +299,10 @@ class NemotronToolParser(ToolParser):
         # if the tool call token is not in the tokens generated so far, append
         # output to contents since it's not a tool
         if self.bot_token not in current_text:
-            # If we were buffering a partial tag and reached here, flush it first.
-            if self._pending_tag_buffer:
-                content_to_flush = self._pending_tag_buffer + delta_text
-                self._pending_tag_buffer = ""
-                return DeltaMessage(content=content_to_flush)
-            return DeltaMessage(content=delta_text)
+            if visible_delta_text:
+                return DeltaMessage(content=visible_delta_text)
+            # still waiting on a potential tag, so emit nothing yet
+            return None
 
         # bit mask flags for partial JSON parsing. If the name hasn't been
         # sent yet, don't allow sending
@@ -272,7 +328,8 @@ class NemotronToolParser(ToolParser):
             try:
                 tool_call_arr: list[dict] = partial_json_parser.loads(
                     parsable_arr, flags)
-            except partial_json_parser.core.exceptions.MalformedJSON:
+            except (partial_json_parser.core.exceptions.MalformedJSON,
+                    json.JSONDecodeError, ValueError):
                 return None
 
             current_tool_call: dict = tool_call_arr[self.current_tool_id] \
@@ -315,6 +372,7 @@ class NemotronToolParser(ToolParser):
                 self.current_tool_id = len(tool_call_arr) - 1
                 self.current_tool_name_sent = False
                 self.streamed_args_for_tool.append("")
+                self.tool_args_emitted.append(False)
                 return delta
 
             # case: update an existing tool - this is handled below
@@ -345,10 +403,6 @@ class NemotronToolParser(ToolParser):
                     self.current_tool_id].get("arguments")
                 cur_arguments = current_tool_call.get("arguments")
 
-                new_text = delta_text.replace("\'", "\"")
-                if ('"}' in new_text):
-                    new_text = new_text[:new_text.rindex('"}')]
-
                 if not cur_arguments and not prev_arguments:
 
                     delta = None
@@ -357,47 +411,11 @@ class NemotronToolParser(ToolParser):
                         "INVARIANT - impossible to have arguments reset "
                         "mid-arguments")
                     delta = None
-                elif cur_arguments and not prev_arguments:
+                elif cur_arguments:
                     cur_arguments_json = json.dumps(cur_arguments,
                                                     ensure_ascii=False)
-                    streamed_prefix = self.streamed_args_for_tool[
-                        self.current_tool_id]
-                    
-                    # The issue: partial JSON parser auto-completes incomplete strings
-                    # e.g., {"location": "  becomes {"location": ""} in parsed result
-                    # We need to handle this by detecting when the parsed result has auto-completed empty strings
-                    
-                    # Check if this looks like an auto-completed partial string
-                    if (cur_arguments_json.endswith('": ""}') and 
-                        not streamed_prefix and 
-                        '": ""' in cur_arguments_json):
-                        # This is likely auto-completed - remove the auto-completed empty string
-                        # e.g., {"location": ""} -> {"location": "
-                        closing_pos = cur_arguments_json.rfind('": ""}')
-                        if closing_pos != -1:
-                            arguments_delta = cur_arguments_json[:closing_pos + 4]  # Keep up to ": "
-                        else:
-                            arguments_delta = cur_arguments_json
-                    else:
-                        # Normal case - use diff calculation
-                        if cur_arguments_json.startswith(streamed_prefix):
-                            arguments_delta = cur_arguments_json[len(streamed_prefix):]
-                        else:
-                            # Fallback: compute diff when prefix does not match.
-                            arguments_delta = extract_intermediate_diff(
-                                cur_arguments_json, streamed_prefix)
-
-                    # Do not include a trailing '}' in the very first
-                    # arguments chunk; defer it to the end-of-call flush to
-                    # avoid prematurely closing the JSON object.
-                    if (not self.streamed_args_for_tool[self.current_tool_id]
-                            and not end_of_call and arguments_delta
-                            and arguments_delta.endswith('}')):
-                        arguments_delta = arguments_delta[:-1]
-                        # if there is an auto-completed closing quote '"' before the }, strip it too
-                        # e.g., {"color_hex": "#"} -> {"color_hex": "#"} -> {"color_hex": "#"}
-                        if arguments_delta.endswith('"'):
-                            arguments_delta = arguments_delta[:-1]
+                    arguments_delta = self._compute_arguments_delta(
+                        cur_arguments_json, end_of_call)
                     if arguments_delta:
                         delta = DeltaMessage(tool_calls=[
                             DeltaToolCall(index=self.current_tool_id,
@@ -407,26 +425,8 @@ class NemotronToolParser(ToolParser):
                         ])
                         self.streamed_args_for_tool[
                             self.current_tool_id] += arguments_delta
-                    else:
-                        delta = None
-
-                elif cur_arguments and prev_arguments:
-                    cur_args_json = json.dumps(cur_arguments,
-                                               ensure_ascii=False)
-                    prev_args_json = json.dumps(prev_arguments,
-                                                ensure_ascii=False)
-
-                    argument_diff = extract_intermediate_diff(
-                        cur_args_json, prev_args_json)
-                    if argument_diff:
-                        delta = DeltaMessage(tool_calls=[
-                            DeltaToolCall(index=self.current_tool_id,
-                                          function=DeltaFunctionCall(
-                                              arguments=argument_diff).model_dump(
-                                                  exclude_none=True))
-                        ])
-                        self.streamed_args_for_tool[
-                            self.current_tool_id] += argument_diff
+                        self.tool_args_emitted[
+                            self.current_tool_id] = True
                     else:
                         # Do not flush final JSON here; let the serving layer
                         # compute a minimal remaining suffix on finish.
@@ -447,19 +447,12 @@ class NemotronToolParser(ToolParser):
                     if cur_arguments is not None:
                         cur_args_json = json.dumps(cur_arguments,
                                                    ensure_ascii=False)
-                        streamed_prefix = self.streamed_args_for_tool[
-                            self.current_tool_id]
-                        
-                        if cur_args_json.startswith(streamed_prefix):
-                            remaining_suffix = cur_args_json[len(
-                                streamed_prefix):]
-                        else:
-                            remaining_suffix = extract_intermediate_diff(
-                                cur_args_json, streamed_prefix)
+                        remaining_suffix = self._compute_arguments_delta(
+                            cur_args_json, end_of_call=True)
 
                         # Only send remaining suffix if it's non-empty and contains meaningful content
                         # (not just whitespace or single characters like closing braces)
-                        if remaining_suffix and remaining_suffix.strip() and len(remaining_suffix.strip()) > 0:
+                        if remaining_suffix and remaining_suffix.strip():
                             extra = DeltaToolCall(
                                 index=self.current_tool_id,
                                 function=DeltaFunctionCall(
@@ -474,6 +467,7 @@ class NemotronToolParser(ToolParser):
                                     delta.tool_calls = [extra]
                             self.streamed_args_for_tool[
                                 self.current_tool_id] += remaining_suffix
+                            self.tool_args_emitted[self.current_tool_id] = True
                         else:
                             pass
                 except Exception: