KeyError: 'token_count' when using Gemini with VideoUrl #2489
Description
Initial Checks
- I confirm that I'm using the latest version of Pydantic AI
- I confirm that I searched for my issue in https://github.com/pydantic/pydantic-ai/issues before opening this issue
Description
Hi Pydantic AI team! Thanks for the great work for creating and maintaining this beautiful library.
We've recently encountered an issue that, when trying to use Gemini (GoogleModel) with VideoUrl in the request, sometimes there is a KeyError: 'token_count' error and it breaks the entire iteration. This is not 100% producible, but still a majority of time, we can have this issue.
File /workspace/.venv/lib/python3.11/site-packages/pydantic_ai/result.py:266, in AgentStream.__aiter__.<locals>.aiter()
262 break
264 # If we broke out of the above loop, we need to yield the rest of the events
265 # If we didn't, this will just be a no-op
--> 266 async for event in usage_checking_stream:
267 yield event
File /workspace/.venv/lib/python3.11/site-packages/pydantic_ai/models/google.py:462, in GeminiStreamedResponse._get_event_iterator(self)
460 async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
461 async for chunk in self._response:
--> 462 self._usage = _metadata_as_usage(chunk)
464 assert chunk.candidates is not None
465 candidate = chunk.candidates[0]
File /workspace/.venv/lib/python3.11/site-packages/pydantic_ai/models/google.py:606, in _metadata_as_usage(response)
604 suffix = key.removesuffix('_details')
605 for detail in metadata_details:
--> 606 details[f'{detail["modality"].lower()}_{suffix}'] = detail['token_count']
608 return usage.Usage(
609 request_tokens=metadata.get('prompt_token_count', 0),
610 response_tokens=metadata.get('candidates_token_count', 0),
611 total_tokens=metadata.get('total_token_count', 0),
612 details=details,
613 )
KeyError: 'token_count'
Initially we think it might be due to a corrupt url - however, we tested one single url with multiple requests back to back, a couple of them succeeded and the remaining failed with the error above. Also interestingly, some urls also never failed. So we're wondering is this an issue from the Pydantic AI library or the Google's gen-ai SDK?
Appreciate any insights that you may have. Thank you in advance!
Example Code
from pydantic_ai import Agent, VideoUrl
agent = Agent(
model=GoogleModel('gemini-2.5-flash-lite'),
)
async with agent.iter(
[
VIDEO_DESCRIPTION_PROMPT,
VideoUrl(url=url, media_type=mime_type),
]
) as agent_run:
async for node in agent_run:
logger.info(node)Python, Pydantic AI & LLM client version
Python 3.11
pydantic-ai 0.6.2
google-genai 1.29.0