@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a3645ed94d
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 14, 05:30 UTC

[Frontend][Responses API] Support reporting tool output tokens and fix reasoning token count (#24285)

Signed-off-by: Ye (Charlotte) Qi <yeq@meta.com>

Ye (Charlotte) Qi committed Sep 6, 2025 at 13:27 UTC a3645ed94d9f587e22a24f13fccd26be9d35d8aa
4 files changed +557 -36
tests/entrypoints/test_context.py new
+425
@@ -0,0 +1,425 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from unittest.mock import MagicMock, patch
5 +
6 +import pytest
7 +from openai_harmony import StreamState
8 +
9 +from vllm.entrypoints.context import HarmonyContext, StreamingHarmonyContext
10 +from vllm.outputs import CompletionOutput, RequestOutput
11 +
12 +
13 +# Helper function for Python < 3.10 compatibility
14 +async def async_next(async_iterator):
15 + """Compatibility function equivalent to Python 3.10's anext()."""
16 + return await async_iterator.__anext__()
17 +
18 +
19 +def create_mock_request_output(
20 + prompt_token_ids=None,
21 + output_token_ids=None,
22 + num_cached_tokens=0,
23 + finished=True,
24 +):
25 + """Helper function to create a mock RequestOutput object for testing."""
26 + outputs = []
27 + token_ids = output_token_ids if output_token_ids is not None else []
28 + outputs = [
29 + CompletionOutput(
30 + index=0,
31 + text="Test output",
32 + token_ids=token_ids,
33 + cumulative_logprob=0.0,
34 + logprobs=None,
35 + finish_reason=None,
36 + stop_reason=None,
37 + )
38 + ]
39 +
40 + return RequestOutput(
41 + request_id="test-id",
42 + prompt="Test prompt",
43 + prompt_token_ids=prompt_token_ids,
44 + prompt_logprobs=None,
45 + outputs=outputs,
46 + finished=finished,
47 + num_cached_tokens=num_cached_tokens,
48 + )
49 +
50 +
51 +async def generate_mock_outputs(num_turns,
52 + prompt_token_counts,
53 + output_token_counts,
54 + cached_token_counts=None):
55 + """Generate a sequence of mock RequestOutput objects to simulate multiple
56 + turns."""
57 + if cached_token_counts is None:
58 + cached_token_counts = [0] * num_turns
59 +
60 + for i in range(num_turns):
61 + # Create mock prompt token IDs and output token IDs
62 + prompt_token_ids = list(range(1, prompt_token_counts[i] + 1))
63 + output_token_ids = list(range(1, output_token_counts[i] + 1))
64 +
65 + # Create and yield the RequestOutput
66 + yield create_mock_request_output(
67 + prompt_token_ids=prompt_token_ids,
68 + output_token_ids=output_token_ids,
69 + num_cached_tokens=cached_token_counts[i],
70 + )
71 +
72 +
73 +@pytest.fixture
74 +def mock_parser():
75 + """Set up a mock parser for tests."""
76 + with patch("vllm.entrypoints.context.get_streamable_parser_for_assistant"
77 + ) as mock_parser_factory:
78 + # Create a mock parser object
79 + parser = MagicMock()
80 + parser.messages = []
81 + parser.current_channel = None
82 + parser.state = StreamState.EXPECT_START
83 + mock_parser_factory.return_value = parser
84 + yield parser
85 +
86 +
87 +def test_single_turn_token_counting():
88 + """Test token counting behavior for a single turn."""
89 + # Create a context
90 + context = HarmonyContext(messages=[], available_tools=[])
91 +
92 + # Create a mock RequestOutput with specific token counts
93 + mock_output = create_mock_request_output(
94 + prompt_token_ids=[1, 2, 3, 4, 5], # 5 prompt tokens
95 + output_token_ids=[6, 7, 8], # 3 output tokens
96 + num_cached_tokens=2, # 2 cached tokens
97 + )
98 +
99 + # Append the output to the context
100 + context.append_output(mock_output)
101 +
102 + # Verify the token counts
103 + assert context.num_prompt_tokens == 5
104 + assert context.num_output_tokens == 3
105 + assert context.num_cached_tokens == 2
106 + assert context.num_tool_output_tokens == 0 # No tool tokens in first turn
107 +
108 + # Verify internal state tracking
109 + assert not context.is_first_turn
110 + assert context.previous_turn.input_tokens == 5
111 + assert context.previous_turn.output_tokens == 3
112 +
113 +
114 +@pytest.mark.asyncio
115 +async def test_multi_turn_token_counting():
116 + """Test token counting behavior across multiple turns with tool output."""
117 + # Create a context
118 + context = HarmonyContext(messages=[], available_tools=["browser"])
119 +
120 + # Simulate a conversation with 3 turns
121 + # Turn 1: prefill 5, decode 3, tool 7
122 + # Turn 2: prefill 15, cached 5, decode 4, tool 1
123 + # Turn 3: prefill 20, cached 15, decode 5
124 + prompt_token_counts = [5, 15, 20]
125 + output_token_counts = [3, 4, 5]
126 + cached_token_counts = [0, 5, 15]
127 + mock_generator = generate_mock_outputs(3, prompt_token_counts,
128 + output_token_counts,
129 + cached_token_counts)
130 +
131 + # First turn - initial prompt and response
132 + mock_output1 = await async_next(mock_generator)
133 + context.append_output(mock_output1)
134 +
135 + # At this point, we should have 5 prompt tokens and 3 output tokens
136 + assert context.num_prompt_tokens == 5
137 + assert context.num_output_tokens == 3
138 + assert context.num_tool_output_tokens == 0
139 +
140 + # Second turn - after tool output
141 + mock_output2 = await async_next(mock_generator)
142 + context.append_output(mock_output2)
143 + # Current prompt tokens (15) - last_turn_input_tokens (5) -
144 + # last_turn_output_tokens (3) = 7
145 + expected_tool_output = 7
146 +
147 + assert context.num_prompt_tokens == 5 + 15
148 + assert context.num_output_tokens == 3 + 4
149 + assert context.num_tool_output_tokens == expected_tool_output
150 + assert context.num_cached_tokens == 5
151 +
152 + # Third turn - final response
153 + mock_output3 = await async_next(mock_generator)
154 + context.append_output(mock_output3)
155 + # Additional tool output tokens from third turn:
156 + # Current prompt (20) - last_turn_input_tokens (15) -
157 + # last_turn_output_tokens (4) = 1
158 + expected_tool_output = 7 + 1
159 +
160 + assert context.num_prompt_tokens == 5 + 15 + 20
161 + assert context.num_output_tokens == 3 + 4 + 5
162 + assert context.num_tool_output_tokens == expected_tool_output
163 + assert context.num_cached_tokens == 5 + 15
164 +
165 +
166 +def test_empty_output_tokens():
167 + """Test behavior when RequestOutput has empty output tokens."""
168 + context = HarmonyContext(messages=[], available_tools=[])
169 +
170 + # Create a RequestOutput with empty output tokens
171 + mock_output = create_mock_request_output(
172 + prompt_token_ids=[1, 2, 3], # 3 prompt tokens
173 + output_token_ids=[], # Empty output tokens list
174 + num_cached_tokens=1,
175 + )
176 +
177 + context.append_output(mock_output)
178 +
179 + # Should handle empty outputs gracefully
180 + assert context.num_prompt_tokens == 3
181 + assert context.num_output_tokens == 0 # No output tokens
182 + assert context.num_cached_tokens == 1
183 + assert context.num_tool_output_tokens == 0
184 +
185 +
186 +def test_missing_prompt_token_ids():
187 + """Test behavior when RequestOutput has None prompt_token_ids."""
188 + context = HarmonyContext(messages=[], available_tools=[])
189 +
190 + mock_output = create_mock_request_output(
191 + prompt_token_ids=None, # No prompt token IDs
192 + output_token_ids=[1, 2], # 2 output tokens
193 + num_cached_tokens=0,
194 + )
195 +
196 + # Logger.error will be called, but we don't need to check for warnings
197 + # here Just ensure it doesn't raise an exception
198 + context.append_output(mock_output)
199 +
200 + # Should handle missing prompt tokens gracefully
201 + assert context.num_prompt_tokens == 0
202 + assert context.num_output_tokens == 2
203 + assert context.num_cached_tokens == 0
204 + assert context.num_tool_output_tokens == 0
205 +
206 +
207 +def test_reasoning_tokens_counting(mock_parser):
208 + """Test that reasoning tokens are counted correctly."""
209 + context = HarmonyContext(messages=[], available_tools=[])
210 +
211 + # Mock parser to simulate reasoning channel
212 + mock_parser.current_channel = "analysis" # Reasoning channel
213 +
214 + mock_output = create_mock_request_output(
215 + prompt_token_ids=[1, 2, 3],
216 + output_token_ids=[4, 5, 6, 7], # 4 tokens, all in reasoning
217 + num_cached_tokens=0,
218 + )
219 +
220 + context.append_output(mock_output)
221 +
222 + # All output tokens should be counted as reasoning
223 + assert context.num_reasoning_tokens == 4
224 + assert context.num_output_tokens == 4
225 +
226 +
227 +def test_zero_tokens_edge_case():
228 + """Test behavior with all zero token counts."""
229 + context = HarmonyContext(messages=[], available_tools=[])
230 +
231 + # Create a request with empty lists (not None) for both prompt and
232 + # output tokens
233 + mock_output = create_mock_request_output(
234 + prompt_token_ids=[], # Empty prompt tokens
235 + output_token_ids=[], # Empty output tokens
236 + num_cached_tokens=0,
237 + )
238 +
239 + context.append_output(mock_output)
240 +
241 + # All counts should be zero
242 + assert context.num_prompt_tokens == 0
243 + assert context.num_output_tokens == 0
244 + assert context.num_cached_tokens == 0
245 + assert context.num_tool_output_tokens == 0
246 + assert context.num_reasoning_tokens == 0
247 +
248 +
249 +@pytest.mark.asyncio
250 +async def test_single_turn_no_tool_output():
251 + """Test that first turn never generates tool output tokens."""
252 + context = HarmonyContext(
253 + messages=[],
254 + available_tools=["browser"] # Tools available
255 + )
256 +
257 + # Even with large prompt in first turn, no tool tokens should be counted
258 + mock_output = create_mock_request_output(
259 + prompt_token_ids=list(range(100)), # 100 tokens
260 + output_token_ids=[1, 2, 3],
261 + num_cached_tokens=0,
262 + )
263 +
264 + context.append_output(mock_output)
265 +
266 + # First turn should never have tool output tokens
267 + assert context.num_tool_output_tokens == 0
268 + assert context.is_first_turn is False # Should be updated after first turn
269 +
270 +
271 +@pytest.mark.asyncio
272 +async def test_negative_tool_tokens_edge_case():
273 + """Test edge case where calculation could result in negative tool
274 + tokens. We should log an error and clamp the value to 0."""
275 + # Use patch to check if logger.error was called
276 + with patch("vllm.entrypoints.context.logger.error") as mock_log:
277 + context = HarmonyContext(messages=[], available_tools=["browser"])
278 +
279 + # First turn
280 + mock_output1 = create_mock_request_output(
281 + prompt_token_ids=list(range(10)), # 10 tokens
282 + output_token_ids=[1, 2, 3, 4, 5], # 5 tokens
283 + )
284 + context.append_output(mock_output1)
285 +
286 + # Second turn with fewer new tokens than previous output
287 + # This could happen in edge cases with aggressive caching
288 + mock_output2 = create_mock_request_output(
289 + prompt_token_ids=list(range(12)), # 12 tokens (only 2 new)
290 + output_token_ids=[6, 7], # 2 tokens
291 + )
292 + context.append_output(mock_output2)
293 +
294 + # Calculated negative tool tokens (12 - 10 - 5 = -3) should be clamped
295 + # to 0 and an error should be logged
296 + assert context.num_tool_output_tokens == 0
297 + assert context.num_prompt_tokens == 10 + 12
298 + assert context.num_output_tokens == 5 + 2
299 +
300 + # Verify the error was logged properly
301 + mock_log.assert_called_once()
302 +
303 + # Extract the actual log message and arguments from the call
304 + args, _ = mock_log.call_args
305 + log_message = args[0]
306 +
307 + # Check for key parts of the message
308 + assert "Negative tool output tokens calculated" in log_message
309 + assert "-3" in str(args) # Check that -3 is in the arguments
310 +
311 +
312 +@pytest.mark.asyncio
313 +async def test_streaming_multi_turn_token_counting(mock_parser):
314 + """Test token counting for streaming multi-turn conversations.
315 +
316 + This test focuses on how StreamingHarmonyContext counts tokens in a
317 + multi-turn conversation with streaming (token-by-token) outputs and
318 + message boundaries.
319 + """
320 + # Create a streaming context
321 + context = StreamingHarmonyContext(messages=[], available_tools=["browser"])
322 +
323 + # Simulate three turns of conversation:
324 + # Turn 1: stream tokens one by one, then finish the message
325 + # Turn 2: new prompt, stream more tokens with a reasoning segment
326 + # Turn 3: new prompt with tool output and cached tokens
327 +
328 + # First turn: 3 tokens streamed one by one
329 + # First token of first turn
330 + context.append_output(
331 + create_mock_request_output(
332 + prompt_token_ids=[1, 2, 3], # 3 prompt tokens
333 + output_token_ids=[101], # Single token
334 + num_cached_tokens=0,
335 + finished=False, # Not end of message yet
336 + ))
337 +
338 + # Second token of first turn
339 + context.append_output(
340 + create_mock_request_output(
341 + output_token_ids=[102],
342 + finished=False,
343 + ))
344 +
345 + # Last token of first turn (finished=True signals end of message)
346 + context.append_output(
347 + create_mock_request_output(
348 + output_token_ids=[103],
349 + finished=True, # End of message
350 + ))
351 +
352 + # Check token counts after first turn
353 + assert context.num_prompt_tokens == 3 # Initial prompt tokens
354 + assert context.num_output_tokens == 3 # Three output tokens
355 + assert context.num_cached_tokens == 0
356 + assert context.num_tool_output_tokens == 0 # No tool output in first turn
357 + assert context.first_tok_of_message is True # Ready for next message
358 +
359 + # Second turn: reasoning tokens in analysis channel
360 + mock_parser.current_channel = "analysis" # Set to reasoning channel
361 +
362 + # First token of second turn
363 + context.append_output(
364 + create_mock_request_output(
365 + prompt_token_ids=[1, 2, 3, 101, 102, 103, 4,
366 + 5], # 8 tokens (includes previous)
367 + output_token_ids=[201],
368 + num_cached_tokens=3, # Some tokens cached
369 + finished=False,
370 + ))
371 +
372 + # More tokens in reasoning channel
373 + context.append_output(
374 + create_mock_request_output(
375 + output_token_ids=[202],
376 + finished=False,
377 + ))
378 +
379 + context.append_output(
380 + create_mock_request_output(
381 + output_token_ids=[203],
382 + finished=True, # End of reasoning message
383 + ))
384 +
385 + # Check counts after second turn (reasoning message)
386 + assert context.num_prompt_tokens == 3 + 8 # Initial + second prompt
387 + assert context.num_output_tokens == 3 + 3 # First turn + second turn
388 + assert context.num_reasoning_tokens == 3 # All tokens in analysis channel
389 + assert context.num_cached_tokens == 3 # Cached tokens from second turn
390 +
391 + # Formula: this turn prompt tokens - last turn prompt - last turn output
392 + expected_tool_tokens = 8 - 3 - 3 # = 2
393 + assert context.num_tool_output_tokens == expected_tool_tokens
394 +
395 + # Third turn: regular output channel
396 + mock_parser.current_channel = "final" # Switch back to regular channel
397 +
398 + # Third turn (with more cached tokens)
399 + context.append_output(
400 + create_mock_request_output(
401 + prompt_token_ids=[
402 + 1, 2, 3, 101, 102, 103, 4, 5, 201, 202, 203, 6, 7
403 + ], # 13 tokens
404 + output_token_ids=[301],
405 + num_cached_tokens=8, # More cached tokens
406 + finished=False,
407 + ))
408 +
409 + context.append_output(
410 + create_mock_request_output(
411 + output_token_ids=[302],
412 + finished=True,
413 + ))
414 +
415 + # Final token counts check
416 + assert context.num_prompt_tokens == 3 + 8 + 13 # All prompts
417 + assert context.num_output_tokens == 3 + 3 + 2 # All outputs
418 + assert context.num_reasoning_tokens == 3 # Unchanged from second turn
419 + assert context.num_cached_tokens == 3 + 8 # Accumulated cached tokens
420 +
421 + # Additional tool tokens from third turn
422 + # Formula: this turn prompt - last turn prompt - last turn output
423 + additional_tool_tokens = 13 - 8 - 3 # = 2
424 + assert context.num_tool_output_tokens == expected_tool_tokens \
425 + + additional_tool_tokens
vllm/entrypoints/context.py
+125 -33
@@ -3,7 +3,6 @@
3 import json
4 import logging
5 from abc import ABC, abstractmethod
6 -from collections.abc import Sequence
6 from contextlib import AsyncExitStack
7 from typing import TYPE_CHECKING, Optional, Union
8
@@ -21,6 +20,23 @@ if TYPE_CHECKING:
20 logger = logging.getLogger(__name__)
21
22
23 +class TurnTokens:
24 + """Tracks token counts for a single conversation turn."""
25 +
26 + def __init__(self, input_tokens=0, output_tokens=0):
27 + self.input_tokens = input_tokens
28 + self.output_tokens = output_tokens
29 +
30 + def reset(self):
31 + """Reset counters for a new turn."""
32 + self.input_tokens = 0
33 + self.output_tokens = 0
34 +
35 + def copy(self):
36 + """Create a copy of this turn's token counts."""
37 + return TurnTokens(self.input_tokens, self.output_tokens)
38 +
39 +
40 class ConversationContext(ABC):
41
42 @abstractmethod
@@ -92,52 +108,124 @@ class HarmonyContext(ConversationContext):
108 self.num_init_messages = len(messages)
109 self.num_prompt_tokens = 0
110 self.num_output_tokens = 0
95 - # TODO(woosuk): Implement the following fields.
111 self.num_cached_tokens = 0
112 self.num_reasoning_tokens = 0
113 + self.num_tool_output_tokens = 0
114
99 - def _update_num_prompt_tokens(self, output: RequestOutput):
100 - if output.prompt_token_ids and len(output.prompt_token_ids) > 0:
101 - # NOTE: with built-in tools, there might be multiple rounds in
102 - # the conversation, with the full conversation being resent
103 - # as new prompt each time. Hence the sum.
104 - self.num_prompt_tokens += len(output.prompt_token_ids)
105 -
106 - def _update_num_cached_tokens(self, output: RequestOutput):
107 - if output.num_cached_tokens is not None:
108 - #Similar to num_prompt_tokens
109 - self.num_cached_tokens += output.num_cached_tokens
110 -
111 - def _update_num_output_tokens(self, token_ids: Sequence[int]):
112 - self.num_output_tokens += len(token_ids)
115 + # Turn tracking - replaces multiple individual tracking variables
116 + self.current_turn = TurnTokens()
117 + self.previous_turn = TurnTokens()
118 + self.is_first_turn = True
119 + self.first_tok_of_message = True # For streaming support
120
114 - def _update_num_reasoning_tokens(self, token_ids: Sequence[int]):
115 - # Count tokens that are part of reasoning content (analysis channel
116 - # or tool-directed messages like python/browser calls)
117 - is_analysis = self.parser.current_channel == "analysis"
118 - is_tool_call = (self.parser.current_recipient is not None and
119 - (self.parser.current_recipient.startswith("python") or
120 - self.parser.current_recipient.startswith("browser.")))
121 - if is_analysis or is_tool_call:
122 - self.num_reasoning_tokens += len(token_ids)
121 + def _update_num_reasoning_tokens(self):
122 + # Count all analysis and commentary channels as reasoning tokens
123 + if self.parser.current_channel in {"analysis", "commentary"}:
124 + self.num_reasoning_tokens += 1
125
126 def append_output(self, output) -> None:
127 if isinstance(output, RequestOutput):
126 - self._update_num_prompt_tokens(output)
127 - self._update_num_cached_tokens(output)
128 output_token_ids = output.outputs[0].token_ids
129 - self._update_num_output_tokens(output_token_ids)
129 self.parser = get_streamable_parser_for_assistant()
130 for token_id in output_token_ids:
131 self.parser.process(token_id)
132 # Check if the current token is part of reasoning content
134 - self._update_num_reasoning_tokens([token_id])
133 + self._update_num_reasoning_tokens()
134 + self._update_prefill_token_usage(output)
135 + # Reset current turn output tokens for this turn
136 + self.current_turn.output_tokens = 0
137 + self._update_decode_token_usage(output)
138 + # Move current turn to previous turn for next turn's calculations
139 + self.previous_turn = self.current_turn.copy()
140 output_msgs = self.parser.messages
141 else:
142 # Tool output.
143 output_msgs = output
144 self._messages.extend(output_msgs)
145
146 + def _update_prefill_token_usage(self, output: RequestOutput) -> None:
147 + """Update token usage statistics for the prefill phase of generation.
148 +
149 + The prefill phase processes the input prompt tokens. This method:
150 + 1. Counts the prompt tokens for this turn
151 + 2. Calculates tool output tokens for multi-turn conversations
152 + 3. Updates cached token counts
153 + 4. Tracks state for next turn calculations
154 +
155 + Tool output tokens are calculated as:
156 + current_prompt_tokens - last_turn_prompt_tokens -
157 + last_turn_output_tokens
158 + This represents tokens added between turns (typically tool responses).
159 +
160 + Args:
161 + output: The RequestOutput containing prompt token information
162 + """
163 + if output.prompt_token_ids is not None:
164 + this_turn_input_tokens = len(output.prompt_token_ids)
165 + else:
166 + this_turn_input_tokens = 0
167 + logger.error(
168 + "RequestOutput appended contains no prompt_token_ids.")
169 +
170 + # Update current turn input tokens
171 + self.current_turn.input_tokens = this_turn_input_tokens
172 + self.num_prompt_tokens += this_turn_input_tokens
173 +
174 + # Calculate tool tokens (except on first turn)
175 + if self.is_first_turn:
176 + self.is_first_turn = False
177 + else:
178 + # start counting tool after first turn
179 + # tool tokens = this turn prefill - last turn prefill -
180 + # last turn decode
181 + this_turn_tool_tokens = (self.current_turn.input_tokens -
182 + self.previous_turn.input_tokens -
183 + self.previous_turn.output_tokens)
184 +
185 + # Handle negative tool token counts (shouldn't happen in normal
186 + # cases)
187 + if this_turn_tool_tokens < 0:
188 + logger.error(
189 + "Negative tool output tokens calculated: %d "
190 + "(current_input=%d, previous_input=%d, "
191 + "previous_output=%d). Setting to 0.",
192 + this_turn_tool_tokens, self.current_turn.input_tokens,
193 + self.previous_turn.input_tokens,
194 + self.previous_turn.output_tokens)
195 + this_turn_tool_tokens = 0
196 +
197 + self.num_tool_output_tokens += this_turn_tool_tokens
198 +
199 + # Update cached tokens
200 + if output.num_cached_tokens is not None:
201 + self.num_cached_tokens += output.num_cached_tokens
202 +
203 + def _update_decode_token_usage(self, output: RequestOutput) -> int:
204 + """Update token usage statistics for the decode phase of generation.
205 +
206 + The decode phase processes the generated output tokens. This method:
207 + 1. Counts output tokens from all completion outputs
208 + 2. Updates the total output token count
209 + 3. Tracks tokens generated in the current turn
210 +
211 + In streaming mode, this is called for each token generated.
212 + In non-streaming mode, this is called once with all output tokens.
213 +
214 + Args:
215 + output: The RequestOutput containing generated token information
216 +
217 + Returns:
218 + int: Number of output tokens processed in this call
219 + """
220 + updated_output_token_count = 0
221 + if output.outputs:
222 + for completion_output in output.outputs:
223 + # only keep last round
224 + updated_output_token_count += len(completion_output.token_ids)
225 + self.num_output_tokens += updated_output_token_count
226 + self.current_turn.output_tokens += updated_output_token_count
227 + return updated_output_token_count
228 +
229 @property
230 def messages(self) -> list:
231 return self._messages
@@ -231,8 +319,8 @@ class StreamingHarmonyContext(HarmonyContext):
319 # append_output is called for each output token in streaming case,
320 # so we only want to add the prompt tokens once for each message.
321 if self.first_tok_of_message:
234 - self._update_num_prompt_tokens(output)
235 - self._update_num_cached_tokens(output)
322 + self._update_prefill_token_usage(output)
323 + self.current_turn.output_tokens = 0
324 # Reset self.first_tok_of_message if needed:
325 # if the current token is the last one of the current message
326 # (finished=True), then the next token processed will mark the
@@ -240,9 +328,13 @@ class StreamingHarmonyContext(HarmonyContext):
328 self.first_tok_of_message = output.finished
329 for tok in output.outputs[0].token_ids:
330 self.parser.process(tok)
243 - self._update_num_output_tokens(output.outputs[0].token_ids)
331 + self._update_decode_token_usage(output)
332 +
333 + # For streaming, update previous turn when message is complete
334 + if output.finished:
335 + self.previous_turn = self.current_turn.copy()
336 # Check if the current token is part of reasoning content
245 - self._update_num_reasoning_tokens(output.outputs[0].token_ids)
337 + self._update_num_reasoning_tokens()
338 self.last_tok = tok
339 else:
340 # Handle the case of tool output in direct message format
vllm/entrypoints/openai/protocol.py
+2 -1
@@ -1841,7 +1841,8 @@ class InputTokensDetails(OpenAIBaseModel):
1841
1842
1843 class OutputTokensDetails(OpenAIBaseModel):
1844 - reasoning_tokens: int
1844 + reasoning_tokens: int = 0
1845 + tool_output_tokens: int = 0
1846
1847
1848 class ResponseUsage(OpenAIBaseModel):
vllm/entrypoints/openai/serving_responses.py
+5 -2
@@ -460,7 +460,7 @@ class OpenAIServingResponses(OpenAIServing):
460 if self.use_harmony:
461 assert isinstance(context, HarmonyContext)
462 output = self._make_response_output_items_with_harmony(context)
463 - # TODO: these are all 0 for now!
463 + num_tool_output_tokens = context.num_tool_output_tokens
464 else:
465 assert isinstance(context, SimpleContext)
466 final_res = context.last_output
@@ -473,6 +473,8 @@ class OpenAIServingResponses(OpenAIServing):
473
474 # Calculate usage.
475 assert final_res.prompt_token_ids is not None
476 + num_tool_output_tokens = 0
477 +
478 assert isinstance(context, (SimpleContext, HarmonyContext))
479 num_prompt_tokens = context.num_prompt_tokens
480 num_generated_tokens = context.num_output_tokens
@@ -486,7 +488,8 @@ class OpenAIServingResponses(OpenAIServing):
488 input_tokens_details=InputTokensDetails(
489 cached_tokens=num_cached_tokens),
490 output_tokens_details=OutputTokensDetails(
489 - reasoning_tokens=num_reasoning_tokens),
491 + reasoning_tokens=num_reasoning_tokens,
492 + tool_output_tokens=num_tool_output_tokens),
493 )
494 response = ResponsesResponse.from_request(
495 request,