@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9affc17a05
Fork 0

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 2, 07:15 UTC

[Refactor] Move unstreamed tool-arg flush from serving layer to parser (#44017)

Signed-off-by: sfeng33 <4florafeng@gmail.com>

Flora Feng committed Jun 1, 2026 at 22:37 UTC 9affc17a05a49d42c19bb67336f1ea9c4f136af8
5 files changed +138 -264
tests/entrypoints/openai/chat_completion/test_serving_chat.py
+11 -137
@@ -1935,8 +1935,10 @@ async def test_streaming_n_gt1_independent_tool_parsers():
1935 finished=True,
1936 )
1937
1938 - # Collect tool-call deltas per choice from the SSE stream.
1938 + # Collect tool-call deltas and finish_reasons per choice from the SSE
1939 + # stream.
1940 tc_deltas_by_choice: dict[int, list[dict]] = {i: [] for i in range(num_choices)}
1941 + finish_reasons_by_choice: dict[int, list[str]] = {i: [] for i in range(num_choices)}
1942 async for chunk_str in serving_chat.chat_completion_stream_generator(
1943 request=request,
1944 result_generator=result_generator(),
@@ -1959,6 +1961,8 @@ async def test_streaming_n_gt1_independent_tool_parsers():
1961 if delta.get("tool_calls"):
1962 for tc in delta["tool_calls"]:
1963 tc_deltas_by_choice[idx].append(tc)
1964 + if choice.get("finish_reason") is not None:
1965 + finish_reasons_by_choice[idx].append(choice["finish_reason"])
1966
1967 # Both choices must independently produce the correct tool call.
1968 for choice_idx in range(num_choices):
@@ -1984,141 +1988,11 @@ async def test_streaming_n_gt1_independent_tool_parsers():
1988 f"Choice {choice_idx}: expected {{'city': 'Tokyo'}}, got {parsed_args}"
1989 )
1990
1987 -
1988 -class TestCreateRemainingArgsDelta:
1989 - """Tests for _create_remaining_args_delta helper function.
1990 -
1991 - This helper is used when streaming tool calls to preserve id/type/name
1992 - fields in the finish chunk, which would otherwise be lost.
1993 - """
1994 -
1995 - def test_preserves_id_type_name(self):
1996 - """Test that id, type, and name are preserved from original delta."""
1997 - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
1998 - from vllm.entrypoints.openai.engine.protocol import (
1999 - DeltaFunctionCall,
2000 - DeltaMessage,
2001 - DeltaToolCall,
2002 - )
2003 -
2004 - original_delta = DeltaMessage(
2005 - tool_calls=[
2006 - DeltaToolCall(
2007 - index=0,
2008 - id="call_abc123",
2009 - type="function",
2010 - function=DeltaFunctionCall(
2011 - name="get_weather",
2012 - arguments='{"location": "Paris"}',
2013 - ),
2014 - )
2015 - ]
1991 + reasons = finish_reasons_by_choice[choice_idx]
1992 + assert len(reasons) == 1, (
1993 + f"Choice {choice_idx}: expected exactly 1 finish_reason, got {reasons}"
1994 )
2017 -
2018 - result = OpenAIServingChat._create_remaining_args_delta(
2019 - original_delta, '", "unit": "celsius"}', 0
2020 - )
2021 -
2022 - assert len(result.tool_calls) == 1
2023 - tc = result.tool_calls[0]
2024 - assert tc.index == 0
2025 - assert tc.id == "call_abc123"
2026 - assert tc.type == "function"
2027 - assert tc.function.name == "get_weather"
2028 - assert tc.function.arguments == '", "unit": "celsius"}'
2029 -
2030 - def test_matches_by_index(self):
2031 - """Test that the correct tool call is matched by index."""
2032 - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
2033 - from vllm.entrypoints.openai.engine.protocol import (
2034 - DeltaFunctionCall,
2035 - DeltaMessage,
2036 - DeltaToolCall,
2037 - )
2038 -
2039 - original_delta = DeltaMessage(
2040 - tool_calls=[
2041 - DeltaToolCall(
2042 - index=0,
2043 - id="call_first",
2044 - type="function",
2045 - function=DeltaFunctionCall(name="func_a", arguments="{}"),
2046 - ),
2047 - DeltaToolCall(
2048 - index=1,
2049 - id="call_second",
2050 - type="function",
2051 - function=DeltaFunctionCall(name="func_b", arguments="{}"),
2052 - ),
2053 - ]
2054 - )
2055 -
2056 - result = OpenAIServingChat._create_remaining_args_delta(
2057 - original_delta, '{"extra": true}', 1
1995 + assert reasons[0] == "tool_calls", (
1996 + f"Choice {choice_idx}: expected finish_reason='tool_calls', "
1997 + f"got '{reasons[0]}'"
1998 )
2059 -
2060 - assert len(result.tool_calls) == 1
2061 - tc = result.tool_calls[0]
2062 - assert tc.index == 1
2063 - assert tc.id == "call_second"
2064 - assert tc.function.name == "func_b"
2065 -
2066 - def test_no_matching_tool_call(self):
2067 - """Test graceful handling when no matching tool call is found."""
2068 - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
2069 - from vllm.entrypoints.openai.engine.protocol import (
2070 - DeltaFunctionCall,
2071 - DeltaMessage,
2072 - DeltaToolCall,
2073 - )
2074 -
2075 - original_delta = DeltaMessage(
2076 - tool_calls=[
2077 - DeltaToolCall(
2078 - index=0,
2079 - id="call_zero",
2080 - type="function",
2081 - function=DeltaFunctionCall(name="func", arguments="{}"),
2082 - )
2083 - ]
2084 - )
2085 -
2086 - result = OpenAIServingChat._create_remaining_args_delta(
2087 - original_delta, '{"arg": 1}', 5
2088 - )
2089 -
2090 - assert len(result.tool_calls) == 1
2091 - tc = result.tool_calls[0]
2092 - assert tc.index == 5
2093 - assert tc.id is None
2094 - assert tc.type is None
2095 - assert tc.function.name is None
2096 - assert tc.function.arguments == '{"arg": 1}'
2097 -
2098 - def test_function_is_none(self):
2099 - """Test handling when original tool call has no function."""
2100 - from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
2101 - from vllm.entrypoints.openai.engine.protocol import DeltaMessage, DeltaToolCall
2102 -
2103 - original_delta = DeltaMessage(
2104 - tool_calls=[
2105 - DeltaToolCall(
2106 - index=0,
2107 - id="call_nofunc",
2108 - type="function",
2109 - function=None,
2110 - )
2111 - ]
2112 - )
2113 -
2114 - result = OpenAIServingChat._create_remaining_args_delta(
2115 - original_delta, '{"data": "value"}', 0
2116 - )
2117 -
2118 - assert len(result.tool_calls) == 1
2119 - tc = result.tool_calls[0]
2120 - assert tc.index == 0
2121 - assert tc.id == "call_nofunc"
2122 - assert tc.type == "function"
2123 - assert tc.function.name is None
2124 - assert tc.function.arguments == '{"data": "value"}'
tests/parser/test_streaming.py
+83
@@ -235,3 +235,86 @@ def test_parse_delta_reasoning_only_thinking_disabled(tokenizer, request_obj):
235 assert "Hello" in content
236 assert "assist" in content
237 assert len(tool_calls) == 0
238 +
239 +
240 +def test_parse_delta_finished_no_flush_without_tool_call_delta(tokenizer, request_obj):
241 + """When finished=True but the final parse_delta produces no
242 + tool-call delta, unstreamed args are not flushed."""
243 + parser = make_parser(tokenizer, reasoning=False, tool=True)
244 +
245 + results = stream_text(
246 + parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[]
247 + )
248 + _, _, tool_calls = collect_fields(results)
249 + assert len(tool_calls) > 0
250 +
251 + streamed = parser._tool_parser.streamed_args_for_tool[0]
252 + assert len(streamed) > 5
253 + parser._tool_parser.streamed_args_for_tool[0] = streamed[:-5]
254 +
255 + # Prevent normal extraction from catching the gap — without a
256 + # tool-call delta to merge into, the flush is skipped.
257 + parser._tool_parser.extract_tool_calls_streaming = lambda *a, **kw: None
258 +
259 + flush_result = parser.parse_delta("", [], request_obj, finished=True)
260 + assert flush_result is None or flush_result.tool_calls is None
261 +
262 +
263 +def test_parse_delta_finished_no_extra_args_when_fully_streamed(tokenizer, request_obj):
264 + """When all args have been streamed, finished=True must not
265 + produce extra or duplicate arguments."""
266 + parser = make_parser(tokenizer, reasoning=False, tool=True)
267 + results = stream_text(
268 + parser, tokenizer, MODEL_OUTPUT, request_obj, prompt_token_ids=[]
269 + )
270 + _, _, tool_calls = collect_fields(results)
271 +
272 + assert len(tool_calls) > 0
273 + assert tool_calls[0].function.name == "get_weather"
274 + tool_args = "".join(
275 + tc.function.arguments for tc in tool_calls if tc.function.arguments
276 + )
277 + assert json.loads(tool_args) == {"city": "Dallas"}
278 +
279 + flush_result = parser.parse_delta("", [], request_obj, finished=True)
280 + assert flush_result is None or flush_result.tool_calls is None
281 +
282 +
283 +def test_parse_delta_finished_appends_remaining_args(tokenizer, request_obj):
284 + """When finished=True and the tool parser has unstreamed args,
285 + parse_delta appends the remaining arguments to the tool-call delta."""
286 + parser = make_parser(tokenizer, reasoning=False, tool=True)
287 + token_ids = tokenizer.encode(MODEL_OUTPUT, add_special_tokens=False)
288 +
289 + remainder = ',"unit":"celsius"}'
290 + prompt_ids: list[int] | None = []
291 + results: list[DeltaMessage | None] = []
292 + for i, tid in enumerate(token_ids):
293 + prev = results[-1] if results else None
294 + prev_had_args = (
295 + prev
296 + and prev.tool_calls
297 + and any(tc.function and tc.function.arguments for tc in prev.tool_calls)
298 + )
299 +
300 + if prev_had_args:
301 + parser._tool_parser.get_remaining_unstreamed_args = lambda: remainder
302 +
303 + result = parser.parse_delta(
304 + tokenizer.decode([tid]),
305 + [tid],
306 + request_obj,
307 + prompt_token_ids=prompt_ids,
308 + finished=prev_had_args,
309 + )
310 + prompt_ids = None
311 + results.append(result)
312 +
313 + if prev_had_args:
314 + break
315 +
316 + _, _, tool_calls = collect_fields(results)
317 + tool_args = "".join(
318 + tc.function.arguments for tc in tool_calls if tc.function.arguments
319 + )
320 + assert tool_args.endswith(remainder)
vllm/entrypoints/openai/chat_completion/serving.py
+4 -127
@@ -3,7 +3,6 @@
3
4 import asyncio
5 import io
6 -import json
6 import time
7 from collections.abc import AsyncGenerator, AsyncIterator
8 from collections.abc import Sequence as GenericSequence
@@ -40,9 +39,7 @@ from vllm.entrypoints.openai.chat_completion.stream_harmony import (
39 extract_harmony_streaming_delta,
40 )
41 from vllm.entrypoints.openai.engine.protocol import (
43 - DeltaFunctionCall,
42 DeltaMessage,
45 - DeltaToolCall,
43 ErrorResponse,
44 FunctionCall,
45 PromptTokenUsageInfo,
@@ -65,7 +62,7 @@ from vllm.entrypoints.utils import get_max_tokens, should_include_usage
62 from vllm.inputs import EngineInput
63 from vllm.logger import init_logger
64 from vllm.logprobs import Logprob
68 -from vllm.outputs import CompletionOutput, RequestOutput
65 +from vllm.outputs import RequestOutput
66 from vllm.parser import ParserManager
67 from vllm.parser.abstract_parser import Parser
68 from vllm.reasoning import ReasoningParser
@@ -715,6 +712,7 @@ class OpenAIServingChat(OpenAIServing):
712 delta_token_ids=as_list(output.token_ids),
713 request=request,
714 prompt_token_ids=res.prompt_token_ids,
715 + finished=output.finish_reason is not None,
716 )
717 if delta_message and delta_message.tool_calls:
718 tools_streamed[i] = True
@@ -805,81 +803,13 @@ class OpenAIServingChat(OpenAIServing):
803 # finish_reason='error' indicates a retryable error
804 self._raise_if_error(output.finish_reason, request_id)
805
808 - # check to make sure we haven't "forgotten" to stream
809 - # any tokens that were generated but previously
810 - # matched by partial json parsing
811 - # only happens if we are NOT using structured outputs
812 - index = 0
813 - auto_tools_called = False
814 - if tool_parser:
815 - auto_tools_called = len(tool_parser.prev_tool_call_arr) > 0
816 - index = (
817 - len(tool_parser.prev_tool_call_arr) - 1
818 - if auto_tools_called
819 - else 0
820 - )
821 - should_check = (
822 - self._should_check_for_unstreamed_tool_arg_tokens(
823 - delta_message, output
824 - )
825 - )
826 - # only check if there are any tool calls
827 - # detected by partial parsing
828 - if should_check and tool_parser and auto_tools_called:
829 - latest_delta_len = 0
830 - if (
831 - isinstance(
832 - delta_message.tool_calls[0].function,
833 - DeltaFunctionCall,
834 - )
835 - ) and isinstance(
836 - delta_message.tool_calls[0].function.arguments, str
837 - ):
838 - latest_delta_len = len(
839 - delta_message.tool_calls[0].function.arguments
840 - )
841 -
842 - # get the expected call based on partial JSON
843 - # parsing which "autocompletes" the JSON.
844 - # Tool parsers (e.g. Qwen3Coder) store
845 - # arguments as a JSON string in
846 - # prev_tool_call_arr. Calling json.dumps()
847 - # on an already-serialized string would
848 - # double-serialize it (e.g. '{"k":1}' becomes
849 - # '"{\\"k\\":1}"'), which then causes the
850 - # replace() below to fail and append the
851 - # entire double-serialized string as a
852 - # spurious final delta.
853 - args = tool_parser.prev_tool_call_arr[index].get(
854 - "arguments", {}
855 - )
856 - if isinstance(args, str):
857 - expected_call = args
858 - else:
859 - expected_call = json.dumps(args, ensure_ascii=False)
860 -
861 - # get what we've streamed so far for arguments
862 - # for the current tool
863 - actual_call = tool_parser.streamed_args_for_tool[index]
864 - if latest_delta_len > 0:
865 - actual_call = actual_call[:-latest_delta_len]
866 -
867 - # check to see if there's anything left to stream
868 - remaining_call = expected_call.replace(actual_call, "", 1)
869 - # set that as a delta message
870 - delta_message = self._create_remaining_args_delta(
871 - delta_message, remaining_call, index
872 - )
873 -
806 # Send the finish response for each request.n only once
807 # In OpenAI's API, when a tool is called, the
808 # finish_reason is:
809 # "tool_calls" for "auto" or "required" tool calls,
810 # and "stop" for named tool calls.
879 - if (
880 - auto_tools_called
881 - or (tools_streamed[i] and not tool_choice_function_name)
882 - or (self.use_harmony and harmony_tools_streamed[i])
811 + if (tools_streamed[i] and not tool_choice_function_name) or (
812 + self.use_harmony and harmony_tools_streamed[i]
813 ):
814 finish_reason_ = "tool_calls"
815 else:
@@ -1535,56 +1465,3 @@ class OpenAIServingChat(OpenAIServing):
1465 and self.enable_auto_tools
1466 and request.tool_choice in ["auto", None]
1467 )
1538 -
1539 - def _should_check_for_unstreamed_tool_arg_tokens(
1540 - self,
1541 - delta_message: DeltaMessage | None,
1542 - output: CompletionOutput,
1543 - ) -> bool:
1544 - """
1545 - Check to see if we should check for unstreamed tool arguments tokens.
1546 - This is only applicable when auto tool parsing is enabled, the delta
1547 - is a tool call with arguments.
1548 - """
1549 -
1550 - return bool(
1551 - # if there is a delta message that includes tool calls which
1552 - # include a function that has arguments
1553 - output.finish_reason is not None
1554 - and self.enable_auto_tools
1555 - and self.tool_parser
1556 - and delta_message
1557 - and delta_message.tool_calls
1558 - and delta_message.tool_calls[0]
1559 - and delta_message.tool_calls[0].function
1560 - and delta_message.tool_calls[0].function.arguments is not None
1561 - )
1562 -
1563 - @staticmethod
1564 - def _create_remaining_args_delta(
1565 - delta_message: DeltaMessage,
1566 - remaining_call: str,
1567 - index: int,
1568 - ) -> DeltaMessage:
1569 - """
1570 - Create a delta message for remaining tool arguments, preserving
1571 - id/type/name from the original delta.
1572 - """
1573 - original_tc = next(
1574 - (tc for tc in delta_message.tool_calls if tc.index == index),
1575 - None,
1576 - )
1577 - original_fn = original_tc.function if original_tc else None
1578 - return DeltaMessage(
1579 - tool_calls=[
1580 - DeltaToolCall(
1581 - index=index,
1582 - id=original_tc.id if original_tc else None,
1583 - type=original_tc.type if original_tc else None,
1584 - function=DeltaFunctionCall(
1585 - name=original_fn.name if original_fn else None,
1586 - arguments=remaining_call,
1587 - ),
1588 - )
1589 - ]
1590 - )
vllm/parser/abstract_parser.py
+21
@@ -320,6 +320,7 @@ class Parser:
320 delta_token_ids: list[int],
321 request: ChatCompletionRequest | ResponsesRequest,
322 prompt_token_ids: list[int] | None = None,
323 + finished: bool = False,
324 ) -> DeltaMessage | None:
325 """Parse a single streaming delta, orchestrating reasoning then
326 tool call extraction via internal stream state.
@@ -656,12 +657,28 @@ class DelegatingParser(Parser):
657 return False
658 return state.reasoning_ended
659
660 + def _append_unstreamed_tool_args(
661 + self,
662 + delta_message: DeltaMessage | None,
663 + ) -> None:
664 + """Append parsed-but-unstreamed tool-call arguments to *delta_message*."""
665 + if (
666 + self._tool_parser is not None
667 + and delta_message
668 + and delta_message.tool_calls
669 + and (last_tc := delta_message.tool_calls[-1]).function
670 + ):
671 + last_tc.function.arguments = (
672 + last_tc.function.arguments or ""
673 + ) + self._tool_parser.get_remaining_unstreamed_args()
674 +
675 def parse_delta(
676 self,
677 delta_text: str,
678 delta_token_ids: list[int],
679 request: ChatCompletionRequest | ResponsesRequest,
680 prompt_token_ids: list[int] | None = None,
681 + finished: bool = False,
682 ) -> DeltaMessage | None:
683 state = self._stream_state
684
@@ -745,6 +762,10 @@ class DelegatingParser(Parser):
762
763 state.previous_text = current_text
764 state.previous_token_ids = current_token_ids
765 +
766 + if finished:
767 + self._append_unstreamed_tool_args(delta_message)
768 +
769 return delta_message
770
771
vllm/tool_parsers/abstract_tool_parser.py
+19
@@ -79,6 +79,25 @@ class ToolParser:
79 else:
80 self.tools = []
81
82 + def get_remaining_unstreamed_args(self) -> str:
83 + """Return tool call arguments parsed but not yet streamed."""
84 + if not self.prev_tool_call_arr:
85 + return ""
86 + index = len(self.prev_tool_call_arr) - 1
87 + args = self.prev_tool_call_arr[index].get("arguments", {})
88 + if isinstance(args, str):
89 + expected = args
90 + else:
91 + expected = json.dumps(args, ensure_ascii=False)
92 + actual = (
93 + self.streamed_args_for_tool[index]
94 + if index < len(self.streamed_args_for_tool)
95 + else ""
96 + )
97 + if expected.startswith(actual):
98 + return expected[len(actual) :]
99 + return ""
100 +
101 @cached_property
102 def vocab(self) -> dict[str, int]:
103 # NOTE: Only PreTrainedTokenizerFast is guaranteed to have .vocab