Updating streaming tool-call parser to return ChoiceDeltaToolCall (#31)

- Updating streaming tool-call parser to return ChoiceDeltaToolCall (ee35ca89f5df3773dc1e1e131000b8a4a138a3c9)


Co-authored-by: Ameya Sunil Mahabaleshwarkar <ameyasunilm@users.noreply.huggingface.co>

nemotron_toolcall_parser_streaming.py

@@ -1,123 +1,166 @@
 import json
-import re
 from collections.abc import Sequence
-from typing import Union, Optional
+from random import choices
+from string import ascii_letters, digits
+from typing import Union
 
 import partial_json_parser
+import regex as re
+from partial_json_parser.core.options import Allow
+from pydantic import Field
 
-from vllm.entrypoints.openai.protocol import (
-    ChatCompletionRequest,
-    DeltaFunctionCall,
-    DeltaMessage,
-    DeltaToolCall,
-    ExtractedToolCallInformation,
-    FunctionCall,
-    ToolCall,
-)
+from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
+                                              DeltaFunctionCall, DeltaMessage,
+                                              DeltaToolCall,
+                                              ExtractedToolCallInformation,
+                                              FunctionCall, ToolCall)
 from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (
-    ToolParser,
-    ToolParserManager,
-)
+    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
-from vllm.utils import random_uuid
+from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
 
 logger = init_logger(__name__)
 
+ALPHANUMERIC = ascii_letters + digits
+
+
+class NemotronToolCall(ToolCall):
+    id: str = Field(
+        default_factory=lambda: NemotronToolCall.generate_random_id())
+
+    @staticmethod
+    def generate_random_id():
+        return "".join(choices(ALPHANUMERIC, k=9))
+
+    @staticmethod
+    def is_valid_id(id: str) -> bool:
+        return id.isalnum() and len(id) == 9
+
+
+def _is_fn_name_regex_support(model_tokenizer: AnyTokenizer) -> bool:
+    return isinstance(model_tokenizer, MistralTokenizer) \
+        and model_tokenizer.version >= 11
+
 
 @ToolParserManager.register_module("nemotron_json")
-class NemotronJSONToolParser(ToolParser):
+class NemotronToolParser(ToolParser):
+    """
+    Tool call parser for Nemotron-Nano-V2
+
+    Used when --enable-auto-tool-choice --tool-call-parser nemotron_json are all set
+    """
 
     def __init__(self, tokenizer: AnyTokenizer):
         super().__init__(tokenizer)
-
-        # Streaming state tracking
-        self.current_tool_name_sent: bool = False
+        # initialize properties used for state when parsing tool calls in
+        # streaming mode
         self.prev_tool_call_arr: list[dict] = []
         self.current_tool_id: int = -1
-        self.streamed_args_for_tool: list[str] = []
-        self.tool_call_ids: list[str] = []  # Track IDs for each tool call
-        
-        # Track what we've sent so far in streaming
-        self.sent_tool_calls_count: int = 0
-        self.sent_args_length: dict[int, int] = {}  # tool_idx -> length of args sent
+        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.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}")
+        self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL)
+        if _is_fn_name_regex_support(self.model_tokenizer):
+            self.fn_name_regex = re.compile(
+                r'([a-zA-Z0-9_-]+)(\{[\s\S]*?\})(?=\s*$|,|\s)', re.DOTALL)
+        else:
+            self.fn_name_regex = None
 
-        self.tool_call_start_token: str = "<TOOLCALL>"
-        self.tool_call_end_token: str = "</TOOLCALL>"
+        # Buffer for partial tag sequences to disambiguate between normal content and
+        # a forthcoming <TOOLCALL> or </TOOLCALL> tag in streaming.
+        self._pending_tag_buffer: str = ""
 
-        self.tool_call_regex = re.compile(r"<TOOLCALL>(.*?)</TOOLCALL>", re.DOTALL)
+    def adjust_request(
+            self, request: ChatCompletionRequest) -> ChatCompletionRequest:
+        if not isinstance(
+                self.model_tokenizer, MistralTokenizer
+        ) and request.tools and request.tool_choice != 'none':
+            # Do not skip special tokens when using chat template
+            # with Mistral parser as TOOL_CALL token is needed
+            # for tool detection.
+            # Note: we don't want skip_special_tokens=False
+            # with MistralTokenizer as it is incompatible
+            request.skip_special_tokens = False
+        return request
 
     def extract_tool_calls(
         self,
         model_output: str,
         request: ChatCompletionRequest,
     ) -> ExtractedToolCallInformation:
-        """Extract tool calls from non-streaming (complete) output."""
-        
-        if self.tool_call_start_token not in model_output:
-            return ExtractedToolCallInformation(
-                tools_called=False,
-                tool_calls=[],
-                content=model_output,
-            )
+        """
+        Extract the tool calls from a complete model response. Requires
+        find-and-replacing single quotes with double quotes for JSON parsing,
+        make sure your tool call arguments don't ever include quotes!
+        """
+
+        # case -- if a tool call token is not present, return a text response
+        if self.bot_token not in model_output:
+            return ExtractedToolCallInformation(tools_called=False,
+                                                tool_calls=[],
+                                                content=model_output)
+
+        # first remove the BOT token
+        tool_content = model_output.replace(self.bot_token, "").strip()
 
         try:
-            # Try to extract complete <TOOLCALL>...</TOOLCALL> blocks
-            tool_call_matches = self.tool_call_regex.findall(model_output)
-            
-            if tool_call_matches:
-                # Complete tool call block found
-                str_tool_calls = tool_call_matches[0].strip()
-            else:
-                # Incomplete - extract everything after <TOOLCALL>
-                start_idx = model_output.find(self.tool_call_start_token) + len(self.tool_call_start_token)
-                str_tool_calls = model_output[start_idx:].strip()
-            
-            # Ensure array brackets
-            if not str_tool_calls.startswith("["):
-                str_tool_calls = "[" + str_tool_calls
-            if not str_tool_calls.endswith("]"):
-                str_tool_calls = str_tool_calls + "]"
-            
-            # Use partial JSON parser for incomplete JSON
-            json_tool_calls = partial_json_parser.loads(str_tool_calls)
-            
-            if not isinstance(json_tool_calls, list):
-                raise ValueError("Tool calls must be a list")
-            
-            tool_calls = []
-            
-            for tool_call in json_tool_calls:
-                if not isinstance(tool_call, dict):
-                    continue
-                try:
-                    tool_calls.append(ToolCall(
-                        type="function",
-                        function=FunctionCall(
-                            name=tool_call.get("name", ""),
-                            arguments=json.dumps(tool_call.get("arguments", {}), ensure_ascii=False) \
-                                if isinstance(tool_call.get("arguments"), dict) else str(tool_call.get("arguments", "")),
-                        ),
-                    ))
-                except Exception as e:
-                    logger.warning(f"Failed to parse tool call: {e}")
-                    continue
-
-            content = model_output[:model_output.find(self.tool_call_start_token)].strip()
+            # we first try to directly load the json as parsing very nested
+            # jsons is difficult
+            try:
+                if self.fn_name_regex:
+                    matches = self.fn_name_regex.findall(tool_content)
 
+                    function_call_arr = []
+                    for match in matches:
+                        fn_name = match[0]
+                        args = match[1]
+
+                        # fn_name is encoded outside serialized json dump
+                        # only arguments are serialized
+                        function_call_arr.append({
+                            "name": fn_name,
+                            "arguments": json.loads(args)
+                        })
+                else:
+                    function_call_arr = json.loads(tool_content)
+            except json.JSONDecodeError:
+                # use a regex to find the part corresponding to the tool call.
+                # NOTE: This use case should not happen if the model is trained
+                # correctly. It's a easy possible fix so it's included, but
+                # can be brittle for very complex / highly nested tool calls
+                raw_tool_call = self.tool_call_regex.findall(tool_content)[0]
+                function_call_arr = json.loads(raw_tool_call)
+
+            # Tool Call
+            tool_calls: list[NemotronToolCall] = [
+                NemotronToolCall(
+                    type="function",
+                    function=FunctionCall(
+                        name=raw_function_call["name"],
+                        # function call args are JSON but as a string
+                        arguments=json.dumps(raw_function_call["arguments"],
+                                             ensure_ascii=False)))
+                for raw_function_call in function_call_arr
+            ]
+
+            # get any content before  the tool call
+            content = model_output.split(self.bot_token)[0]
             return ExtractedToolCallInformation(
-                tools_called=True if tool_calls else False,
+                tools_called=True,
                 tool_calls=tool_calls,
-                content=content if content else None,
-            )
+                content=content if len(content) > 0 else None)
 
-        except Exception as e:
-            logger.exception(f"Error extracting tool calls. Response: {model_output}")
-            return ExtractedToolCallInformation(
-                tools_called=False,
-                tool_calls=[],
-                content=model_output,
-            )
+        except Exception:
+            logger.exception("Error in extracting tool call from response.")
+            # return information to just treat the tool call as regular JSON
+            return ExtractedToolCallInformation(tools_called=False,
+                                                tool_calls=[],
+                                                content=tool_content)
 
     def extract_tool_calls_streaming(
         self,
@@ -129,108 +172,315 @@ class NemotronJSONToolParser(ToolParser):
         delta_token_ids: Sequence[int],
         request: ChatCompletionRequest,
     ) -> Union[DeltaMessage, None]:
-        """Extract tool calls from streaming output.
-        
-        This incrementally parses the <TOOLCALL> JSON as it streams in,
-        sending delta updates for each tool call and its arguments.
-        """
-        
-        # Check if we just started tool calling
-        if self.tool_call_start_token in delta_text and self.tool_call_start_token not in previous_text:
-            # First time seeing <TOOLCALL>, return content before it
-            content_before = delta_text.split(self.tool_call_start_token)[0]
-            if content_before:
-                return DeltaMessage(content=content_before)
-            # Start of tool call section - no delta yet
-            return None
-        
-        # Check if we're not in tool call mode yet
-        if self.tool_call_start_token not in current_text:
-            # Regular content, no tool calls
-            return DeltaMessage(content=delta_text) if delta_text else None
-        
-        # We're inside <TOOLCALL>...</TOOLCALL>
-        # For Nemotron, the entire TOOLCALL block is generated at once
-        # So we should only parse when we have the complete </TOOLCALL>
-        
-        # Check if we have the complete tool call block yet
-        if self.tool_call_end_token not in current_text:
-            # Incomplete tool call, don't send deltas yet
-            return None
-        
-        # We have the complete tool call block, parse it
-        start_idx = current_text.find(self.tool_call_start_token) + len(self.tool_call_start_token)
-        end_idx = current_text.find(self.tool_call_end_token)
-        json_str = current_text[start_idx:end_idx].strip()
-        
-        # Parse the complete JSON
+        # 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.
         try:
-            # Ensure we have array brackets
-            if not json_str.startswith("["):
-                json_str = "[" + json_str
-            if not json_str.endswith("]"):
-                json_str = json_str + "]"
-            
-            # Parse complete JSON
-            tool_calls_arr = json.loads(json_str)
-            
-            if not isinstance(tool_calls_arr, list):
+            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
             
-            # Generate delta updates for new/updated tool calls
-            delta_tool_calls = []
-            
-            for idx, tool_call in enumerate(tool_calls_arr):
-                if not isinstance(tool_call, dict):
-                    continue
+            # 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
                 
-                # Ensure we have a tool ID for this call
-                while len(self.tool_call_ids) <= idx:
-                    self.tool_call_ids.append(random_uuid())
+                # 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
                 
-                tool_id = self.tool_call_ids[idx]
-                tool_name = tool_call.get("name", "")
-                tool_args = tool_call.get("arguments", {})
                 
-                # Convert arguments to JSON string
-                if isinstance(tool_args, dict):
-                    args_str = json.dumps(tool_args, ensure_ascii=False)
+                # 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:
-                    args_str = str(tool_args)
-                
-                # Check if this is a new tool call
-                if idx >= self.sent_tool_calls_count:
-                    # New tool call - send ID, name, and complete arguments all at once
-                    # This matches how other models (Llama, etc.) send tool calls
-                    delta_tool_calls.append(DeltaToolCall(
-                        index=idx,
-                        id=tool_id,
-                        type="function",
-                        function=DeltaFunctionCall(
-                            name=tool_name,
-                            arguments=args_str  # Send complete JSON string
-                        )
-                    ))
-                    self.sent_tool_calls_count = idx + 1
-                    self.sent_args_length[idx] = len(args_str)
-                
-                # NOTE: We don't send incremental updates for arguments
-                # because Nemotron generates complete tool calls in one shot
-                # Unlike thinking models that stream arguments token-by-token
-            
-            if delta_tool_calls:
-                return DeltaMessage(tool_calls=delta_tool_calls)
+                    # 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
+        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'):
+                return None
+
+        # 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)
+
+        # bit mask flags for partial JSON parsing. If the name hasn't been
+        # sent yet, don't allow sending
+        # an incomplete string since OpenAI only ever (as far as I have
+        # seen) allows sending the entire tool/ function name at once.
+        flags = Allow.ALL if self.current_tool_name_sent \
+            else Allow.ALL & ~Allow.STR
+        end_of_call: bool = False
+        try:
+
+            # replace BOT token with empty string, and convert single quotes
+            # to double to allow parsing as JSON since mistral uses single
+            # quotes instead of double for tool calls
+            parsable_arr = current_text.split(self.bot_token)[-1]
             
-        except Exception as e:
-            # JSON parsing failed (expected for incomplete JSON)
-            logger.debug(f"Partial JSON parse failed (expected during streaming): {e}")
-            pass
-        
-        # Check if we just completed the tool calls (end tag in this delta)
-        if self.tool_call_end_token in delta_text and self.tool_call_end_token not in previous_text:
-            # We just completed - reset state for next potential tool call
-            self.sent_tool_calls_count = 0
-            self.sent_args_length = {}
-            self.tool_call_ids = []
-        
-        return None
+            # Check if we're at the end of the tool call
+            if '</TOOLCALL>' in parsable_arr:
+                end_of_call = True
+                parsable_arr = parsable_arr.split('</TOOLCALL>')[0]
+
+            # tool calls are generated in an array, so do partial JSON
+            # parsing on the entire array
+            try:
+                tool_call_arr: list[dict] = partial_json_parser.loads(
+                    parsable_arr, flags)
+            except partial_json_parser.core.exceptions.MalformedJSON:
+                return None
+
+            current_tool_call: dict = tool_call_arr[self.current_tool_id] \
+                if len(tool_call_arr) > 0 else {}
+
+            # case -- if no tokens have been streamed for the tool, e.g.
+            #   only the array brackets, stream nothing
+            if len(tool_call_arr) == 0:
+                return None
+
+            # case: we are starting a new tool in the array
+            #   -> array has > 0 length AND length has moved past cursor
+            elif (len(tool_call_arr) > 0
+                  and len(tool_call_arr) > self.current_tool_id + 1):
+
+                # if we're moving on to a new call, first make sure we
+                # haven't missed anything in the previous one that was
+                # auto-generated due to JSON completions, but wasn't
+                # streamed to the client yet.
+                if self.current_tool_id >= 0:
+                    diff: Union[str, None] = current_tool_call.get("arguments")
+
+                    if diff:
+                        diff = json.dumps(diff, ensure_ascii=False).replace(
+                            self.streamed_args_for_tool[self.current_tool_id],
+                            "")
+                        delta = DeltaMessage(tool_calls=[
+                            DeltaToolCall(index=self.current_tool_id,
+                                          function=DeltaFunctionCall(
+                                              arguments=diff).model_dump(
+                                                  exclude_none=True))
+                        ])
+                        self.streamed_args_for_tool[
+                            self.current_tool_id] += diff
+                    else:
+                        delta = None
+                else:
+                    delta = None
+                # re-set stuff pertaining to progress in the current tool
+                self.current_tool_id = len(tool_call_arr) - 1
+                self.current_tool_name_sent = False
+                self.streamed_args_for_tool.append("")
+                return delta
+
+            # case: update an existing tool - this is handled below
+
+            # if the current tool name hasn't been sent, send if available
+            # - otherwise send nothing
+            if not self.current_tool_name_sent:
+                function_name = current_tool_call.get("name")
+                if function_name:
+
+                    delta = DeltaMessage(tool_calls=[
+                        DeltaToolCall(index=self.current_tool_id,
+                                      type="function",
+                                      id=NemotronToolCall.generate_random_id(),
+                                      function=DeltaFunctionCall(
+                                          name=function_name).model_dump(
+                                              exclude_none=True))
+                    ])
+                    self.current_tool_name_sent = True
+                else:
+                    delta = None
+
+            # now we know we're on the same tool call and we're streaming
+            # arguments
+            else:
+
+                prev_arguments = self.prev_tool_call_arr[
+                    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
+                elif not cur_arguments and prev_arguments:
+                    logger.error(
+                        "INVARIANT - impossible to have arguments reset "
+                        "mid-arguments")
+                    delta = None
+                elif cur_arguments and not prev_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]
+                    if arguments_delta:
+                        delta = DeltaMessage(tool_calls=[
+                            DeltaToolCall(index=self.current_tool_id,
+                                          function=DeltaFunctionCall(
+                                              arguments=arguments_delta).
+                                          model_dump(exclude_none=True))
+                        ])
+                        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
+                    else:
+                        # Do not flush final JSON here; let the serving layer
+                        # compute a minimal remaining suffix on finish.
+                        delta = None
+                else:
+                    # End-of-call or equal state; do not force a final flush here.
+                    delta = None
+
+            # check to see if the name is defined and has been sent. if so,
+            # stream the name - otherwise keep waiting
+            # finish by setting old and returning None as base case
+            self.prev_tool_call_arr = tool_call_arr
+            # If we've reached the end of a tool call, flush any remaining
+            # suffix (including a final '}') that hasn't been streamed yet.
+            if end_of_call and self.current_tool_id >= 0:
+                try:
+                    cur_arguments = current_tool_call.get("arguments")
+                    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)
+
+                        # 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:
+                            extra = DeltaToolCall(
+                                index=self.current_tool_id,
+                                function=DeltaFunctionCall(
+                                    arguments=remaining_suffix).model_dump(
+                                        exclude_none=True))
+                            if delta is None:
+                                delta = DeltaMessage(tool_calls=[extra])
+                            else:
+                                if getattr(delta, "tool_calls", None):
+                                    delta.tool_calls.append(extra)
+                                else:
+                                    delta.tool_calls = [extra]
+                            self.streamed_args_for_tool[
+                                self.current_tool_id] += remaining_suffix
+                        else:
+                            pass
+                except Exception:
+                    pass
+
+            return delta
+
+        except Exception:
+            logger.exception("Error trying to handle streaming tool call.")
+            return None