diff --git a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/main/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParser.java b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/main/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParser.java index 605844f7c2..50ff9f782a 100644 --- a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/main/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParser.java +++ b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/main/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParser.java @@ -94,14 +94,26 @@ public ChatResponse parseResponse(GenerateContentResponse response, Instant star if (response.usageMetadata().isPresent()) { GenerateContentResponseUsageMetadata metadata = response.usageMetadata().get(); - int inputTokens = metadata.promptTokenCount().orElse(0); + int inputTokens = + metadata.promptTokenCount().orElse(0) + + metadata.toolUsePromptTokenCount().orElse(0); int cachedTokens = metadata.cachedContentTokenCount().orElse(0); - int totalOutputTokens = metadata.candidatesTokenCount().orElse(0); int thinkingTokens = metadata.thoughtsTokenCount().orElse(0); - - // Output tokens exclude thinking tokens (following DashScope behavior) - // In Gemini, candidatesTokenCount includes thinking, so we subtract it - int outputTokens = totalOutputTokens - thinkingTokens; + int outputTokens = + metadata.candidatesTokenCount() + .map(candidateTokens -> candidateTokens + thinkingTokens) + .orElseGet( + () -> + metadata.totalTokenCount() + .map( + total -> + Math.max( + 0, + total + - inputTokens)) + // Without candidate or total counts, + // thoughts are the only reported output. + .orElse(thinkingTokens)); usage = ChatUsage.builder() diff --git a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/test/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParserTest.java b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/test/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParserTest.java index 0da20af442..6d1b7c307b 100644 --- a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/test/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParserTest.java +++ b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-gemini/src/test/java/io/agentscope/extensions/model/gemini/formatter/GeminiResponseParserTest.java @@ -216,9 +216,9 @@ void testParseUsageMetadata() { GenerateContentResponseUsageMetadata usageMetadata = GenerateContentResponseUsageMetadata.builder() .promptTokenCount(100) - .candidatesTokenCount(60) // Includes thinking + .candidatesTokenCount(60) .thoughtsTokenCount(10) // Thinking tokens - .totalTokenCount(160) + .totalTokenCount(170) .build(); GenerateContentResponse response = @@ -238,13 +238,71 @@ void testParseUsageMetadata() { // Input tokens = promptTokenCount assertEquals(100, usage.getInputTokens()); - // Output tokens = candidatesTokenCount - thoughtsTokenCount - assertEquals(50, usage.getOutputTokens()); + // Output tokens include candidate and model-generated thinking tokens. + assertEquals(70, usage.getOutputTokens()); // Time should be > 0 assertTrue(usage.getTime() >= 0); } + @Test + void testParseUsageMetadataClassifiesToolUseTokensAsInput() { + GenerateContentResponseUsageMetadata usageMetadata = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(500) + .candidatesTokenCount(120) + .toolUsePromptTokenCount(300) + .thoughtsTokenCount(10) + .totalTokenCount(930) + .build(); + + GenerateContentResponse response = + GenerateContentResponse.builder().usageMetadata(usageMetadata).build(); + + ChatUsage usage = parser.parseResponse(response, startTime).getUsage(); + + assertNotNull(usage); + assertEquals(800, usage.getInputTokens()); + assertEquals(130, usage.getOutputTokens()); + } + + @Test + void testParseUsageMetadataUsesTotalWhenCandidateCountIsMissing() { + GenerateContentResponseUsageMetadata usageMetadata = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(500) + .toolUsePromptTokenCount(300) + .totalTokenCount(930) + .build(); + + GenerateContentResponse response = + GenerateContentResponse.builder().usageMetadata(usageMetadata).build(); + + ChatUsage usage = parser.parseResponse(response, startTime).getUsage(); + + assertNotNull(usage); + assertEquals(800, usage.getInputTokens()); + assertEquals(130, usage.getOutputTokens()); + } + + @Test + void testParseUsageMetadataUsesThinkingWhenCandidateAndTotalCountsAreMissing() { + GenerateContentResponseUsageMetadata usageMetadata = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(500) + .thoughtsTokenCount(10) + .build(); + + GenerateContentResponse response = + GenerateContentResponse.builder().usageMetadata(usageMetadata).build(); + + ChatUsage usage = parser.parseResponse(response, startTime).getUsage(); + + assertNotNull(usage); + assertEquals(500, usage.getInputTokens()); + assertEquals(10, usage.getOutputTokens()); + } + @Test void testParseUsageMetadataReadsCachedContentTokenCount() { // Gemini 报告的 cachedContentTokenCount 必须透传到 ChatUsage.cachedTokens,