Fix Google Generative AI token usage statistics tracing (#169222)

homeassistant/components/google_generative_ai_conversation/entity.py

@@ -338,6 +338,7 @@ def _convert_content(
 
 
 async def _transform_stream(
+    chat_log: conversation.ChatLog,
     result: AsyncIterator[GenerateContentResponse],
 ) -> AsyncGenerator[conversation.AssistantContentDeltaDict]:
     new_message = True
@@ -346,6 +347,19 @@ async def _transform_stream(
         async for response in result:
             LOGGER.debug("Received response chunk: %s", response)
 
+            if (usage := response.usage_metadata) is not None:
+                chat_log.async_trace(
+                    {
+                        "stats": {
+                            "input_tokens": usage.prompt_token_count,
+                            "cached_input_tokens": (
+                                usage.cached_content_token_count or 0
+                            ),
+                            "output_tokens": usage.candidates_token_count,
+                        }
+                    }
+                )
+
             if new_message:
                 if part_details:
                     yield {"native": ContentDetails(part_details=part_details)}
@@ -623,7 +637,7 @@ class GoogleGenerativeAILLMBaseEntity(Entity):
                         content
                         async for content in chat_log.async_add_delta_content_stream(
                             self.entity_id,
-                            _transform_stream(chat_response_generator),
+                            _transform_stream(chat_log, chat_response_generator),
                         )
                         if isinstance(content, conversation.ToolResultContent)
                     ]

tests/components/google_generative_ai_conversation/test_conversation.py

@@ -13,6 +13,7 @@ from homeassistant.components.conversation import (
     AssistantContent,
     ToolResultContent,
     UserContent,
+    trace,
 )
 from homeassistant.components.google_generative_ai_conversation.entity import (
     ERROR_GETTING_RESPONSE,
@@ -795,3 +796,72 @@ async def test_history_always_user_first_turn(
         == "Garage door left open, do you want to close it?"
     )
     assert actual_history[1].role == "model"
+
+
+@pytest.mark.usefixtures("mock_init_component")
+async def test_token_stats_reported(
+    hass: HomeAssistant,
+    mock_config_entry: MockConfigEntry,
+    mock_chat_log: MockChatLog,  # noqa: F811
+    mock_send_message_stream: AsyncMock,
+) -> None:
+    """Test that token stats are reported to the chat log."""
+    trace.async_clear_traces()
+
+    agent_id = "conversation.google_ai_conversation"
+    context = Context()
+
+    messages = [
+        [
+            GenerateContentResponse(
+                candidates=[
+                    {
+                        "content": {
+                            "parts": [{"text": "Hello! "}],
+                            "role": "model",
+                        },
+                    }
+                ],
+            ),
+            GenerateContentResponse(
+                candidates=[
+                    {
+                        "content": {
+                            "parts": [{"text": "How can I help you?"}],
+                            "role": "model",
+                        },
+                        "finish_reason": "STOP",
+                    }
+                ],
+                usage_metadata={
+                    "prompt_token_count": 10,
+                    "candidates_token_count": 20,
+                    "cached_content_token_count": 5,
+                },
+            ),
+        ],
+    ]
+
+    mock_send_message_stream.return_value = messages
+
+    await conversation.async_converse(
+        hass,
+        "Hello",
+        mock_chat_log.conversation_id,
+        context,
+        agent_id=agent_id,
+    )
+
+    traces = trace.async_get_traces()
+    trace_obj = next(iter(traces))
+    events = trace_obj.as_dict().get("events", [])
+    stats = next(
+        e["data"]["stats"]
+        for e in events
+        if e.get("event_type") == "agent_detail" and e.get("data", {}).get("stats")
+    )
+    assert stats == {
+        "input_tokens": 10,
+        "cached_input_tokens": 5,
+        "output_tokens": 20,
+    }