diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index aa446c35d0..12208e7b02 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/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, ) diff --git a/pydantic_ai_slim/pyproject.toml b/pydantic_ai_slim/pyproject.toml index 1b5909140d..d408bf88cc 100644 --- a/pydantic_ai_slim/pyproject.toml +++ b/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] diff --git a/tests/models/test_google.py b/tests/models/test_google.py index e729ace8fa..4dee75bd29 100644 --- a/tests/models/test_google.py +++ b/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] diff --git a/uv.lock b/uv.lock index 62d3f22778..deec3732f8 100644 --- a/uv.lock +++ b/uv.lock @@ -1933,16 +1933,16 @@ http = [ [[package]] name = "genai-prices" -version = "0.0.35" +version = "0.0.40" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "eval-type-backport", marker = "python_full_version < '3.11'" }, { name = "httpx" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0d/f4/d1d790c9d1c01e9d1d7ec17985702d2cf209bdf0b7c528a5d99d09054633/genai_prices-0.0.35.tar.gz", hash = "sha256:79b1cdfeef9a0ee7b2895781c3b2beafd22a0aec76a8864aab8aa6ff6013a9e3", size = 45930, upload-time = "2025-10-19T10:05:45.831Z" } +sdist = { url = "https://files.pythonhosted.org/packages/50/86/61ae6a03d5199a5cb3bdc4faa0a9299c1414b32eaa67c2345266efa129a3/genai_prices-0.0.40.tar.gz", hash = "sha256:5b540ede52b484a4ff2876f365c1273286564fc3350a3c3152948e0d6c90d2a8", size = 48024, upload-time = "2025-11-18T18:23:47.521Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/89/84ae62f3ba944ee7bf9009840b7918249cba7f83b9b16f25b85fdfeaeaa8/genai_prices-0.0.35-py3-none-any.whl", hash = "sha256:4e53f19ffe4151074bf7e60f2dd1a0b65593b4c4a9134149bd940e3e77467db2", size = 48565, upload-time = "2025-10-19T10:05:44.464Z" }, + { url = "https://files.pythonhosted.org/packages/af/1f/71e89d07cc56f0c365b12b119a5e0645918e098a6b01bfa616227ca459e9/genai_prices-0.0.40-py3-none-any.whl", hash = "sha256:dd06a067599166c1dc95d46ac9392f2306d7a84e2a6cb1494249d4a02adba54d", size = 50688, upload-time = "2025-11-18T18:23:46.3Z" }, ] [[package]] @@ -5668,7 +5668,7 @@ requires-dist = [ { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, { name = "fasta2a", marker = "extra == 'a2a'", specifier = ">=0.4.1" }, { name = "fastmcp", marker = "extra == 'fastmcp'", specifier = ">=2.12.0" }, - { name = "genai-prices", specifier = ">=0.0.35" }, + { name = "genai-prices", specifier = ">=0.0.40" }, { name = "google-auth", marker = "extra == 'vertexai'", specifier = ">=2.36.0" }, { name = "google-genai", marker = "extra == 'google'", specifier = ">=1.50.1" }, { name = "griffe", specifier = ">=1.3.2" }, @@ -6742,6 +6742,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6b/fa/3234f913fe9a6525a7b97c6dad1f51e72b917e6872e051a5e2ffd8b16fbb/ruamel.yaml.clib-0.2.14-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:70eda7703b8126f5e52fcf276e6c0f40b0d314674f896fc58c47b0aef2b9ae83", size = 137970, upload-time = "2025-09-22T19:51:09.472Z" }, { url = "https://files.pythonhosted.org/packages/ef/ec/4edbf17ac2c87fa0845dd366ef8d5852b96eb58fcd65fc1ecf5fe27b4641/ruamel.yaml.clib-0.2.14-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:a0cb71ccc6ef9ce36eecb6272c81afdc2f565950cdcec33ae8e6cd8f7fc86f27", size = 739639, upload-time = "2025-09-22T19:51:10.566Z" }, { url = "https://files.pythonhosted.org/packages/15/18/b0e1fafe59051de9e79cdd431863b03593ecfa8341c110affad7c8121efc/ruamel.yaml.clib-0.2.14-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e7cb9ad1d525d40f7d87b6df7c0ff916a66bc52cb61b66ac1b2a16d0c1b07640", size = 764456, upload-time = "2025-09-22T19:51:11.736Z" }, + { url = "https://files.pythonhosted.org/packages/e7/cd/150fdb96b8fab27fe08d8a59fe67554568727981806e6bc2677a16081ec7/ruamel_yaml_clib-0.2.14-cp314-cp314-win32.whl", hash = "sha256:9b4104bf43ca0cd4e6f738cb86326a3b2f6eef00f417bd1e7efb7bdffe74c539", size = 102394, upload-time = "2025-11-14T21:57:36.703Z" }, + { url = "https://files.pythonhosted.org/packages/bd/e6/a3fa40084558c7e1dc9546385f22a93949c890a8b2e445b2ba43935f51da/ruamel_yaml_clib-0.2.14-cp314-cp314-win_amd64.whl", hash = "sha256:13997d7d354a9890ea1ec5937a219817464e5cc344805b37671562a401ca3008", size = 122673, upload-time = "2025-11-14T21:57:38.177Z" }, ] [[package]]