Skip to content

Commit efbb1e7

Browse files
alexmojakialexmojaki
authored
Include thoughts tokens in output_tokens for Google models (#2634)
1 parent 29867c7 commit efbb1e7

File tree

6 files changed

+21
-21
lines changed

6 files changed

+21
-21
lines changed

‎pydantic_ai_slim/pydantic_ai/models/gemini.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -890,7 +890,7 @@ def _metadata_as_usage(response: _GeminiResponse) -> usage.RequestUsage:
890890

891891
return usage.RequestUsage(
892892
input_tokens=metadata.get('prompt_token_count', 0),
893-
output_tokens=metadata.get('candidates_token_count', 0),
893+
output_tokens=metadata.get('candidates_token_count', 0) + thoughts_token_count,
894894
cache_read_tokens=cached_content_token_count,
895895
input_audio_tokens=input_audio_tokens,
896896
output_audio_tokens=output_audio_tokens,

‎pydantic_ai_slim/pydantic_ai/models/google.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -678,7 +678,7 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
678678
if cached_content_token_count := metadata.cached_content_token_count:
679679
details['cached_content_tokens'] = cached_content_token_count
680680

681-
if thoughts_token_count := metadata.thoughts_token_count:
681+
if thoughts_token_count := (metadata.thoughts_token_count or 0):
682682
details['thoughts_tokens'] = thoughts_token_count
683683

684684
if tool_use_prompt_token_count := metadata.tool_use_prompt_token_count:
@@ -711,7 +711,7 @@ def _metadata_as_usage(response: GenerateContentResponse) -> usage.RequestUsage:
711711

712712
return usage.RequestUsage(
713713
input_tokens=metadata.prompt_token_count or 0,
714-
output_tokens=metadata.candidates_token_count or 0,
714+
output_tokens=(metadata.candidates_token_count or 0) + thoughts_token_count,
715715
cache_read_tokens=cached_content_token_count or 0,
716716
input_audio_tokens=input_audio_tokens,
717717
output_audio_tokens=output_audio_tokens,

‎tests/models/test_gemini.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1230,7 +1230,7 @@ async def get_image() -> BinaryContent:
12301230
ToolCallPart(tool_name='get_image', args={}, tool_call_id=IsStr()),
12311231
],
12321232
usage=RequestUsage(
1233-
input_tokens=38, output_tokens=28, details={'thoughts_tokens': 361, 'text_prompt_tokens': 38}
1233+
input_tokens=38, output_tokens=389, details={'thoughts_tokens': 361, 'text_prompt_tokens': 38}
12341234
),
12351235
model_name='gemini-2.5-pro-preview-03-25',
12361236
timestamp=IsDatetime(),
@@ -1257,7 +1257,7 @@ async def get_image() -> BinaryContent:
12571257
parts=[TextPart(content='The image shows a kiwi fruit, sliced in half.')],
12581258
usage=RequestUsage(
12591259
input_tokens=360,
1260-
output_tokens=11,
1260+
output_tokens=212,
12611261
details={'thoughts_tokens': 201, 'text_prompt_tokens': 102, 'image_prompt_tokens': 258},
12621262
),
12631263
model_name='gemini-2.5-pro-preview-03-25',
@@ -1570,7 +1570,7 @@ async def test_gemini_model_thinking_part(allow_model_requests: None, gemini_api
15701570
),
15711571
],
15721572
usage=RequestUsage(
1573-
input_tokens=801, output_tokens=1519, details={'thoughts_tokens': 794, 'text_prompt_tokens': 801}
1573+
input_tokens=801, output_tokens=2313, details={'thoughts_tokens': 794, 'text_prompt_tokens': 801}
15741574
),
15751575
model_name='gemini-2.5-flash-preview-04-17',
15761576
timestamp=IsDatetime(),
@@ -1844,7 +1844,7 @@ def upcase(text: str) -> str:
18441844
)
18451845
],
18461846
usage=RequestUsage(
1847-
input_tokens=9, output_tokens=44, details={'thoughts_tokens': 545, 'text_prompt_tokens': 9}
1847+
input_tokens=9, output_tokens=589, details={'thoughts_tokens': 545, 'text_prompt_tokens': 9}
18481848
),
18491849
model_name='models/gemini-2.5-pro-preview-05-06',
18501850
timestamp=IsDatetime(),
@@ -2063,7 +2063,7 @@ async def get_user_country() -> str:
20632063
ModelResponse(
20642064
parts=[ToolCallPart(tool_name='get_user_country', args={}, tool_call_id=IsStr())],
20652065
usage=RequestUsage(
2066-
input_tokens=123, output_tokens=12, details={'thoughts_tokens': 318, 'text_prompt_tokens': 123}
2066+
input_tokens=123, output_tokens=330, details={'thoughts_tokens': 318, 'text_prompt_tokens': 123}
20672067
),
20682068
model_name='models/gemini-2.5-pro-preview-05-06',
20692069
timestamp=IsDatetime(),
@@ -2090,7 +2090,7 @@ async def get_user_country() -> str:
20902090
ModelResponse(
20912091
parts=[TextPart(content='{"city": "Mexico City", "country": "Mexico"}')],
20922092
usage=RequestUsage(
2093-
input_tokens=154, output_tokens=13, details={'thoughts_tokens': 94, 'text_prompt_tokens': 154}
2093+
input_tokens=154, output_tokens=107, details={'thoughts_tokens': 94, 'text_prompt_tokens': 154}
20942094
),
20952095
model_name='models/gemini-2.5-pro-preview-05-06',
20962096
timestamp=IsDatetime(),
@@ -2175,7 +2175,7 @@ def test_map_usage():
21752175
RequestUsage(
21762176
input_tokens=1,
21772177
cache_read_tokens=9100,
2178-
output_tokens=2,
2178+
output_tokens=9502,
21792179
input_audio_tokens=9200,
21802180
cache_audio_read_tokens=9300,
21812181
output_audio_tokens=9400,

‎tests/models/test_google.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -257,7 +257,7 @@ async def get_capital(country: str) -> str:
257257
ModelResponse(
258258
parts=[ToolCallPart(tool_name='get_capital', args={'country': 'France'}, tool_call_id=IsStr())],
259259
usage=RequestUsage(
260-
input_tokens=57, output_tokens=15, details={'thoughts_tokens': 155, 'text_prompt_tokens': 57}
260+
input_tokens=57, output_tokens=170, details={'thoughts_tokens': 155, 'text_prompt_tokens': 57}
261261
),
262262
model_name='models/gemini-2.5-pro',
263263
timestamp=IsDatetime(),
@@ -282,7 +282,7 @@ async def get_capital(country: str) -> str:
282282
)
283283
],
284284
usage=RequestUsage(
285-
input_tokens=104, output_tokens=22, details={'thoughts_tokens': 178, 'text_prompt_tokens': 104}
285+
input_tokens=104, output_tokens=200, details={'thoughts_tokens': 178, 'text_prompt_tokens': 104}
286286
),
287287
model_name='models/gemini-2.5-pro',
288288
timestamp=IsDatetime(),
@@ -943,7 +943,7 @@ async def test_google_model_thinking_part(allow_model_requests: None, google_pro
943943
ModelResponse(
944944
parts=[IsInstance(ThinkingPart), IsInstance(TextPart)],
945945
usage=RequestUsage(
946-
input_tokens=15, output_tokens=1041, details={'thoughts_tokens': 1647, 'text_prompt_tokens': 15}
946+
input_tokens=15, output_tokens=2688, details={'thoughts_tokens': 1647, 'text_prompt_tokens': 15}
947947
),
948948
model_name='models/gemini-2.5-pro',
949949
timestamp=IsDatetime(),
@@ -1352,7 +1352,7 @@ async def get_user_country() -> str:
13521352
ModelResponse(
13531353
parts=[ToolCallPart(tool_name='get_user_country', args={}, tool_call_id=IsStr())],
13541354
usage=RequestUsage(
1355-
input_tokens=49, output_tokens=12, details={'thoughts_tokens': 264, 'text_prompt_tokens': 49}
1355+
input_tokens=49, output_tokens=276, details={'thoughts_tokens': 264, 'text_prompt_tokens': 49}
13561356
),
13571357
model_name='models/gemini-2.5-pro',
13581358
timestamp=IsDatetime(),
@@ -1373,7 +1373,7 @@ async def get_user_country() -> str:
13731373
ModelResponse(
13741374
parts=[TextPart(content='The largest city in Mexico is Mexico City.')],
13751375
usage=RequestUsage(
1376-
input_tokens=80, output_tokens=9, details={'thoughts_tokens': 150, 'text_prompt_tokens': 80}
1376+
input_tokens=80, output_tokens=159, details={'thoughts_tokens': 150, 'text_prompt_tokens': 80}
13771377
),
13781378
model_name='models/gemini-2.5-pro',
13791379
timestamp=IsDatetime(),
@@ -1587,7 +1587,7 @@ async def get_user_country() -> str:
15871587
ModelResponse(
15881588
parts=[ToolCallPart(tool_name='get_user_country', args={}, tool_call_id=IsStr())],
15891589
usage=RequestUsage(
1590-
input_tokens=123, output_tokens=12, details={'thoughts_tokens': 132, 'text_prompt_tokens': 123}
1590+
input_tokens=123, output_tokens=144, details={'thoughts_tokens': 132, 'text_prompt_tokens': 123}
15911591
),
15921592
model_name='models/gemini-2.5-pro',
15931593
timestamp=IsDatetime(),
@@ -1615,7 +1615,7 @@ async def get_user_country() -> str:
16151615
ModelResponse(
16161616
parts=[TextPart(content='{"city": "Mexico City", "country": "Mexico"}')],
16171617
usage=RequestUsage(
1618-
input_tokens=154, output_tokens=13, details={'thoughts_tokens': 153, 'text_prompt_tokens': 154}
1618+
input_tokens=154, output_tokens=166, details={'thoughts_tokens': 153, 'text_prompt_tokens': 154}
16191619
),
16201620
model_name='models/gemini-2.5-pro',
16211621
timestamp=IsDatetime(),
@@ -1754,7 +1754,7 @@ def test_map_usage():
17541754
RequestUsage(
17551755
input_tokens=1,
17561756
cache_read_tokens=9100,
1757-
output_tokens=2,
1757+
output_tokens=9502,
17581758
input_audio_tokens=9200,
17591759
cache_audio_read_tokens=9300,
17601760
output_audio_tokens=9400,

‎tests/test_mcp.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -683,7 +683,7 @@ async def test_tool_returning_audio_resource(
683683
ModelResponse(
684684
parts=[ToolCallPart(tool_name='get_audio_resource', args={}, tool_call_id=IsStr())],
685685
usage=RequestUsage(
686-
input_tokens=383, output_tokens=12, details={'thoughts_tokens': 125, 'text_prompt_tokens': 383}
686+
input_tokens=383, output_tokens=137, details={'thoughts_tokens': 125, 'text_prompt_tokens': 383}
687687
),
688688
model_name='models/gemini-2.5-pro-preview-05-06',
689689
timestamp=IsDatetime(),
@@ -745,7 +745,7 @@ async def test_tool_returning_audio_resource_link(
745745
ToolCallPart(tool_name='get_audio_resource_link', args={}, tool_call_id=IsStr()),
746746
],
747747
usage=RequestUsage(
748-
input_tokens=561, output_tokens=41, details={'thoughts_tokens': 195, 'text_prompt_tokens': 561}
748+
input_tokens=561, output_tokens=236, details={'thoughts_tokens': 195, 'text_prompt_tokens': 561}
749749
),
750750
model_name='models/gemini-2.5-pro',
751751
timestamp=IsDatetime(),

‎uv.lock

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)