Extract google model usage using genai-prices (#3466)

Co-authored-by: Douwe Maan <[email protected]>

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -484,7 +484,7 @@ def _process_response(self, response: GenerateContentResponse) -> ModelResponse:
         else:
             parts = candidate.content.parts or []
 
-        usage = _metadata_as_usage(response)
+        usage = _metadata_as_usage(response, provider=self._provider.name, provider_url=self._provider.base_url)
         return _process_response_from_parts(
             parts,
             candidate.grounding_metadata,
@@ -511,6 +511,7 @@ async def _process_streamed_response(
             _response=peekable_response,
             _timestamp=first_chunk.create_time or _utils.now_utc(),
             _provider_name=self._provider.name,
+            _provider_url=self._provider.base_url,
         )
 
     async def _map_messages(
@@ -628,11 +629,12 @@ class GeminiStreamedResponse(StreamedResponse):
     _response: AsyncIterator[GenerateContentResponse]
     _timestamp: datetime
     _provider_name: str
+    _provider_url: str
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:  # noqa: C901
         code_execution_tool_call_id: str | None = None
         async for chunk in self._response:
-            self._usage = _metadata_as_usage(chunk)
+            self._usage = _metadata_as_usage(chunk, self._provider_name, self._provider_url)
 
             if not chunk.candidates:
                 continue  # pragma: no cover
@@ -881,7 +883,7 @@ def _tool_config(function_names: list[str]) -> ToolConfigDict:
     return ToolConfigDict(function_calling_config=function_calling_config)
 
 
-def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
+def _metadata_as_usage(response: GenerateContentResponse, provider: str, provider_url: str) -> usage.RequestUsage:
     metadata = response.usage_metadata
     if metadata is None:
         return usage.RequestUsage()
@@ -895,9 +897,6 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
     if tool_use_prompt_token_count := metadata.tool_use_prompt_token_count:
         details['tool_use_prompt_tokens'] = tool_use_prompt_token_count
 
-    input_audio_tokens = 0
-    output_audio_tokens = 0
-    cache_audio_read_tokens = 0
     for prefix, metadata_details in [
         ('prompt', metadata.prompt_tokens_details),
         ('cache', metadata.cache_tokens_details),
@@ -911,22 +910,12 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
             if not detail.modality or not detail.token_count:
                 continue
             details[f'{detail.modality.lower()}_{prefix}_tokens'] = detail.token_count
-            if detail.modality != 'AUDIO':
-                continue
-            if metadata_details is metadata.prompt_tokens_details:
-                input_audio_tokens = detail.token_count
-            elif metadata_details is metadata.candidates_tokens_details:
-                output_audio_tokens = detail.token_count
-            elif metadata_details is metadata.cache_tokens_details:  # pragma: no branch
-                cache_audio_read_tokens = detail.token_count
-
-    return usage.RequestUsage(
-        input_tokens=metadata.prompt_token_count or 0,
-        output_tokens=(metadata.candidates_token_count or 0) + thoughts_token_count,
-        cache_read_tokens=cached_content_token_count or 0,
-        input_audio_tokens=input_audio_tokens,
-        output_audio_tokens=output_audio_tokens,
-        cache_audio_read_tokens=cache_audio_read_tokens,
+
+    return usage.RequestUsage.extract(
+        response.model_dump(include={'model_version', 'usage_metadata'}, by_alias=True),
+        provider=provider,
+        provider_url=provider_url,
+        provider_fallback='google',
         details=details,
     )
 

pydantic_ai_slim/pyproject.toml

@@ -60,7 +60,7 @@ dependencies = [
     "exceptiongroup; python_version < '3.11'",
     "opentelemetry-api>=1.28.0",
     "typing-inspection>=0.4.0",
-    "genai-prices>=0.0.35",
+    "genai-prices>=0.0.40",
 ]
 
 [tool.hatch.metadata.hooks.uv-dynamic-versioning.optional-dependencies]

tests/models/test_google.py

@@ -362,7 +362,7 @@ async def test_google_model_builtin_code_execution_stream(
                 ],
                 usage=RequestUsage(
                     input_tokens=46,
-                    output_tokens=528,
+                    output_tokens=1429,
                     details={
                         'thoughts_tokens': 396,
                         'tool_use_prompt_tokens': 901,
@@ -1001,7 +1001,7 @@ async def test_google_model_web_search_tool(allow_model_requests: None, google_p
                 ],
                 usage=RequestUsage(
                     input_tokens=17,
-                    output_tokens=414,
+                    output_tokens=533,
                     details={
                         'thoughts_tokens': 213,
                         'tool_use_prompt_tokens': 119,
@@ -1078,7 +1078,7 @@ async def test_google_model_web_search_tool(allow_model_requests: None, google_p
                 ],
                 usage=RequestUsage(
                     input_tokens=209,
-                    output_tokens=337,
+                    output_tokens=623,
                     details={
                         'thoughts_tokens': 131,
                         'tool_use_prompt_tokens': 286,
@@ -1145,7 +1145,7 @@ async def test_google_model_web_search_tool_stream(allow_model_requests: None, g
                 ],
                 usage=RequestUsage(
                     input_tokens=17,
-                    output_tokens=653,
+                    output_tokens=755,
                     details={
                         'thoughts_tokens': 412,
                         'tool_use_prompt_tokens': 102,
@@ -1322,7 +1322,7 @@ async def test_google_model_web_search_tool_stream(allow_model_requests: None, g
                 ],
                 usage=RequestUsage(
                     input_tokens=249,
-                    output_tokens=541,
+                    output_tokens=860,
                     details={
                         'thoughts_tokens': 301,
                         'tool_use_prompt_tokens': 319,
@@ -1411,7 +1411,7 @@ async def test_google_model_code_execution_tool(allow_model_requests: None, goog
                 ],
                 usage=RequestUsage(
                     input_tokens=15,
-                    output_tokens=660,
+                    output_tokens=1335,
                     details={
                         'thoughts_tokens': 483,
                         'tool_use_prompt_tokens': 675,
@@ -1467,7 +1467,7 @@ async def test_google_model_code_execution_tool(allow_model_requests: None, goog
                 ],
                 usage=RequestUsage(
                     input_tokens=39,
-                    output_tokens=598,
+                    output_tokens=1235,
                     details={
                         'thoughts_tokens': 540,
                         'tool_use_prompt_tokens': 637,
@@ -2719,7 +2719,15 @@ async def get_user_country() -> str:
 
 
 def test_map_usage():
-    assert _metadata_as_usage(GenerateContentResponse()) == RequestUsage()
+    assert (
+        _metadata_as_usage(
+            GenerateContentResponse(),
+            # Test the 'google' provider fallback
+            provider='',
+            provider_url='',
+        )
+        == RequestUsage()
+    )
 
     response = GenerateContentResponse(
         usage_metadata=GenerateContentResponseUsageMetadata(
@@ -2732,7 +2740,7 @@ def test_map_usage():
             candidates_tokens_details=[ModalityTokenCount(modality=MediaModality.AUDIO, token_count=9400)],
         )
     )
-    assert _metadata_as_usage(response) == snapshot(
+    assert _metadata_as_usage(response, provider='', provider_url='') == snapshot(
         RequestUsage(
             input_tokens=1,
             cache_read_tokens=9100,
@@ -3463,6 +3471,7 @@ async def response_iterator() -> AsyncIterator[GenerateContentResponse]:
         _response=response_iterator(),
         _timestamp=datetime.datetime.now(datetime.timezone.utc),
         _provider_name='test-provider',
+        _provider_url='',
     )
 
     events = [event async for event in streamed_response._get_event_iterator()]  # pyright: ignore[reportPrivateUsage]