diff --git a/src/agentscope/model/_gemini_model.py b/src/agentscope/model/_gemini_model.py index 6dfeddc8de..18f99fa4e3 100644 --- a/src/agentscope/model/_gemini_model.py +++ b/src/agentscope/model/_gemini_model.py @@ -303,6 +303,36 @@ async def __call__( return parsed_response + def _extract_usage( + self, + usage_metadata: Any, + start_datetime: datetime, + ) -> ChatUsage | None: + """Extract ChatUsage from usage_metadata safely, returning None if + unavailable or if token counts are None. + + Args: + usage_metadata: + The usage metadata object from the Gemini response. + start_datetime (`datetime`): + The start datetime of the generation. + + Returns: + `ChatUsage | None`: + A ChatUsage object, or None if data is unavailable. + """ + if not usage_metadata: + return None + prompt_tokens = usage_metadata.prompt_token_count + total_tokens = usage_metadata.total_token_count + if prompt_tokens is not None and total_tokens is not None: + return ChatUsage( + input_tokens=prompt_tokens, + output_tokens=total_tokens - prompt_tokens, + time=(datetime.now() - start_datetime).total_seconds(), + ) + return None + async def _parse_gemini_stream_generation_response( self, start_datetime: datetime, @@ -383,14 +413,7 @@ async def _parse_gemini_stream_generation_response( if text and structured_model: metadata = _json_loads_with_repair(text) - usage = None - if chunk.usage_metadata: - usage = ChatUsage( - input_tokens=chunk.usage_metadata.prompt_token_count, - output_tokens=chunk.usage_metadata.total_token_count - - chunk.usage_metadata.prompt_token_count, - time=(datetime.now() - start_datetime).total_seconds(), - ) + usage = self._extract_usage(chunk.usage_metadata, start_datetime) # The content blocks for the current chunk content_blocks: list = [] @@ -502,16 +525,7 @@ def _parse_gemini_generation_response( if response.text and structured_model: metadata = _json_loads_with_repair(response.text) - if response.usage_metadata: - usage = ChatUsage( - input_tokens=response.usage_metadata.prompt_token_count, - output_tokens=response.usage_metadata.total_token_count - - response.usage_metadata.prompt_token_count, - time=(datetime.now() - start_datetime).total_seconds(), - ) - - else: - usage = None + usage = self._extract_usage(response.usage_metadata, start_datetime) return ChatResponse( content=content_blocks + tool_calls,