From 3740d3d1c36a5e6fc196919e8b35e6d20444ebb4 Mon Sep 17 00:00:00 2001 From: Lalith Vardhan Date: Sat, 31 Oct 2026 09:59:32 -0700 Subject: [PATCH 1/2] Python: include thinking and tool-use prompt tokens in Gemini usage totals GeminiChatClient mapped candidates_token_count to output_token_count and prompt_token_count to input_token_count. Gemini reports thinking tokens (thoughts_token_count) and tool-use prompt tokens (tool_use_prompt_token_count) outside those two fields, so on thinking models output_token_count was smaller than reasoning_output_token_count and input + output did not add up to total_token_count. Fold both into the totals, matching the other connectors and the GenAI semconv (reasoning.output_tokens SHOULD be included in output_tokens). --- .../agent_framework_gemini/_chat_client.py | 15 ++- .../gemini/tests/test_gemini_client.py | 101 +++++++++++++++++- 2 files changed, 110 insertions(+), 6 deletions(-) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index fd7b15449c9..f0961ada6c1 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -1270,16 +1270,21 @@ def _parse_usage(self, usage: types.GenerateContentResponseUsageMetadata | None) if not usage: return None details: UsageDetails = {} + # Gemini reports tool-use prompt tokens and thinking tokens outside prompt_token_count and + # candidates_token_count (total = prompt + tool_use_prompt + candidates + thoughts). Fold them + # in so input/output are full totals, as in the other connectors and the GenAI semconv. + tool_use_prompt_tokens = usage.tool_use_prompt_token_count + thoughts_tokens = usage.thoughts_token_count if (v := usage.prompt_token_count) is not None: - details["input_token_count"] = v - if (v := usage.candidates_token_count) is not None: - details["output_token_count"] = v + details["input_token_count"] = v + (tool_use_prompt_tokens or 0) + if (v := usage.candidates_token_count) is not None or thoughts_tokens is not None: + details["output_token_count"] = (v or 0) + (thoughts_tokens or 0) if (v := usage.total_token_count) is not None: details["total_token_count"] = v if (v := usage.cached_content_token_count) is not None: details["cache_read_input_token_count"] = v - if (v := usage.thoughts_token_count) is not None: - details["reasoning_output_token_count"] = v + if thoughts_tokens is not None: + details["reasoning_output_token_count"] = thoughts_tokens return details or None def _map_finish_reason(self, reason: str | None) -> FinishReasonLiteral | FinishReason | None: diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index c13f00b4bd2..9357516f19a 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -197,6 +197,7 @@ def _make_response( total_tokens: int | None = 15, cached_tokens: int | None = None, thoughts_tokens: int | None = None, + tool_use_prompt_tokens: int | None = None, ) -> MagicMock: """Build a mock types.GenerateContentResponse.""" response = MagicMock() @@ -219,6 +220,7 @@ def _make_response( usage.total_token_count = total_tokens usage.cached_content_token_count = cached_tokens usage.thoughts_token_count = thoughts_tokens + usage.tool_use_prompt_token_count = tool_use_prompt_tokens response.usage_metadata = usage else: response.usage_metadata = None @@ -721,7 +723,7 @@ async def test_get_response_usage_details_includes_cached_and_reasoning_tokens() [_make_part(text="Hi")], prompt_tokens=20, output_tokens=8, - total_tokens=28, + total_tokens=34, cached_tokens=12, thoughts_tokens=6, ) @@ -734,6 +736,74 @@ async def test_get_response_usage_details_includes_cached_and_reasoning_tokens() assert response.usage_details["reasoning_output_token_count"] == 6 +async def test_get_response_output_token_count_includes_thinking_tokens() -> None: + """Gemini reports thinking tokens outside candidates_token_count; output_token_count must include them.""" + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock( + return_value=_make_response( + [_make_part(text="Hi")], + prompt_tokens=1200, + output_tokens=80, + total_tokens=2180, + thoughts_tokens=900, + ) + ) + + response = await client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])]) + + usage = response.usage_details + assert usage is not None + assert usage["output_token_count"] == 980 + assert usage["reasoning_output_token_count"] == 900 + assert usage["input_token_count"] + usage["output_token_count"] == usage["total_token_count"] + + +async def test_get_response_input_token_count_includes_tool_use_prompt_tokens() -> None: + """Tool-use prompt tokens (e.g. Google Search results fed back to the model) count as input.""" + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock( + return_value=_make_response( + [_make_part(text="Hi")], + prompt_tokens=1200, + output_tokens=80, + total_tokens=2530, + cached_tokens=1024, + thoughts_tokens=900, + tool_use_prompt_tokens=350, + ) + ) + + response = await client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])]) + + usage = response.usage_details + assert usage is not None + assert usage["input_token_count"] == 1550 + assert usage["output_token_count"] == 980 + assert usage["total_token_count"] == 2530 + assert usage["cache_read_input_token_count"] == 1024 + + +async def test_get_response_output_token_count_with_only_thinking_tokens() -> None: + """A response cut off while thinking has no candidates_token_count but still produced output tokens.""" + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock( + return_value=_make_response( + [_make_part(text="")], + finish_reason="MAX_TOKENS", + prompt_tokens=10, + output_tokens=None, + total_tokens=110, + thoughts_tokens=100, + ) + ) + + response = await client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])]) + + assert response.usage_details is not None + assert response.usage_details["output_token_count"] == 100 + assert response.usage_details["reasoning_output_token_count"] == 100 + + async def test_get_response_no_usage_when_metadata_absent() -> None: """Returns None for usage_details when the API response includes no usage metadata.""" client, mock = _make_gemini_client() @@ -2870,6 +2940,35 @@ async def test_streaming_get_final_response() -> None: assert final.usage_details["output_token_count"] == 5 +async def test_streaming_usage_includes_thinking_and_tool_use_prompt_tokens() -> None: + """The final streamed usage folds thinking and tool-use prompt tokens into output and input.""" + client, mock = _make_gemini_client() + chunks = [ + _make_response([_make_part(text="Hello ")], finish_reason=None, prompt_tokens=None, output_tokens=None), + _make_response( + [_make_part(text="world!")], + finish_reason="STOP", + prompt_tokens=1200, + output_tokens=80, + total_tokens=2530, + thoughts_tokens=900, + tool_use_prompt_tokens=350, + ), + ] + mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks)) + + stream = client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])], stream=True) + async for _ in stream: + pass + final = await stream.get_final_response() + + assert final.usage_details is not None + assert final.usage_details["input_token_count"] == 1550 + assert final.usage_details["output_token_count"] == 980 + assert final.usage_details["reasoning_output_token_count"] == 900 + assert final.usage_details["total_token_count"] == 2530 + + async def test_streaming_updates_have_assistant_role() -> None: """Streamed updates must carry the assistant role, matching the non-streaming response.""" client, mock = _make_gemini_client() From d968f39e033c8bc061edd1c4330c439ac728dbf3 Mon Sep 17 00:00:00 2001 From: Lalith Vardhan Date: Tue, 3 Nov 2026 09:58:20 -0800 Subject: [PATCH 2/2] Python: count tool-use prompt tokens when Gemini omits prompt_token_count; fix test typing --- .../agent_framework_gemini/_chat_client.py | 4 +-- .../gemini/tests/test_gemini_client.py | 26 ++++++++++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index f0961ada6c1..e513138275a 100644 --- a/python/packages/gemini/agent_framework_gemini/_chat_client.py +++ b/python/packages/gemini/agent_framework_gemini/_chat_client.py @@ -1275,8 +1275,8 @@ def _parse_usage(self, usage: types.GenerateContentResponseUsageMetadata | None) # in so input/output are full totals, as in the other connectors and the GenAI semconv. tool_use_prompt_tokens = usage.tool_use_prompt_token_count thoughts_tokens = usage.thoughts_token_count - if (v := usage.prompt_token_count) is not None: - details["input_token_count"] = v + (tool_use_prompt_tokens or 0) + if (v := usage.prompt_token_count) is not None or tool_use_prompt_tokens is not None: + details["input_token_count"] = (v or 0) + (tool_use_prompt_tokens or 0) if (v := usage.candidates_token_count) is not None or thoughts_tokens is not None: details["output_token_count"] = (v or 0) + (thoughts_tokens or 0) if (v := usage.total_token_count) is not None: diff --git a/python/packages/gemini/tests/test_gemini_client.py b/python/packages/gemini/tests/test_gemini_client.py index 9357516f19a..8957b7fe4c3 100644 --- a/python/packages/gemini/tests/test_gemini_client.py +++ b/python/packages/gemini/tests/test_gemini_client.py @@ -755,7 +755,31 @@ async def test_get_response_output_token_count_includes_thinking_tokens() -> Non assert usage is not None assert usage["output_token_count"] == 980 assert usage["reasoning_output_token_count"] == 900 - assert usage["input_token_count"] + usage["output_token_count"] == usage["total_token_count"] + input_tokens = usage["input_token_count"] + output_tokens = usage["output_token_count"] + assert input_tokens is not None and output_tokens is not None + assert input_tokens + output_tokens == usage["total_token_count"] + + +async def test_get_response_input_token_count_set_when_only_tool_use_prompt_tokens_reported() -> None: + """If prompt_token_count is missing but tool-use prompt tokens are present, input still counts them.""" + client, mock = _make_gemini_client() + mock.aio.models.generate_content = AsyncMock( + return_value=_make_response( + [_make_part(text="Hi")], + prompt_tokens=None, + output_tokens=80, + total_tokens=430, + tool_use_prompt_tokens=350, + ) + ) + + response = await client.get_response(messages=[Message(role="user", contents=[Content.from_text("Hi")])]) + + usage = response.usage_details + assert usage is not None + assert usage["input_token_count"] == 350 + assert usage["output_token_count"] == 80 async def test_get_response_input_token_count_includes_tool_use_prompt_tokens() -> None: