feat: ensure tokens generated during multi-round built-in tools are summed in the usage statistics

Signed-off-by: Guillaume Calmettes <[email protected]>

Author: gcalmettes

vllm/entrypoints/context.py

@@ -3,6 +3,7 @@
 import json
 import logging
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Union
 
 from openai_harmony import Author, Message, Role, StreamState, TextContent
@@ -66,24 +67,30 @@ def __init__(
         self.tool_sessions = tool_sessions
 
         self.parser = get_streamable_parser_for_assistant()
+        self.encoding = get_encoding()
         self.num_init_messages = len(messages)
         self.num_prompt_tokens = 0
         self.num_output_tokens = 0
         # TODO(woosuk): Implement the following fields.
         self.num_cached_tokens = 0
         self.num_reasoning_tokens = 0
 
-    def _update_prompt_tokens(self, output: RequestOutput):
-        if output.prompt_token_ids and len(
-                output.prompt_token_ids) > 0 and self.num_prompt_tokens == 0:
-            self.num_prompt_tokens = len(output.prompt_token_ids)
+    def _update_num_prompt_tokens(self, output: RequestOutput):
+        if output.prompt_token_ids and len(output.prompt_token_ids) > 0:
+            # NOTE: with built-in tools, there might be multiple rounds in
+            # the conversation, with the full conversation being resent
+            # as new prompt each time. Hence the sum.
+            self.num_prompt_tokens += len(output.prompt_token_ids)
+
+    def _update_num_output_tokens(self, token_ids: Sequence[int]):
+        self.num_output_tokens += len(token_ids)
 
     def append_output(self, output) -> None:
         if isinstance(output, RequestOutput):
-            self._update_prompt_tokens(output)
+            self._update_num_prompt_tokens(output)
             output_token_ids = output.outputs[0].token_ids
+            self._update_num_output_tokens(output_token_ids)
             self.parser = get_streamable_parser_for_assistant()
-            self.num_output_tokens += len(output_token_ids)
             for token_id in output_token_ids:
                 self.parser.process(token_id)
             output_msgs = self.parser.messages
@@ -163,7 +170,6 @@ def __init__(self, *args, **kwargs):
         self.last_output = None
 
         self.parser = get_streamable_parser_for_assistant()
-        self.encoding = get_encoding()
         self.last_tok = None
 
     @property
@@ -172,10 +178,10 @@ def messages(self) -> list:
 
     def append_output(self, output) -> None:
         if isinstance(output, RequestOutput):
-            self._update_prompt_tokens(output)
+            self._update_num_prompt_tokens(output)
             tok = output.outputs[0].token_ids[0]
             self.parser.process(tok)
-            self.num_output_tokens += 1
+            self._update_num_output_tokens(output.outputs[0].token_ids)
             self.last_tok = tok
         else:
             # Handle the case of tool output in direct message format