This immutable data class tracks the number of tokens used during a chat completion, - * including input tokens (prompt), output tokens (generated response), cached input tokens, and - * execution time. + * including input tokens (prompt), output tokens (generated response, including reasoning/thinking + * tokens when the provider reports them separately), cached input tokens, and execution time. */ public class ChatUsage { @@ -39,7 +39,8 @@ public class ChatUsage { * #ChatUsage(int, int, int, double)} with {@code cachedTokens} defaulting to {@code 0}. * * @param inputTokens the number of tokens used for the input/prompt - * @param outputTokens the number of tokens used for the output/generated response + * @param outputTokens the number of tokens used for model-generated output, including + * reasoning/thinking tokens when reported separately by the provider * @param time the execution time in seconds */ public ChatUsage(int inputTokens, int outputTokens, double time) { @@ -50,7 +51,8 @@ public ChatUsage(int inputTokens, int outputTokens, double time) { * Creates a new ChatUsage instance. * * @param inputTokens the number of tokens used for the input/prompt - * @param outputTokens the number of tokens used for the output/generated response + * @param outputTokens the number of tokens used for model-generated output, including + * reasoning/thinking tokens when reported separately by the provider * @param cachedTokens the number of input tokens served from the prompt cache (a subset of * {@code inputTokens}); {@code 0} when the provider does not report cache information * @param time the execution time in seconds @@ -79,7 +81,8 @@ public int getInputTokens() { /** * Gets the number of output tokens used. * - * @return the number of tokens used for the output/generated response + * @return the number of model-generated output tokens, including reasoning/thinking tokens when + * reported separately by the provider */ public int getOutputTokens() { return outputTokens; @@ -148,7 +151,8 @@ public Builder inputTokens(int inputTokens) { /** * Sets the number of output tokens. * - * @param outputTokens the number of tokens used for the output/generated response + * @param outputTokens the number of tokens used for model-generated output, including + * reasoning/thinking tokens when reported separately by the provider * @return this builder instance */ public Builder outputTokens(int outputTokens) { diff --git a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-dashscope/src/main/java/io/agentscope/extensions/model/dashscope/dto/DashScopeUsage.java b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-dashscope/src/main/java/io/agentscope/extensions/model/dashscope/dto/DashScopeUsage.java index d80eec1877..f1c3e379e8 100644 --- a/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-dashscope/src/main/java/io/agentscope/extensions/model/dashscope/dto/DashScopeUsage.java +++ b/agentscope-extensions/agentscope-extensions-model/agentscope-extensions-model-dashscope/src/main/java/io/agentscope/extensions/model/dashscope/dto/DashScopeUsage.java @@ -39,7 +39,12 @@ public class DashScopeUsage { @JsonProperty("input_tokens") private Integer inputTokens; - /** Number of tokens in the output. */ + /** + * Number of tokens in the output. + * + *
DashScope includes reasoning tokens in this total; any separately reported reasoning-token + * count is a subset rather than an additional amount. + */ @JsonProperty("output_tokens") private Integer outputTokens; 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..4679b87f9b 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,32 @@ public ChatResponse parseResponse(GenerateContentResponse response, Instant star if (response.usageMetadata().isPresent()) { GenerateContentResponseUsageMetadata metadata = response.usageMetadata().get(); - int inputTokens = metadata.promptTokenCount().orElse(0); + // Server-side tool results are fed back to the model as additional input. + 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; + // Gemini reports candidate and thinking tokens separately; both are output. + // The total already includes thinking, so do not add it again in the fallback. + int outputTokens; + if (metadata.candidatesTokenCount().isPresent()) { + outputTokens = metadata.candidatesTokenCount().get() + thinkingTokens; + } else if (metadata.totalTokenCount().isPresent()) { + int totalTokens = metadata.totalTokenCount().get(); + int reportedOutputTokens = totalTokens - inputTokens; + if (reportedOutputTokens < 0) { + log.debug( + "Gemini usage totalTokenCount ({}) is smaller than input token" + + " count ({}); clamping outputTokens to zero", + totalTokens, + inputTokens); + } + outputTokens = Math.max(0, reportedOutputTokens); + } else { + outputTokens = 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..8bece8b6c3 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 @@ -19,6 +19,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertTrue; import com.google.genai.types.Candidate; @@ -38,6 +39,8 @@ import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; /** * Unit tests for GeminiResponseParser. @@ -216,9 +219,9 @@ void testParseUsageMetadata() { GenerateContentResponseUsageMetadata usageMetadata = GenerateContentResponseUsageMetadata.builder() .promptTokenCount(100) - .candidatesTokenCount(60) // Includes thinking + .candidatesTokenCount(60) // Excludes thinking .thoughtsTokenCount(10) // Thinking tokens - .totalTokenCount(160) + .totalTokenCount(170) .build(); GenerateContentResponse response = @@ -238,8 +241,9 @@ void testParseUsageMetadata() { // Input tokens = promptTokenCount assertEquals(100, usage.getInputTokens()); - // Output tokens = candidatesTokenCount - thoughtsTokenCount - assertEquals(50, usage.getOutputTokens()); + // Output tokens include both candidate and thinking tokens. + assertEquals(70, usage.getOutputTokens()); + assertEquals(170, usage.getTotalTokens()); // Time should be > 0 assertTrue(usage.getTime() >= 0); @@ -260,9 +264,10 @@ void testParseUsageMetadataReadsCachedContentTokenCount() { // cachedContentTokenCount 是 promptTokenCount 的子集(Gemini SDK 文档: // promptTokenCount 包含 cachedContentTokenCount),故 prompt 必须 > cached .promptTokenCount(500) + .toolUsePromptTokenCount(300) .candidatesTokenCount(60) .thoughtsTokenCount(10) - .totalTokenCount(560) + .totalTokenCount(870) .cachedContentTokenCount(300) .build(); @@ -277,6 +282,89 @@ void testParseUsageMetadataReadsCachedContentTokenCount() { assertNotNull(chatResponse.getUsage()); assertEquals(300, chatResponse.getUsage().getCachedTokens()); + assertEquals(800, chatResponse.getUsage().getInputTokens()); + assertEquals(70, chatResponse.getUsage().getOutputTokens()); + assertEquals(870, chatResponse.getUsage().getTotalTokens()); + } + + @Test + void testUsageTokenAccountingFromIssuePayloadJson() { + // Preserve the issue #3033 reproduction in provider wire format so SDK field mapping cannot + // drift together with the parser's arithmetic fixtures. + GenerateContentResponse response = + GenerateContentResponse.fromJson( + """ + { + "usageMetadata": { + "promptTokenCount": 500, + "candidatesTokenCount": 120, + "toolUsePromptTokenCount": 300, + "thoughtsTokenCount": 10, + "totalTokenCount": 930 + } + } + """); + + ChatUsage usage = parser.parseResponse(response, startTime).getUsage(); + + assertNotNull(usage); + assertEquals(800, usage.getInputTokens()); + assertEquals(130, usage.getOutputTokens()); + assertEquals(930, usage.getTotalTokens()); + } + + @ParameterizedTest(name = "{0}") + @CsvSource({ + "server-side tools and thinking, 500, 300, 120, 10, 930, 800, 130", + "server-side tools without thinking, 500, 300, 120, , 920, 800, 120", + "thinking exceeds candidates, 100, , 10, 60, 170, 100, 70", + "missing candidates with tools, 500, 300, , 10, 930, 800, 130", + "missing candidates without tools, 100, , , 10, 170, 100, 70", + "explicit zero candidates, 100, , 0, 10, 170, 100, 10", + "missing total with candidates, 100, , 60, 10, , 100, 70", + "missing total and candidates, 100, , , 10, , 100, 10", + "total smaller than input, 100, 50, , , 120, 150, 0", + "prompt only, 100, , , , , 100, 0", + "no prompt with tool input, , 300, 120, 10, 430, 300, 130", + "empty metadata, , , , , , 0, 0" + }) + void testUsageTokenAccounting( + String scenario, + Integer prompt, + Integer toolPrompt, + Integer candidates, + Integer thoughts, + Integer total, + int expectedInput, + int expectedOutput) { + GenerateContentResponseUsageMetadata.Builder metadata = + GenerateContentResponseUsageMetadata.builder(); + if (prompt != null) { + metadata.promptTokenCount(prompt); + } + if (toolPrompt != null) { + metadata.toolUsePromptTokenCount(toolPrompt); + } + if (candidates != null) { + metadata.candidatesTokenCount(candidates); + } + if (thoughts != null) { + metadata.thoughtsTokenCount(thoughts); + } + if (total != null) { + metadata.totalTokenCount(total); + } + // Streaming responses may carry usage without candidate content. + GenerateContentResponse response = + GenerateContentResponse.builder().usageMetadata(metadata.build()).build(); + + ChatUsage usage = parser.parseResponse(response, startTime).getUsage(); + + assertNotNull(usage); + assertEquals(expectedInput, usage.getInputTokens(), scenario); + assertEquals(expectedOutput, usage.getOutputTokens(), scenario); + assertEquals(expectedInput + expectedOutput, usage.getTotalTokens(), scenario); + assertEquals(0, usage.getCachedTokens()); } @Test @@ -291,6 +379,7 @@ void testParseEmptyResponse() { // Verify assertNotNull(chatResponse); assertEquals("response-empty", chatResponse.getId()); + assertNull(chatResponse.getUsage()); assertEquals(0, chatResponse.getContent().size()); }