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 10, 12:45 UTC
[OpenAI] Fix tool_choice=required streaming when output has trailing extra data (#31610)
Signed-off-by: maylikenoother <ogedengbemary19@gmail.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
Mary committed
Jan 8, 2026 at 13:01 UTC
7645bc524b2fec6ae8e75924f99d19225aa779a9
2 files changed
+41
-2
tests/tool_use/test_tool_choice_required.py
+33
@@ -311,6 +311,7 @@ def test_streaming_output_valid(output, empty_params, delta_len):
311
previous_text = current_text
312
313
assert len(messages) > 0
314
+
315
combined_messages = "["
316
for message in messages:
317
if message.tool_calls[0].function.name:
@@ -328,3 +329,35 @@ def test_streaming_output_valid(output, empty_params, delta_len):
329
combined_messages += "}]"
330
assert json.loads(combined_messages) == output
331
assert json.dumps(json.loads(combined_messages)) == output_json
332
+
333
+
334
+def test_streaming_output_valid_with_trailing_extra_data():
335
+ self = MagicMock()
336
+
337
+ output = [{"name": "get_current_weather", "parameters": {"city": "Vienna"}}]
338
+ output_json = json.dumps(output) + "\nDONE"
339
+
340
+ previous_text = ""
341
+ function_name_returned = False
342
+ messages = []
343
+ delta_len = 3
344
+ for i in range(0, len(output_json), delta_len):
345
+ delta_text = output_json[i : i + delta_len]
346
+ current_text = previous_text + delta_text
347
+
348
+ delta_message, function_name_returned = (
349
+ OpenAIServingChat.extract_tool_call_required_streaming(
350
+ self,
351
+ previous_text=previous_text,
352
+ current_text=current_text,
353
+ delta_text=delta_text,
354
+ function_name_returned=function_name_returned,
355
+ )
356
+ )
357
+
358
+ if delta_message:
359
+ messages.append(delta_message)
360
+
361
+ previous_text = current_text
362
+
363
+ assert len(messages) > 0
vllm/entrypoints/openai/serving_chat.py
+8
-2
@@ -13,6 +13,7 @@ import partial_json_parser
13
import regex as re
14
from fastapi import Request
15
from openai_harmony import Message as OpenAIMessage
16
+from partial_json_parser.core.options import Allow
17
18
from vllm.engine.protocol import EngineClient
19
from vllm.entrypoints.chat_utils import (
@@ -76,6 +77,7 @@ from vllm.tokenizers.mistral import (
77
)
78
from vllm.tool_parsers import ToolParser
79
from vllm.tool_parsers.mistral_tool_parser import MistralToolCall
80
+from vllm.tool_parsers.utils import partial_json_loads
81
from vllm.utils.collection_utils import as_list
82
from vllm.v1.sample.logits_processor import validate_logits_processors_parameters
83
@@ -511,8 +513,12 @@ class OpenAIServingChat(OpenAIServing):
513
# if the current text is empty, we cannot parse it
514
return None, function_name_returned
515
try:
514
- obj = partial_json_parser.loads(current_text)
515
- except partial_json_parser.core.exceptions.MalformedJSON:
516
+ flags = Allow.ALL
517
+ obj, _ = partial_json_loads(current_text, flags)
518
+ except (
519
+ partial_json_parser.core.exceptions.MalformedJSON,
520
+ json.JSONDecodeError,
521
+ ):
522
logger.debug("not enough tokens to parse into JSON yet")
523
obj = None
524