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 13, 08:15 UTC
[Frontend] split append tool output (#28333)
Signed-off-by: Andrew Xia <axia@fb.com> Co-authored-by: Andrew Xia <axia@fb.com>
Andrew Xia committed
Nov 12, 2025 at 20:03 UTC
7c38ed0f1c83077111bac6dd044d441c6e4fab67
3 files changed
+73
-63
tests/entrypoints/openai/test_serving_responses.py
+3
@@ -34,6 +34,9 @@ class MockConversationContext(ConversationContext):
34
def append_output(self, output) -> None:
35
pass
36
37
+ def append_tool_output(self, output) -> None:
38
+ pass
39
+
40
async def call_tool(self):
41
return []
42
vllm/entrypoints/context.py
+69
-62
@@ -80,7 +80,11 @@ class TurnMetrics:
80
81
class ConversationContext(ABC):
82
@abstractmethod
83
- def append_output(self, output) -> None:
83
+ def append_output(self, output: RequestOutput) -> None:
84
+ pass
85
+
86
+ @abstractmethod
87
+ def append_tool_output(self, output) -> None:
88
pass
89
90
@abstractmethod
@@ -151,6 +155,9 @@ class SimpleContext(ConversationContext):
155
self.num_cached_tokens = output.num_cached_tokens or 0
156
self.num_output_tokens += len(output.outputs[0].token_ids or [])
157
158
+ def append_tool_output(self, output) -> None:
159
+ raise NotImplementedError("Should not be called.")
160
+
161
def need_builtin_tool_call(self) -> bool:
162
return False
163
@@ -205,28 +212,28 @@ class HarmonyContext(ConversationContext):
212
if self.parser.current_channel in {"analysis", "commentary"}:
213
self.num_reasoning_tokens += 1
214
208
- def append_output(self, output: RequestOutput | list[Message]) -> None:
209
- if isinstance(output, RequestOutput):
210
- output_token_ids = output.outputs[0].token_ids
211
- self.parser = get_streamable_parser_for_assistant()
212
- for token_id in output_token_ids:
213
- self.parser.process(token_id)
214
- # Check if the current token is part of reasoning content
215
- self._update_num_reasoning_tokens()
216
- self._update_prefill_token_usage(output)
217
- self._update_decode_token_usage(output)
218
- # Append current turn to all turn list for next turn's calculations
219
- self.all_turn_metrics.append(self.current_turn_metrics.copy())
220
- self.current_turn_metrics.reset()
221
- # append_output is called only once before tool calling
222
- # in non-streaming case
223
- # so we can append all the parser messages to _messages
224
- output_msgs = self.parser.messages
225
- # The responses finish reason is set in the last message
226
- self.finish_reason = output.outputs[0].finish_reason
227
- else:
228
- # Tool output.
229
- output_msgs = output
215
+ def append_output(self, output: RequestOutput) -> None:
216
+ output_token_ids = output.outputs[0].token_ids
217
+ self.parser = get_streamable_parser_for_assistant()
218
+ for token_id in output_token_ids:
219
+ self.parser.process(token_id)
220
+ # Check if the current token is part of reasoning content
221
+ self._update_num_reasoning_tokens()
222
+ self._update_prefill_token_usage(output)
223
+ self._update_decode_token_usage(output)
224
+ # Append current turn to all turn list for next turn's calculations
225
+ self.all_turn_metrics.append(self.current_turn_metrics.copy())
226
+ self.current_turn_metrics.reset()
227
+ # append_output is called only once before tool calling
228
+ # in non-streaming case
229
+ # so we can append all the parser messages to _messages
230
+ output_msgs = self.parser.messages
231
+ # The responses finish reason is set in the last message
232
+ self.finish_reason = output.outputs[0].finish_reason
233
+ self._messages.extend(output_msgs)
234
+
235
+ def append_tool_output(self, output: list[Message]) -> None:
236
+ output_msgs = output
237
self._messages.extend(output_msgs)
238
239
def _update_prefill_token_usage(self, output: RequestOutput) -> None:
@@ -502,45 +509,45 @@ class StreamingHarmonyContext(HarmonyContext):
509
def messages(self) -> list:
510
return self._messages
511
505
- def append_output(self, output: RequestOutput | list[Message]) -> None:
506
- if isinstance(output, RequestOutput):
507
- # append_output is called for each output token in streaming case,
508
- # so we only want to add the prompt tokens once for each message.
509
- if self.first_tok_of_message:
510
- self._update_prefill_token_usage(output)
511
- # Reset self.first_tok_of_message if needed:
512
- # if the current token is the last one of the current message
513
- # (finished=True), then the next token processed will mark the
514
- # beginning of a new message
515
- self.first_tok_of_message = output.finished
516
- for tok in output.outputs[0].token_ids:
517
- self.parser.process(tok)
518
- self._update_decode_token_usage(output)
519
-
520
- # For streaming, update previous turn when message is complete
521
- if output.finished:
522
- self.all_turn_metrics.append(self.current_turn_metrics.copy())
523
- self.current_turn_metrics.reset()
524
- # Check if the current token is part of reasoning content
525
- self._update_num_reasoning_tokens()
526
- self.last_tok = tok
527
- if len(self._messages) - self.num_init_messages < len(self.parser.messages):
528
- self._messages.extend(
529
- self.parser.messages[len(self._messages) - self.num_init_messages :]
530
- )
531
- else:
532
- # Handle the case of tool output in direct message format
533
- assert len(output) == 1, "Tool output should be a single message"
534
- msg = output[0]
535
- # Sometimes the recipient is not set for tool messages,
536
- # so we set it to "assistant"
537
- if msg.author.role == Role.TOOL and msg.recipient is None:
538
- msg.recipient = "assistant"
539
- toks = self.encoding.render(msg)
540
- for tok in toks:
541
- self.parser.process(tok)
542
- self.last_tok = toks[-1]
543
- # TODO: add tool_output messages to self._messages
512
+ def append_output(self, output: RequestOutput) -> None:
513
+ # append_output is called for each output token in streaming case,
514
+ # so we only want to add the prompt tokens once for each message.
515
+ if self.first_tok_of_message:
516
+ self._update_prefill_token_usage(output)
517
+ # Reset self.first_tok_of_message if needed:
518
+ # if the current token is the last one of the current message
519
+ # (finished=True), then the next token processed will mark the
520
+ # beginning of a new message
521
+ self.first_tok_of_message = output.finished
522
+ for tok in output.outputs[0].token_ids:
523
+ self.parser.process(tok)
524
+ self._update_decode_token_usage(output)
525
+
526
+ # For streaming, update previous turn when message is complete
527
+ if output.finished:
528
+ self.all_turn_metrics.append(self.current_turn_metrics.copy())
529
+ self.current_turn_metrics.reset()
530
+ # Check if the current token is part of reasoning content
531
+ self._update_num_reasoning_tokens()
532
+ self.last_tok = tok
533
+ if len(self._messages) - self.num_init_messages < len(self.parser.messages):
534
+ self._messages.extend(
535
+ self.parser.messages[len(self._messages) - self.num_init_messages :]
536
+ )
537
+
538
+ def append_tool_output(self, output: list[Message]) -> None:
539
+ # Handle the case of tool output in direct message format
540
+ assert len(output) == 1, "Tool output should be a single message"
541
+ msg = output[0]
542
+ # Sometimes the recipient is not set for tool messages,
543
+ # so we set it to "assistant"
544
+ if msg.author.role == Role.TOOL and msg.recipient is None:
545
+ msg.recipient = "assistant"
546
+ toks = self.encoding.render(msg)
547
+ for tok in toks:
548
+ self.parser.process(tok)
549
+ self.last_tok = toks[-1]
550
+ # TODO: add tool_output messages to self._messages
551
552
def is_expecting_start(self) -> bool:
553
return self.parser.state == StreamState.EXPECT_START
vllm/entrypoints/openai/serving_engine.py
+1
-1
@@ -1227,7 +1227,7 @@ class OpenAIServing:
1227
1228
# Call the tool and update the context with the result.
1229
tool_output = await context.call_tool()
1230
- context.append_output(tool_output)
1230
+ context.append_tool_output(tool_output)
1231
1232
# TODO: uncomment this and enable tool output streaming
1233
# yield context