A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 06:15 UTC
[BUGFIX] Support streamed_args_for_tool in MistralToolParser (#41730)
Signed-off-by: juliendenize <julien.denize@mistral.ai>
Julien Denize committed
May 5, 2026 at 19:48 UTC
c6235ed1803e31bafedd90d53766e5705c570780
2 files changed
+55
-15
tests/tool_parsers/test_mistral_tool_parser.py
+29
@@ -590,6 +590,33 @@ def _test_extract_tool_calls_streaming(
590
]
591
assert_tool_calls(actual_tool_calls, expected_tool_calls)
592
593
+ if expected_tool_calls:
594
+ assert len(tool_parser.streamed_args_for_tool) == len(expected_tool_calls)
595
+ assert len(tool_parser.prev_tool_call_arr) == len(expected_tool_calls)
596
+ for i in range(len(expected_tool_calls)):
597
+ assert (
598
+ tool_parser.prev_tool_call_arr[i]["arguments"]
599
+ == tool_parser.streamed_args_for_tool[i]
600
+ )
601
+ assert tool_parser.streamed_args_for_tool[i] == function_args_strs[i]
602
+ assert (
603
+ tool_parser.prev_tool_call_arr[i]["name"]
604
+ == expected_tool_calls[i].function.name
605
+ )
606
+
607
+ # Simulate the serving layer's unstreamed-args check
608
+ index = len(tool_parser.prev_tool_call_arr) - 1
609
+ args = tool_parser.prev_tool_call_arr[index].get("arguments", {})
610
+ expected_call = (
611
+ args if isinstance(args, str) else json.dumps(args, ensure_ascii=False)
612
+ )
613
+ actual_call = tool_parser.streamed_args_for_tool[index]
614
+ remaining_call = expected_call.replace(actual_call, "", 1)
615
+ assert remaining_call == ""
616
+ else:
617
+ assert len(tool_parser.streamed_args_for_tool) == 0
618
+ assert len(tool_parser.prev_tool_call_arr) == 0
619
+
620
621
@pytest.mark.parametrize(
622
ids=[
@@ -855,6 +882,8 @@ def test_extract_tool_calls_streaming_v11_no_tools(
882
previous_text = current_text
883
884
assert collected_content == model_output
885
+ assert len(mistral_tool_parser.streamed_args_for_tool) == 0
886
+ assert len(mistral_tool_parser.prev_tool_call_arr) == 0
887
888
889
@pytest.mark.parametrize(
vllm/tool_parsers/mistral_tool_parser.py
+26
-15
@@ -623,13 +623,6 @@ class MistralToolParser(ToolParser):
623
if len(delta_tool_calls) > 0:
624
delta.tool_calls = delta_tool_calls
625
626
- # HACK: serving_chat.py inspects the internal state of tool parsers
627
- # when determining its final streaming delta, automatically
628
- # adding autocompleted JSON.
629
- # These two lines avoid that nonsense while ensuring finish_reason
630
- # is set to tool_calls when at least one tool is called.
631
- if delta_tool_calls and not self.prev_tool_call_arr:
632
- self.prev_tool_call_arr = [{"arguments": {}}]
626
return delta
627
628
def _generate_delta_tool_call(self, delta_text: str) -> list[DeltaToolCall]:
@@ -642,6 +635,8 @@ class MistralToolParser(ToolParser):
635
StreamingState.PARSING_ARGUMENTS,
636
] and delta_text.startswith(self.bot_token):
637
self.current_tool_id += 1
638
+ self.streamed_args_for_tool.append("")
639
+ self.prev_tool_call_arr.append({})
640
self.streaming_state = StreamingState.PARSING_NAME
641
delta_text = delta_text.replace(self.bot_token, "", 1)
642
if self.streaming_state == StreamingState.PARSING_NAME:
@@ -655,6 +650,9 @@ class MistralToolParser(ToolParser):
650
self.current_tool_name += delta_function_name
651
# HF tokenizers may include [ARGS] in the text
652
self.current_tool_name = self.current_tool_name.replace("[ARGS]", "")
653
+ self.prev_tool_call_arr[self.current_tool_id]["name"] = (
654
+ self.current_tool_name
655
+ )
656
delta_text = delta_text[len(delta_function_name) :]
657
self.streaming_state = StreamingState.PARSING_ARGUMENTS
658
else:
@@ -671,6 +669,10 @@ class MistralToolParser(ToolParser):
669
self.streaming_state = StreamingState.TOOL_COMPLETE
670
else:
671
delta_arguments = delta_text
672
+ self.streamed_args_for_tool[self.current_tool_id] += delta_arguments
673
+ self.prev_tool_call_arr[self.current_tool_id]["arguments"] = (
674
+ self.streamed_args_for_tool[self.current_tool_id]
675
+ )
676
ret = []
677
if self.current_tool_name or delta_arguments:
678
ret += [
@@ -820,9 +822,12 @@ class MistralToolParser(ToolParser):
822
if self.current_tool_mistral_id is not None:
823
current_tool_call.id = self.current_tool_mistral_id
824
self.current_tool_mistral_id = None
825
+ self._track_streamed_args_pre_v11(current_tool_call)
826
delta_tool_calls.append(current_tool_call)
827
current_tool_call_modified = False
828
self.current_tool_id += 1
829
+ self.streamed_args_for_tool.append("")
830
+ self.prev_tool_call_arr.append({})
831
self.current_tool_mistral_id = MistralToolCall.generate_random_id()
832
current_tool_call = DeltaToolCall(
833
index=self.current_tool_id,
@@ -835,6 +840,9 @@ class MistralToolParser(ToolParser):
840
# we have the complete tool name
841
current_tool_call_modified = True
842
current_tool_call.function.name = self.current_tool_name
843
+ self.prev_tool_call_arr[self.current_tool_id]["name"] = (
844
+ self.current_tool_name
845
+ )
846
self.current_tool_name = None
847
if self.streaming_state == StreamingState.PARSING_NAME_COMPLETED:
848
self.streaming_state = StreamingState.WAITING_FOR_TOOL_KEY
@@ -860,16 +868,9 @@ class MistralToolParser(ToolParser):
868
if self.current_tool_mistral_id is not None:
869
current_tool_call.id = self.current_tool_mistral_id
870
self.current_tool_mistral_id = None
871
+ self._track_streamed_args_pre_v11(current_tool_call)
872
delta_tool_calls.append(current_tool_call)
873
865
- # HACK: serving_chat.py inspects the internal state of tool parsers
866
- # when determining it's final streaming delta, automatically
867
- # adding autocompleted JSON.
868
- # These two lines avoid that nonsense while ensuring finish_reason
869
- # is set to tool_calls when at least one tool is called.
870
- if delta_tool_calls and not self.prev_tool_call_arr:
871
- self.prev_tool_call_arr = [{"arguments": {}}]
872
-
874
if content or len(delta_tool_calls) > 0:
875
delta_message = DeltaMessage()
876
if content:
@@ -883,6 +884,16 @@ class MistralToolParser(ToolParser):
884
else:
885
return None
886
887
+ def _track_streamed_args_pre_v11(self, tool_call: DeltaToolCall) -> None:
888
+ r"""Accumulate `tool_call` arguments into the streaming state."""
889
+ if tool_call.function is not None and tool_call.function.arguments is not None:
890
+ self.streamed_args_for_tool[self.current_tool_id] += (
891
+ tool_call.function.arguments
892
+ )
893
+ self.prev_tool_call_arr[self.current_tool_id]["arguments"] = (
894
+ self.streamed_args_for_tool[self.current_tool_id]
895
+ )
896
+
897
def _split_delta(
898
self,
899
delta_text: str,