Skip to content

Commit 7e8965d

Browse files
allen-stephencopybara-github
authored andcommitted
fix(live): propagate output token count in live API usage metadata
Merge #6064 **Please ensure you have read the [contribution guide](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) before creating a pull request.** ### Link to Issue or Description of Change **Problem:** Live API output token count in `usage_metadata` is always `0` (input/total are correct), so everything reading `candidates_token_count`reports zero output tokens for live sessions. Root cause: live usage is a `types.UsageMetadata` (`response_token_count` / `response_tokens_details`), but `LlmResponse.usage_metadata` is a `types.GenerateContentResponseUsageMetadata` (`candidates_token_count` / `candidates_tokens_details`). The raw object was assigned directly; Pydantic re-validated into the target type and silently dropped the two renamed fields. **Solution:** Remap the live usage into `GenerateContentResponseUsageMetadata` in `gemini_llm_connection.receive()` via a `_to_generate_content_usage_metadata()` helper. Updated related tests. ### Testing Plan **Unit Tests:** - [X] I have added or updated unit tests for my change. - [X] All unit tests pass locally. $ uv run pytest tests/unittests/models/test_gemini_llm_connection.py -q 45 passed, 17 warnings in 1.25s **Manual End-to-End (E2E) Tests:** `uv run adk web contributing/samples/live --log_level DEBUG` → `live_bidi_streaming_single_agent`, audio stream. Confirmed live response's `usage_metadata.candidates_token_count` is now non-zero (was `0`), input/total unchanged. ### Checklist - [X] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document. - [X] I have performed a self-review of my own code. - [X] I have commented my code, particularly in hard-to-understand areas. - [X] I have added tests that prove my fix is effective or that my feature works. - [X] New and existing unit tests pass locally with my changes. - [X] I have manually tested my changes end-to-end. - [X] Any dependent changes have been merged and published in downstream modules. ### Additional context N/A COPYBARA_INTEGRATE_REVIEW=#6064 from allen-stephen:fix/live-output-tokens 7ae8c29 PiperOrigin-RevId: 930500943
1 parent 792775f commit 7e8965d

2 files changed

Lines changed: 99 additions & 4 deletions

File tree

src/google/adk/models/gemini_llm_connection.py

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -216,6 +216,35 @@ def __build_full_text_response(
216216
live_session_id=self._gemini_session.session_id,
217217
)
218218

219+
def _to_generate_content_usage_metadata(
220+
self, usage_metadata: types.UsageMetadata
221+
) -> types.GenerateContentResponseUsageMetadata:
222+
"""Converts live API usage metadata to GenerateContentResponse usage metadata.
223+
224+
The live API names output tokens `response_token_count`/
225+
`response_tokens_details`, whereas `GenerateContentResponseUsageMetadata`
226+
names them `candidates_token_count`/`candidates_tokens_details`.
227+
228+
Args:
229+
usage_metadata: The live API usage metadata.
230+
231+
Returns:
232+
The converted usage metadata.
233+
"""
234+
return types.GenerateContentResponseUsageMetadata(
235+
prompt_token_count=usage_metadata.prompt_token_count,
236+
cached_content_token_count=usage_metadata.cached_content_token_count,
237+
candidates_token_count=usage_metadata.response_token_count,
238+
total_token_count=usage_metadata.total_token_count,
239+
thoughts_token_count=usage_metadata.thoughts_token_count,
240+
tool_use_prompt_token_count=usage_metadata.tool_use_prompt_token_count,
241+
prompt_tokens_details=usage_metadata.prompt_tokens_details,
242+
cache_tokens_details=usage_metadata.cache_tokens_details,
243+
candidates_tokens_details=usage_metadata.response_tokens_details,
244+
tool_use_prompt_tokens_details=usage_metadata.tool_use_prompt_tokens_details,
245+
traffic_type=usage_metadata.traffic_type,
246+
)
247+
219248
async def receive(self) -> AsyncGenerator[LlmResponse, None]:
220249
"""Receives the model response using the llm server connection.
221250
@@ -234,9 +263,11 @@ async def receive(self) -> AsyncGenerator[LlmResponse, None]:
234263
logger.debug('Got LLM Live message: %s', message)
235264
live_session_id = self._gemini_session.session_id
236265
if message.usage_metadata:
237-
# Tracks token usage data per model.
266+
# Remap live token usage to GenerateContentResponse usage metadata.
238267
yield LlmResponse(
239-
usage_metadata=message.usage_metadata,
268+
usage_metadata=self._to_generate_content_usage_metadata(
269+
message.usage_metadata
270+
),
240271
model_version=self._model_version,
241272
live_session_id=live_session_id,
242273
)

tests/unittests/models/test_gemini_llm_connection.py

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -231,10 +231,12 @@ async def mock_receive_generator():
231231
content_response = next((r for r in responses if r.content), None)
232232
assert content_response is not None
233233

234+
# The live API's `response_token_count`/`response_tokens_details` are remapped
235+
# to `candidates_token_count`/`candidates_tokens_details`.
234236
expected_usage = types.GenerateContentResponseUsageMetadata(
235237
prompt_token_count=10,
236238
cached_content_token_count=5,
237-
candidates_token_count=None,
239+
candidates_token_count=20,
238240
total_token_count=35,
239241
thoughts_token_count=2,
240242
prompt_tokens_details=[
@@ -243,12 +245,74 @@ async def mock_receive_generator():
243245
cache_tokens_details=[
244246
types.ModalityTokenCount(modality='text', token_count=5)
245247
],
246-
candidates_tokens_details=None,
248+
candidates_tokens_details=[
249+
types.ModalityTokenCount(modality='text', token_count=20)
250+
],
247251
)
248252
assert usage_response.usage_metadata == expected_usage
249253
assert content_response.content == mock_content
250254

251255

256+
async def test_receive_usage_metadata_remaps_output_tokens(
257+
gemini_connection, mock_gemini_session
258+
):
259+
"""Test that live API output tokens are remapped to candidates_token_count."""
260+
usage_metadata = types.UsageMetadata(
261+
prompt_token_count=10,
262+
cached_content_token_count=5,
263+
response_token_count=20,
264+
total_token_count=35,
265+
thoughts_token_count=2,
266+
tool_use_prompt_token_count=3,
267+
prompt_tokens_details=[
268+
types.ModalityTokenCount(modality='text', token_count=10)
269+
],
270+
cache_tokens_details=[
271+
types.ModalityTokenCount(modality='text', token_count=5)
272+
],
273+
response_tokens_details=[
274+
types.ModalityTokenCount(modality='text', token_count=20)
275+
],
276+
)
277+
278+
mock_message = mock.AsyncMock()
279+
mock_message.usage_metadata = usage_metadata
280+
mock_message.server_content = None
281+
mock_message.tool_call = None
282+
mock_message.session_resumption_update = None
283+
mock_message.go_away = None
284+
285+
async def mock_receive_generator():
286+
yield mock_message
287+
288+
receive_mock = mock.Mock(return_value=mock_receive_generator())
289+
mock_gemini_session.receive = receive_mock
290+
291+
responses = [resp async for resp in gemini_connection.receive()]
292+
293+
usage_response = next((r for r in responses if r.usage_metadata), None)
294+
assert usage_response is not None
295+
result = usage_response.usage_metadata
296+
assert isinstance(result, types.GenerateContentResponseUsageMetadata)
297+
# Output tokens are remapped from response_* to candidates_*.
298+
assert result.candidates_token_count == 20
299+
assert result.candidates_tokens_details == [
300+
types.ModalityTokenCount(modality='text', token_count=20)
301+
]
302+
# Shared fields are carried over unchanged.
303+
assert result.prompt_token_count == 10
304+
assert result.cached_content_token_count == 5
305+
assert result.total_token_count == 35
306+
assert result.thoughts_token_count == 2
307+
assert result.tool_use_prompt_token_count == 3
308+
assert result.prompt_tokens_details == [
309+
types.ModalityTokenCount(modality='text', token_count=10)
310+
]
311+
assert result.cache_tokens_details == [
312+
types.ModalityTokenCount(modality='text', token_count=5)
313+
]
314+
315+
252316
async def test_receive_populates_live_session_id(
253317
gemini_connection, mock_gemini_session
254318
):

0 commit comments

Comments
 (0)