From 22acf979d4f07173961aefa748dbdfec4bdc9f0a Mon Sep 17 00:00:00 2001 From: Daniel Peng Date: Sun, 18 Oct 2026 01:43:26 +0800 Subject: [PATCH 1/2] fix(gemini): include thinking tokens in completion usage --- ag2/ag_ui/usage.py | 26 ++------ ag2/config/gemini/mappers.py | 7 +- test/config/gemini/test_ag_ui_usage.py | 2 +- test/config/gemini/test_gemini_usage.py | 85 ++++++++++++++++++++++++- 4 files changed, 96 insertions(+), 24 deletions(-) diff --git a/ag2/ag_ui/usage.py b/ag2/ag_ui/usage.py index 4bc0252fe2d5..22a95aec8586 100644 --- a/ag2/ag_ui/usage.py +++ b/ag2/ag_ui/usage.py @@ -63,7 +63,7 @@ def _with_total(entry: TokenUsage) -> TokenUsage: def _token_usage(record: UsageRecord) -> TokenUsage: usage = record.usage input_tokens = _token_count(_input_total(record.provider, usage)) - output_tokens = _token_count(_output_total(record.provider, usage)) + output_tokens = _token_count(usage.completion_tokens) return TokenUsage( provider=record.provider, model=record.model, @@ -78,14 +78,11 @@ def _token_usage(record: UsageRecord) -> TokenUsage: ) -# The cache and reasoning corrections to AG-UI 1.0's accounting are made here, -# where usage leaves for the wire. `Usage` keeps those counts as the provider -# reported them, because budgets and limiters read them that way. (Gemini's -# tool-use prompt tokens are the exception: they are billed as prompt, and -# `normalize_usage` folds them into `prompt_tokens` itself.) -# AG-UI's input and output are totals, and its cache and reasoning counts parts -# of them, so where a provider reports those beside a smaller count, they are -# added in here. +# Cache corrections to AG-UI 1.0's accounting are made here, where usage leaves +# for the wire. `Usage` keeps cache counts as the provider reported them. +# Completion already includes reasoning tokens, normalized by the provider mapper. +# AG-UI's input is a total, so cache counts reported beside a smaller prompt +# count are added in here. # Providers whose prompt count leaves out the tokens read from and written to # the cache. Bedrock's Converse counts like Anthropic: its prompt-caching guide @@ -93,11 +90,6 @@ def _token_usage(record: UsageRecord) -> TokenUsage: # `cacheWriteInputTokens`. _CACHE_OUTSIDE_PROMPT = frozenset({"anthropic", "bedrock"}) -# Providers whose completion count leaves out the reasoning tokens: -# Gemini's `thoughts_token_count` sits beside `candidates_token_count`. xAI's -# reasoning count is not known to do either, so it is left as reported. -_REASONING_OUTSIDE_COMPLETION = frozenset({"google"}) - def _input_total(provider: str | None, usage: Usage) -> float | None: if usage.prompt_tokens is None or provider not in _CACHE_OUTSIDE_PROMPT: @@ -105,12 +97,6 @@ def _input_total(provider: str | None, usage: Usage) -> float | None: return usage.prompt_tokens + (usage.cache_read_input_tokens or 0) + (usage.cache_creation_input_tokens or 0) -def _output_total(provider: str | None, usage: Usage) -> float | None: - if usage.completion_tokens is None or provider not in _REASONING_OUTSIDE_COMPLETION: - return usage.completion_tokens - return usage.completion_tokens + (usage.thinking_tokens or 0) - - def _token_count(value: float | None) -> int | None: # The wire type admits only non-negative integers, and this runs on the # failure path while the run is being reported — so a value the wire would diff --git a/ag2/config/gemini/mappers.py b/ag2/config/gemini/mappers.py index f2eddd69d133..5a84ad00b283 100644 --- a/ag2/config/gemini/mappers.py +++ b/ag2/config/gemini/mappers.py @@ -397,9 +397,14 @@ def normalize_usage(metadata: types.GenerateContentResponseUsageMetadata) -> Usa tool_use_prompt = _to_float(getattr(metadata, "tool_use_prompt_token_count", None)) if tool_use_prompt is not None: prompt = (prompt or 0.0) + tool_use_prompt + completion = _to_float(metadata.candidates_token_count) + # Gemini reports thoughts beside candidates; completion includes both, + # with thinking_tokens retained as the reasoning breakdown. + if thinking is not None: + completion = (completion or 0.0) + thinking return Usage( prompt_tokens=prompt, - completion_tokens=_to_float(metadata.candidates_token_count), + completion_tokens=completion, total_tokens=_to_float(metadata.total_token_count), cache_read_input_tokens=cache_read, thinking_tokens=thinking, diff --git a/test/config/gemini/test_ag_ui_usage.py b/test/config/gemini/test_ag_ui_usage.py index 526a476e9442..e4ae92ee530f 100644 --- a/test/config/gemini/test_ag_ui_usage.py +++ b/test/config/gemini/test_ag_ui_usage.py @@ -54,7 +54,7 @@ async def _finished(agent: Agent) -> RunFinishedEvent: async def test_the_output_a_client_sees_counts_what_gemini_spent_thinking() -> None: - """Gemini reports its thoughts beside the candidates, so they are added in on the way out.""" + """Gemini thoughts are already in normalized completion and must not be added twice.""" agent = _agent({ "promptTokenCount": 12, "candidatesTokenCount": 10, diff --git a/test/config/gemini/test_gemini_usage.py b/test/config/gemini/test_gemini_usage.py index 43ce689d67b6..1f3e86c07cfb 100644 --- a/test/config/gemini/test_gemini_usage.py +++ b/test/config/gemini/test_gemini_usage.py @@ -6,6 +6,12 @@ from unittest.mock import MagicMock +import httpx +import pytest +from google.genai import types + +from ag2 import Agent +from ag2.config.gemini import GeminiConfig from ag2.config.gemini.mappers import normalize_usage from ag2.events import Usage @@ -69,7 +75,7 @@ def test_includes_thinking_tokens(self): result = normalize_usage(_make_metadata(thoughts=296, total=416)) assert result == Usage( prompt_tokens=100, - completion_tokens=20, + completion_tokens=316, total_tokens=416, thinking_tokens=296, ) @@ -93,4 +99,79 @@ def test_handles_metadata_without_thoughts_field(self): m.cached_content_token_count = None result = normalize_usage(m) - assert result.thinking_tokens is None + assert result == Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120) + + +@pytest.mark.parametrize( + ("candidates", "thoughts", "completion"), + [(5, 40, 45), (None, 40, 40), (0, 40, 40), (5, None, 5), (5, 0, 5), (0, 0, 0), (None, 0, None), (None, None, None)], +) +def test_real_sdk_completion_includes_thinking( + candidates: int | None, thoughts: int | None, completion: int | None +) -> None: + metadata = types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, + candidates_token_count=candidates, + thoughts_token_count=thoughts, + total_token_count=55, + ) + + assert normalize_usage(metadata) == Usage( + prompt_tokens=10, + completion_tokens=completion, + total_tokens=55, + thinking_tokens=thoughts or None, + ) + + +def _usage_response(request: httpx.Request) -> httpx.Response: + response = types.GenerateContentResponse( + candidates=[types.Candidate(content=types.Content(role="model", parts=[types.Part(text="done")]))], + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, + candidates_token_count=5, + thoughts_token_count=40, + total_token_count=55, + ), + ) + if request.url.path.endswith(":streamGenerateContent"): + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=f"data: {response.model_dump_json(by_alias=True)}\n\n", + ) + return httpx.Response(200, json=response.model_dump(by_alias=True, exclude_none=True)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_agent_usage_counts_thinking_as_completion(streaming: bool) -> None: + async with httpx.AsyncClient(transport=httpx.MockTransport(_usage_response)) as http_client: + agent = Agent( + "test_agent", + config=GeminiConfig(model="gemini-3-flash", api_key="test", http_client=http_client, streaming=streaming), + ) + reply = await agent.ask("Hello") + + assert (await reply.usage()).total == Usage( + prompt_tokens=10, completion_tokens=45, total_tokens=55, thinking_tokens=40 + ) + + +def test_real_sdk_thinking_preserves_cache_and_tool_prompt_counts() -> None: + metadata = types.GenerateContentResponseUsageMetadata( + prompt_token_count=10, + candidates_token_count=5, + thoughts_token_count=40, + total_token_count=62, + tool_use_prompt_token_count=7, + cached_content_token_count=8, + ) + + assert normalize_usage(metadata) == Usage( + prompt_tokens=17, + completion_tokens=45, + total_tokens=62, + thinking_tokens=40, + cache_read_input_tokens=8, + ) From 7d23d35967ec4acdaf982f704737c9be82db3861 Mon Sep 17 00:00:00 2001 From: Semen Frolov Date: Sat, 24 Oct 2026 13:41:47 +0300 Subject: [PATCH 2/2] fix(xai): include reasoning tokens in completion usage xAI reports reasoning_tokens beside completion_tokens (total = prompt + completion + reasoning). Fold them into Usage.completion_tokens and keep thinking_tokens as the breakdown, matching the OpenAI, Anthropic and Gemini mappers. --- ag2/config/xai/mappers.py | 8 +++++++- test/ag_ui/test_usage.py | 4 ++-- test/config/xai/test_xai_usage.py | 18 +++++++++++++++++- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/ag2/config/xai/mappers.py b/ag2/config/xai/mappers.py index 5557f327c0d2..8acfe771dc96 100644 --- a/ag2/config/xai/mappers.py +++ b/ag2/config/xai/mappers.py @@ -313,9 +313,15 @@ def normalize_usage(usage: usage_pb2.SamplingUsage | None) -> Usage: if usage is None: return Usage() + # xAI reports reasoning beside the completion (total = prompt + completion + + # reasoning); completion includes it, with thinking_tokens as the breakdown. + completion = float(usage.completion_tokens) if usage.completion_tokens else None + if usage.reasoning_tokens: + completion = (completion or 0.0) + usage.reasoning_tokens + return Usage( prompt_tokens=float(usage.prompt_tokens) if usage.prompt_tokens else None, - completion_tokens=float(usage.completion_tokens) if usage.completion_tokens else None, + completion_tokens=completion, total_tokens=float(usage.total_tokens) if usage.total_tokens else None, cache_read_input_tokens=float(usage.cached_prompt_text_tokens) if usage.cached_prompt_text_tokens else None, thinking_tokens=float(usage.reasoning_tokens) if usage.reasoning_tokens else None, diff --git a/test/ag_ui/test_usage.py b/test/ag_ui/test_usage.py index 19569c210cad..ecdb2be77dbb 100644 --- a/test/ag_ui/test_usage.py +++ b/test/ag_ui/test_usage.py @@ -629,8 +629,8 @@ async def test_bedrock_cache_is_added_to_its_input(self) -> None: @pytest.mark.parametrize("provider", ["openai", "xai"]) async def test_other_providers_are_reported_as_they_count(self, provider: str) -> None: - """OpenAI counts cache and reasoning inside its totals already. xAI's reasoning - is not known to sit outside them, so it is not added in. + """OpenAI and xAI count reasoning inside the normalized completion already, + so it is not added in again. """ usage = Usage( prompt_tokens=100, diff --git a/test/config/xai/test_xai_usage.py b/test/config/xai/test_xai_usage.py index 99f64606b9c3..1b2932872194 100644 --- a/test/config/xai/test_xai_usage.py +++ b/test/config/xai/test_xai_usage.py @@ -32,13 +32,29 @@ def test_cached_prompt_tokens_map_to_cache_read() -> None: def test_reasoning_tokens_map_to_thinking_tokens() -> None: - raw = usage_pb2.SamplingUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15, reasoning_tokens=3) + raw = usage_pb2.SamplingUsage(prompt_tokens=10, completion_tokens=5, total_tokens=18, reasoning_tokens=3) result = normalize_usage(raw) assert result.thinking_tokens == 3.0 +def test_completion_includes_reasoning_tokens() -> None: + # xAI reports reasoning beside the completion: total = prompt + completion + reasoning. + raw = usage_pb2.SamplingUsage(prompt_tokens=204, completion_tokens=32, total_tokens=536, reasoning_tokens=300) + + result = normalize_usage(raw) + + assert result == Usage(prompt_tokens=204.0, completion_tokens=332.0, total_tokens=536.0, thinking_tokens=300.0) + assert result.prompt_tokens + result.completion_tokens == result.total_tokens + + +def test_reasoning_without_completion_tokens() -> None: + raw = usage_pb2.SamplingUsage(prompt_tokens=10, total_tokens=50, reasoning_tokens=40) + + assert normalize_usage(raw).completion_tokens == 40.0 + + def test_zero_or_missing_fields_become_none() -> None: raw = usage_pb2.SamplingUsage()