diff --git a/python/packages/gemini/agent_framework_gemini/_chat_client.py b/python/packages/gemini/agent_framework_gemini/_chat_client.py index fd7b15449c9..e513138275a 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 = {} - 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 + # 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 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: 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..8957b7fe4c3 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,98 @@ 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 + 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: + """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 +2964,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()