From b051aa713a3d4c044fb9453935879f09d91eea69 Mon Sep 17 00:00:00 2001 From: Marty Sullivan Date: Mon, 7 Sep 2026 04:54:13 -0400 Subject: [PATCH] fix(vertex-live): sum tool-use prompt tokens across a session's turns toolUsePromptTokenCount was the one prompt-side total not named in _AGGREGATED_FIELDS, so it rode the unknown-key pass-through and took the first frame's value while promptTokenCount, candidatesTokenCount and totalTokenCount beside it were summed. Live's frames grow over a session, so the first frame is the smallest number in the series and a grounded session under-reported its tool-use tokens by everything after turn one. It is now summed like its three neighbours. This is reporting only, and pricing these tokens is deliberately left out. Google charges tool-use prompt tokens at the input token rate, but generic_cost_per_token reads the input bill out of prompt_tokens_details and only falls back to prompt_tokens when the details carry no text or a cache hit overlaps them. Measured on the native-audio entry with 500 tool-use tokens: adding them to prompt_tokens moves an ordinary Live turn's bill by $0.0000000000, and on a turn with a cache hit it moves it by $0.0002650000 where the tokens are worth $0.0002500000, because it perturbs the cache-overlap correction. Pricing them belongs beside the modality terms in the shared input-cost path, in its own change that fixes the same latent no-op on the ordinary Gemini path. Not verified against a live capture: no Vertex Live session we have captured reported toolUsePromptTokenCount at all, so the summing convention is inferred from the three prompt-side totals that accumulate the same way. --- ...tex_ai_live_passthrough_logging_handler.py | 2 + .../test_vertex_ai_live_passthrough.py | 46 ++++++++++++++----- 2 files changed, 37 insertions(+), 11 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 72e8c3f7657..2cd336db3d6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -31,6 +31,7 @@ _AGGREGATED_FIELDS: Final = frozenset( "promptTokenCount", "candidatesTokenCount", "totalTokenCount", + "toolUsePromptTokenCount", "promptTokensDetails", "candidatesTokensDetails", } @@ -159,6 +160,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler): "promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots), "candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots), "totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots), + "toolUsePromptTokenCount": sum(snapshot.get("toolUsePromptTokenCount", 0) for snapshot in snapshots), "promptTokensDetails": [ {"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0 ], diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index ba9167e2b13..bfff52de673 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -363,22 +363,46 @@ class TestVertexAILivePassthroughLoggingHandler: cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) - def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj): - """toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used. + TOOL_USE_PER_TURN = (100, 250, 400) - It is reported, not billed. Nothing in the shared Gemini input-cost path prices - tool-use tokens, and folding them into prompt_tokens here would suppress that - path's cache-overlap correction and raise the bill instead. + def _grounded_messages(self): + """The three-turn session again, with each turn's own toolUsePromptTokenCount attached.""" + messages = self._live_messages(self.AUDIO_SESSION[:3]) + head, turns = messages[0], messages[1:] + return [head] + [ + {**message, "usageMetadata": {**message["usageMetadata"], "toolUsePromptTokenCount": tool_use}} + for message, tool_use in zip(turns, self.TOOL_USE_PER_TURN) + ] + + def test_server_side_tool_use_prompt_tokens_are_summed_over_the_session(self, handler, mock_logging_obj): + """toolUsePromptTokenCount rode the unknown-key pass-through, so it took the first turn only. + + Every other total beside it is summed across the session, and the first turn is the + smallest number in the series, so a grounded session logged far fewer tool-use tokens + than it used. This session's turns are deliberately distinct, so 750 can only come from + summing: first-turn selection gives 100, last-turn or max gives 400. """ - messages = self._live_messages(self.AUDIO_SESSION[:1]) - grounded = [dict(message) for message in messages] - grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500} + grounded = self._grounded_messages() usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) - assert usage.prompt_tokens_details.tool_use_tokens == 500 + assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN) - plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL) - grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL) + def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj): + """Deliberate boundary: these tokens are reported here, and priced nowhere. + + generic_cost_per_token reads the input bill out of prompt_tokens_details, and falls + back to prompt_tokens only when the details carry no text or a cache hit overlaps them, + so adding tool-use tokens to prompt_tokens is worth nothing on an ordinary Live turn and + over-charges against the cache-overlap correction when it is not. Pricing them belongs + in the shared input-cost path, beside the modality terms that already read the details. + """ + turns = self.AUDIO_SESSION[:3] + plain_cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL) + grounded_cost = self._session_cost( + handler, mock_logging_obj, self._grounded_messages(), self.NATIVE_AUDIO_MODEL + ) + + assert plain_cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9) assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill" def test_a_malformed_details_entry_does_not_cost_the_whole_session(self, handler, mock_logging_obj):