@@ -371,10 +371,26 @@ def _usage_to_metadata(usage: Any) -> Any:
371371
372372 if not usage :
373373 return None
374+
375+ usage_dict = usage .model_dump ()
376+
377+ prompt_tokens = usage_dict .get ("input_tokens" , 0 )
378+ output_tokens = usage_dict .get ("output_tokens" , 0 )
379+
380+ input_tokens_details = usage_dict .get ("input_tokens_details" , {})
381+ cached_tokens = input_tokens_details .get ("cached_tokens" , 0 )
382+
383+ output_tokens_details = usage_dict .get ("output_tokens_details" , {})
384+ reasoning_tokens = output_tokens_details .get ("reasoning_tokens" , 0 )
385+ completion_tokens = output_tokens - reasoning_tokens
386+ total_tokens = usage_dict .get ("total_tokens" , 0 )
387+
374388 return genai_types .GenerateContentResponseUsageMetadata (
375- prompt_token_count = getattr (usage , "input_tokens" , 0 ),
376- candidates_token_count = getattr (usage , "output_tokens" , 0 ),
377- total_token_count = getattr (usage , "total_tokens" , 0 ),
389+ prompt_token_count = prompt_tokens ,
390+ cached_content_token_count = cached_tokens ,
391+ candidates_token_count = completion_tokens ,
392+ thoughts_token_count = reasoning_tokens ,
393+ total_token_count = total_tokens ,
378394 )
379395
380396
0 commit comments