Skip to content

Extract google model usage using genai-prices #3466

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Nov 18, 2025
+36 −36
Merged

Extract google model usage using genai-prices #3466

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
955792a
Extract google model usage using genai-prices
alexmojaki Nov 18, 2025
f8cf009
Merge branch 'main' into alex/google-genai-prices
DouweM Nov 18, 2025
989dd3e
Merge branch 'main' into alex/google-genai-prices
alexmojaki Nov 18, 2025
41f46f8
Merge branch 'main' into alex/google-genai-prices
DouweM Nov 18, 2025
39dcec5
Merge branch 'main' into alex/google-genai-prices
DouweM Nov 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 11 additions & 22 deletions pydantic_ai_slim/pydantic_ai/models/google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand All @@ -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),
Expand All @@ -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,
)

Expand Down
2 changes: 1 addition & 1 deletion pydantic_ai_slim/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
27 changes: 18 additions & 9 deletions tests/models/test_google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -362,7 +362,7 @@ async def test_google_model_builtin_code_execution_stream(
],
usage=RequestUsage(
input_tokens=46,
output_tokens=528,
output_tokens=1429,
Copy link
Contributor Author

@alexmojaki alexmojaki Nov 18, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All these increases are because toolUsePromptTokenCount is meant to be counted too, see pydantic/genai-prices#200. For example here 1429-528==901 which is the value of tool_use_prompt_tokens just below.

details={
'thoughts_tokens': 396,
'tool_use_prompt_tokens': 901,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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,
Expand Down Expand Up @@ -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]
Expand Down
10 changes: 6 additions & 4 deletions uv.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.