From a50ed629a5fd99c6bf2caef3847719db8ef35af9 Mon Sep 17 00:00:00 2001 From: Joel Almeida Date: Wed, 23 Apr 2025 19:00:28 -0700 Subject: [PATCH 1/2] remove extra usage additions during yielding --- pydantic_ai_slim/pydantic_ai/models/gemini.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/gemini.py b/pydantic_ai_slim/pydantic_ai/models/gemini.py index b98bfb9b1..16a24f8fc 100644 --- a/pydantic_ai_slim/pydantic_ai/models/gemini.py +++ b/pydantic_ai_slim/pydantic_ai/models/gemini.py @@ -420,7 +420,6 @@ async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]: responses_to_yield = gemini_responses[:-1] for r in responses_to_yield[current_gemini_response_index:]: current_gemini_response_index += 1 - self._usage += _metadata_as_usage(r) yield r # Now yield the final response, which should be complete From 5c4884a080a5ee0a71790325413b097dfd358b66 Mon Sep 17 00:00:00 2001 From: Joel Almeida Date: Wed, 23 Apr 2025 19:44:31 -0700 Subject: [PATCH 2/2] keep usage along the way --- pydantic_ai_slim/pydantic_ai/models/gemini.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/gemini.py b/pydantic_ai_slim/pydantic_ai/models/gemini.py index 16a24f8fc..e49803d01 100644 --- a/pydantic_ai_slim/pydantic_ai/models/gemini.py +++ b/pydantic_ai_slim/pydantic_ai/models/gemini.py @@ -420,12 +420,13 @@ async def _get_gemini_responses(self) -> AsyncIterator[_GeminiResponse]: responses_to_yield = gemini_responses[:-1] for r in responses_to_yield[current_gemini_response_index:]: current_gemini_response_index += 1 + self._usage = _metadata_as_usage(r) yield r # Now yield the final response, which should be complete if gemini_responses: r = gemini_responses[-1] - self._usage += _metadata_as_usage(r) + self._usage = _metadata_as_usage(r) yield r @property