From cc8be0f4010fff246cfb9d73eeb51d2a8aaa84d3 Mon Sep 17 00:00:00 2001 From: liugy789 Date: Tue, 8 Sep 2026 11:50:03 +0800 Subject: [PATCH 1/2] fix(gemini): correct token usage accounting --- .../formatter/GeminiResponseParser.java | 18 ++++--- .../formatter/GeminiResponseParserTest.java | 48 +++++++++++++++++-- 2 files changed, 56 insertions(+), 10 deletions(-) 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..37d2c9f0c5 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,20 @@ 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( + () -> + Math.max( + 0, + metadata.totalTokenCount().orElse(0) + - inputTokens)); 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..e144085ae7 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,53 @@ 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 testParseUsageMetadataReadsCachedContentTokenCount() { // Gemini 报告的 cachedContentTokenCount 必须透传到 ChatUsage.cachedTokens, From 04f3a72c3df40ee71265a6f4379bc0c8f53c7199 Mon Sep 17 00:00:00 2001 From: liugy789 Date: Tue, 8 Sep 2026 13:03:11 +0800 Subject: [PATCH 2/2] fix(gemini): preserve thinking token fallback --- .../gemini/formatter/GeminiResponseParser.java | 14 ++++++++++---- .../formatter/GeminiResponseParserTest.java | 18 ++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) 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 37d2c9f0c5..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 @@ -104,10 +104,16 @@ public ChatResponse parseResponse(GenerateContentResponse response, Instant star .map(candidateTokens -> candidateTokens + thinkingTokens) .orElseGet( () -> - Math.max( - 0, - metadata.totalTokenCount().orElse(0) - - inputTokens)); + 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 e144085ae7..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 @@ -285,6 +285,24 @@ void testParseUsageMetadataUsesTotalWhenCandidateCountIsMissing() { 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,