From b29ecc633a4c4f867529ed08e99f61be9ca58843 Mon Sep 17 00:00:00 2001 From: Sandeep Belgavi Date: Thu, 18 Jun 2026 13:16:37 +0530 Subject: [PATCH 1/5] Reapply "feat: capture live audio tokens in ChatCompletionsResponse" This reverts commit b48b4fc9c28fd41f3e366d81307e6d93174e79da. --- .../models/chat/ChatCompletionsResponse.java | 19 ++++++ ...hatCompletionsResponseAudioTokensTest.java | 68 +++++++++++++++++++ 2 files changed, 87 insertions(+) create mode 100644 core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseAudioTokensTest.java diff --git a/core/src/main/java/com/google/adk/models/chat/ChatCompletionsResponse.java b/core/src/main/java/com/google/adk/models/chat/ChatCompletionsResponse.java index 6cb25f38f..94176d519 100644 --- a/core/src/main/java/com/google/adk/models/chat/ChatCompletionsResponse.java +++ b/core/src/main/java/com/google/adk/models/chat/ChatCompletionsResponse.java @@ -30,6 +30,8 @@ import com.google.genai.types.FinishReason.Known; import com.google.genai.types.FunctionCall; import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.MediaModality; +import com.google.genai.types.ModalityTokenCount; import com.google.genai.types.Part; import java.util.Base64; import java.util.HashMap; @@ -84,6 +86,23 @@ private ChatCompletionsResponse() {} && usage.completionTokensDetails.reasoningTokens != null) { builder.thoughtsTokenCount(usage.completionTokensDetails.reasoningTokens); } + if (usage.promptTokensDetails != null && usage.promptTokensDetails.audioTokens != null) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(usage.promptTokensDetails.audioTokens) + .build())); + } + if (usage.completionTokensDetails != null + && usage.completionTokensDetails.audioTokens != null) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(usage.completionTokensDetails.audioTokens) + .build())); + } return builder.build(); } diff --git a/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseAudioTokensTest.java b/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseAudioTokensTest.java new file mode 100644 index 000000000..5964864b2 --- /dev/null +++ b/core/src/test/java/com/google/adk/models/chat/ChatCompletionsResponseAudioTokensTest.java @@ -0,0 +1,68 @@ +package com.google.adk.models.chat; + +import static com.google.common.truth.Truth.assertThat; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.adk.models.LlmResponse; +import com.google.adk.models.chat.ChatCompletionsResponse.ChatCompletion; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.MediaModality; +import com.google.genai.types.ModalityTokenCount; +import java.util.List; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public class ChatCompletionsResponseAudioTokensTest { + private final ObjectMapper objectMapper = new ObjectMapper(); + + @Test + public void testDeserializeChatCompletion_withAudioTokens() throws Exception { + String json = + """ + { + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o-audio-preview", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!" + }, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + "prompt_tokens_details": { + "audio_tokens": 5 + }, + "completion_tokens_details": { + "audio_tokens": 7 + } + } + } + """; + + ChatCompletion completion = objectMapper.readValue(json, ChatCompletion.class); + LlmResponse response = completion.toLlmResponse(); + + GenerateContentResponseUsageMetadata usage = response.usageMetadata().get(); + + List promptDetails = usage.promptTokensDetails().get(); + assertThat(promptDetails).hasSize(1); + assertThat(promptDetails.get(0).modality().get()) + .isEqualTo(new MediaModality(MediaModality.Known.AUDIO)); + assertThat(promptDetails.get(0).tokenCount().get()).isEqualTo(5); + + List completionDetails = usage.candidatesTokensDetails().get(); + assertThat(completionDetails).hasSize(1); + assertThat(completionDetails.get(0).modality().get()) + .isEqualTo(new MediaModality(MediaModality.Known.AUDIO)); + assertThat(completionDetails.get(0).tokenCount().get()).isEqualTo(7); + } +} From d8fe55511237ffca8d4f55193a6b916a82f0bca4 Mon Sep 17 00:00:00 2001 From: Sandeep Belgavi Date: Thu, 18 Jun 2026 14:07:30 +0530 Subject: [PATCH 2/5] test: add tutorials for live audio tokens Adds ListModels and LiveAudioTokensTest to demonstrate and verify the capture of audio tokens in ChatCompletionsResponse. Co-authored-by: Cursor --- .../com/google/adk/tutorials/ListModels.java | 22 +++++ .../adk/tutorials/LiveAudioTokensTest.java | 83 +++++++++++++++++++ 2 files changed, 105 insertions(+) create mode 100644 tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java create mode 100644 tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java new file mode 100644 index 000000000..314be3ae2 --- /dev/null +++ b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java @@ -0,0 +1,22 @@ +package com.google.adk.tutorials; + +import com.google.genai.Client; +import com.google.genai.types.Model; +import java.util.List; + +public class ListModels { + public static void main(String[] args) throws Exception { + Client client = Client.builder().build(); + List models = + client.models().list(com.google.genai.types.ListModelsConfig.builder().build()).getPage(); + for (Model m : models) { + if (m.name().get().contains("live") + || m.name().get().contains("audio") + || m.name().get().contains("2.0") + || m.name().get().contains("2.5") + || m.name().get().contains("3.1")) { + System.out.println(m.name().get()); + } + } + } +} diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java new file mode 100644 index 000000000..6a8030d96 --- /dev/null +++ b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java @@ -0,0 +1,83 @@ +package com.google.adk.tutorials; + +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.runner.Runner; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.Modality; +import com.google.genai.types.ModalityTokenCount; +import com.google.genai.types.Part; + +public class LiveAudioTokensTest { + public static void main(String[] args) { + LlmAgent agent = + LlmAgent.builder() + .name("audio_agent") + .model("gemini-2.5-pro") // using the requested model + .instruction( + "You are a helpful assistant. Please say 'Hello, how can I help you today?'") + .build(); + + Runner runner = Runner.builder().agent(agent).appName("audio_test").build(); + + RunConfig runConfig = + RunConfig.builder() + .autoCreateSession(true) + .responseModalities(ImmutableList.of(new Modality(Modality.Known.AUDIO))) + .build(); + + Content userMessage = + Content.builder() + .role("user") + .parts(ImmutableList.of(Part.fromText("Please introduce yourself."))) + .build(); + + System.out.println("Sending request to model..."); + + runner + .runAsync("user1", "session1", userMessage, runConfig) + .doOnNext( + event -> { + if (event.author() != null && event.author().equals("model")) { + if (event.content().isPresent()) { + Content c = event.content().get(); + for (Part p : c.parts().get()) { + if (p.text().isPresent()) { + System.out.println("Text: " + p.text().get()); + } + } + } + if (event.usageMetadata().isPresent()) { + GenerateContentResponseUsageMetadata usage = event.usageMetadata().get(); + System.out.println("Total Tokens: " + usage.totalTokenCount().orElse(0)); + + if (usage.promptTokensDetails().isPresent()) { + for (ModalityTokenCount mtc : usage.promptTokensDetails().get()) { + System.out.println( + "Prompt Modality: " + + mtc.modality().get() + + " Tokens: " + + mtc.tokenCount().get()); + } + } + if (usage.candidatesTokensDetails().isPresent()) { + for (ModalityTokenCount mtc : usage.candidatesTokensDetails().get()) { + System.out.println( + "Completion Modality: " + + mtc.modality().get() + + " Tokens: " + + mtc.tokenCount().get()); + } + } + } + } + }) + .doOnError(Throwable::printStackTrace) + .blockingSubscribe(); + + System.out.println("Done."); + System.exit(0); + } +} From 3fd78372adb0b300d2b1ed4250f61228fe6de4c4 Mon Sep 17 00:00:00 2001 From: Sandeep Belgavi Date: Thu, 18 Jun 2026 16:02:52 +0530 Subject: [PATCH 3/5] chore: add author and date, fix ListModels compilation Added Sandeep Belgavi and the current date to the headers of the new tutorial files. Also fixed the compilation issue in ListModels.java. Co-authored-by: Cursor --- .../main/java/com/google/adk/tutorials/ListModels.java | 9 +++++---- .../com/google/adk/tutorials/LiveAudioTokensTest.java | 4 ++++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java index 314be3ae2..705ce320b 100644 --- a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java +++ b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/ListModels.java @@ -1,15 +1,16 @@ +/* + * Author: Sandeep Belgavi + * Date: June 18, 2026 + */ package com.google.adk.tutorials; import com.google.genai.Client; import com.google.genai.types.Model; -import java.util.List; public class ListModels { public static void main(String[] args) throws Exception { Client client = Client.builder().build(); - List models = - client.models().list(com.google.genai.types.ListModelsConfig.builder().build()).getPage(); - for (Model m : models) { + for (Model m : client.models.list(com.google.genai.types.ListModelsConfig.builder().build())) { if (m.name().get().contains("live") || m.name().get().contains("audio") || m.name().get().contains("2.0") diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java index 6a8030d96..4bf39c0de 100644 --- a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java +++ b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveAudioTokensTest.java @@ -1,3 +1,7 @@ +/* + * Author: Sandeep Belgavi + * Date: June 18, 2026 + */ package com.google.adk.tutorials; import com.google.adk.agents.LlmAgent; From 9689fd0f8041c02e754e8084187887af78664f6f Mon Sep 17 00:00:00 2001 From: Sandeep Belgavi Date: Thu, 18 Jun 2026 16:51:55 +0530 Subject: [PATCH 4/5] feat: add audio token usage parsing for Bedrock, Ollama, and Sarvam models Extracts `prompt_tokens_details.audio_tokens` and `completion_tokens_details.audio_tokens` from the API responses of Bedrock, Ollama, and Sarvam models, mapping them into the standard `GenerateContentResponseUsageMetadata` object. This matches the existing support for OpenAI-compatible and Gemini models. Co-authored-by: Cursor --- .../com/google/adk/models/BedrockBaseLM.java | 95 +++++++++++++-- .../com/google/adk/models/OllamaBaseLM.java | 109 ++++++++++++++++-- .../com/google/adk/models/SarvamBaseLM.java | 95 ++++++++++++--- 3 files changed, 261 insertions(+), 38 deletions(-) diff --git a/core/src/main/java/com/google/adk/models/BedrockBaseLM.java b/core/src/main/java/com/google/adk/models/BedrockBaseLM.java index 835c12afd..52bd472f1 100644 --- a/core/src/main/java/com/google/adk/models/BedrockBaseLM.java +++ b/core/src/main/java/com/google/adk/models/BedrockBaseLM.java @@ -19,6 +19,8 @@ import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.MediaModality; +import com.google.genai.types.ModalityTokenCount; import com.google.genai.types.Part; import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; @@ -626,6 +628,8 @@ private Flowable createRobustStreamingResponse( final AtomicInteger inputTokens = new AtomicInteger(0); final AtomicInteger outputTokens = new AtomicInteger(0); final AtomicInteger totalTokens = new AtomicInteger(0); + final AtomicInteger promptAudioTokens = new AtomicInteger(0); + final AtomicInteger completionAudioTokens = new AtomicInteger(0); return Flowable.generate( () -> callLLMChatStream(modelId, messages, functions), @@ -642,7 +646,12 @@ private Flowable createRobustStreamingResponse( if (accumulatedText.length() > 0) { // Create usage metadata from accumulated token counts GenerateContentResponseUsageMetadata usageMetadata = - getUsageMetadata(inputTokens.get(), outputTokens.get(), totalTokens.get()); + getUsageMetadata( + inputTokens.get(), + outputTokens.get(), + totalTokens.get(), + promptAudioTokens.get(), + completionAudioTokens.get()); LlmResponse.Builder finalResponseBuilder = LlmResponse.builder() @@ -693,6 +702,18 @@ private Flowable createRobustStreamingResponse( int total = usage.getInt("totalTokens"); totalTokens.set(total); } + if (usage.has("prompt_tokens_details")) { + JSONObject pDetails = usage.optJSONObject("prompt_tokens_details"); + if (pDetails != null && pDetails.has("audio_tokens")) { + promptAudioTokens.set(pDetails.getInt("audio_tokens")); + } + } + if (usage.has("completion_tokens_details")) { + JSONObject cDetails = usage.optJSONObject("completion_tokens_details"); + if (cDetails != null && cDetails.has("audio_tokens")) { + completionAudioTokens.set(cDetails.getInt("audio_tokens")); + } + } } JSONObject message = null; @@ -792,7 +813,12 @@ private Flowable createRobustStreamingResponse( // Create usage metadata from accumulated token counts GenerateContentResponseUsageMetadata usageMetadata = - getUsageMetadata(inputTokens.get(), outputTokens.get(), totalTokens.get()); + getUsageMetadata( + inputTokens.get(), + outputTokens.get(), + totalTokens.get(), + promptAudioTokens.get(), + completionAudioTokens.get()); // Handle function call completion if (inFunctionCall.get() && functionCallName.length() > 0) { @@ -1284,18 +1310,41 @@ public Flowable generateContent( // Add overloaded method for streaming token usage private GenerateContentResponseUsageMetadata getUsageMetadata( - int promptTokens, int completionTokens, int totalTokens) { + int promptTokens, + int completionTokens, + int totalTokens, + int promptAudioTokens, + int completionAudioTokens) { if (totalTokens > 0 || promptTokens > 0 || completionTokens > 0) { logger.info( "Streaming token counts: prompt={}, completion={}, total={}", promptTokens, completionTokens, totalTokens); - return GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(promptTokens) - .candidatesTokenCount(completionTokens) - .totalTokenCount(totalTokens > 0 ? totalTokens : promptTokens + completionTokens) - .build(); + GenerateContentResponseUsageMetadata.Builder builder = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokens) + .candidatesTokenCount(completionTokens) + .totalTokenCount(totalTokens > 0 ? totalTokens : promptTokens + completionTokens); + + if (promptAudioTokens > 0) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(promptAudioTokens) + .build())); + } + if (completionAudioTokens > 0) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(completionAudioTokens) + .build())); + } + + return builder.build(); } return null; } @@ -1322,12 +1371,36 @@ private GenerateContentResponseUsageMetadata getUsageMetadata(JSONObject agentRe promptTokens, completionTokens, totalTokens); - return Optional.of( + GenerateContentResponseUsageMetadata.Builder builder = GenerateContentResponseUsageMetadata.builder() .promptTokenCount(promptTokens) .candidatesTokenCount(completionTokens) - .totalTokenCount(totalTokens) - .build()); + .totalTokenCount(totalTokens); + + if (usage.has("prompt_tokens_details")) { + JSONObject pDetails = usage.optJSONObject("prompt_tokens_details"); + if (pDetails != null && pDetails.has("audio_tokens")) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(pDetails.getInt("audio_tokens")) + .build())); + } + } + if (usage.has("completion_tokens_details")) { + JSONObject cDetails = usage.optJSONObject("completion_tokens_details"); + if (cDetails != null && cDetails.has("audio_tokens")) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(cDetails.getInt("audio_tokens")) + .build())); + } + } + + return Optional.of(builder.build()); } } } diff --git a/core/src/main/java/com/google/adk/models/OllamaBaseLM.java b/core/src/main/java/com/google/adk/models/OllamaBaseLM.java index 99d0772e0..983c9420b 100644 --- a/core/src/main/java/com/google/adk/models/OllamaBaseLM.java +++ b/core/src/main/java/com/google/adk/models/OllamaBaseLM.java @@ -19,6 +19,8 @@ import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.MediaModality; +import com.google.genai.types.ModalityTokenCount; import com.google.genai.types.Part; import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; @@ -461,6 +463,8 @@ private Flowable createRobustStreamingResponse( final AtomicInteger outputTokens = new AtomicInteger(0); final AtomicLong promptEvalDuration = new AtomicLong(0); final AtomicLong evalDuration = new AtomicLong(0); + final AtomicInteger promptAudioTokens = new AtomicInteger(0); + final AtomicInteger completionAudioTokens = new AtomicInteger(0); return Flowable.generate( () -> callLLMChatStream(modelId, messages, functions), @@ -530,7 +534,33 @@ private Flowable createRobustStreamingResponse( if (responseJson.optBoolean("done", false)) { streamCompleted.set(true); - GenerateContentResponseUsageMetadata usageMetadata = getUsageMetadata(responseJson); + if (responseJson.has("prompt_eval_count")) { + inputTokens.set(responseJson.getInt("prompt_eval_count")); + } + if (responseJson.has("eval_count")) { + outputTokens.set(responseJson.getInt("eval_count")); + } + // Check for audio tokens if Ollama adds them in the future + if (responseJson.has("prompt_tokens_details")) { + JSONObject pDetails = responseJson.optJSONObject("prompt_tokens_details"); + if (pDetails != null && pDetails.has("audio_tokens")) { + promptAudioTokens.set(pDetails.getInt("audio_tokens")); + } + } + if (responseJson.has("completion_tokens_details")) { + JSONObject cDetails = responseJson.optJSONObject("completion_tokens_details"); + if (cDetails != null && cDetails.has("audio_tokens")) { + completionAudioTokens.set(cDetails.getInt("audio_tokens")); + } + } + + GenerateContentResponseUsageMetadata usageMetadata = + getUsageMetadata( + inputTokens.get(), + outputTokens.get(), + inputTokens.get() + outputTokens.get(), + promptAudioTokens.get(), + completionAudioTokens.get()); if (accumulatedText.length() > 0 && !inFunctionCall.get()) { LlmResponse.Builder aggregatedResponseBuilder = @@ -612,13 +642,35 @@ private LlmResponse createTextResponse(String text, boolean partial) { } private GenerateContentResponseUsageMetadata getUsageMetadata( - int promptTokens, int completionTokens, int totalTokens) { + int promptTokens, + int completionTokens, + int totalTokens, + int promptAudioTokens, + int completionAudioTokens) { if (totalTokens > 0 || promptTokens > 0 || completionTokens > 0) { - return GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(promptTokens) - .candidatesTokenCount(completionTokens) - .totalTokenCount(totalTokens > 0 ? totalTokens : promptTokens + completionTokens) - .build(); + GenerateContentResponseUsageMetadata.Builder builder = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokens) + .candidatesTokenCount(completionTokens) + .totalTokenCount(totalTokens > 0 ? totalTokens : promptTokens + completionTokens); + + if (promptAudioTokens > 0) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(promptAudioTokens) + .build())); + } + if (completionAudioTokens > 0) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(completionAudioTokens) + .build())); + } + return builder.build(); } return null; } @@ -632,6 +684,8 @@ private GenerateContentResponseUsageMetadata getUsageMetadata(JSONObject agentRe int promptTokens = 0; int completionTokens = 0; int totalTokens = 0; + int promptAudioTokens = 0; + int completionAudioTokens = 0; if (agentResponse.has("prompt_eval_count")) { promptTokens = agentResponse.getInt("prompt_eval_count"); @@ -642,17 +696,48 @@ private GenerateContentResponseUsageMetadata getUsageMetadata(JSONObject agentRe } totalTokens = promptTokens + completionTokens; + if (agentResponse.has("prompt_tokens_details")) { + JSONObject pDetails = agentResponse.optJSONObject("prompt_tokens_details"); + if (pDetails != null && pDetails.has("audio_tokens")) { + promptAudioTokens = pDetails.getInt("audio_tokens"); + } + } + if (agentResponse.has("completion_tokens_details")) { + JSONObject cDetails = agentResponse.optJSONObject("completion_tokens_details"); + if (cDetails != null && cDetails.has("audio_tokens")) { + completionAudioTokens = cDetails.getInt("audio_tokens"); + } + } + if (totalTokens > 0 || promptTokens > 0 || completionTokens > 0) { logger.info( "Ollama token counts: prompt={}, completion={}, total={}", promptTokens, completionTokens, totalTokens); - return GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(promptTokens) - .candidatesTokenCount(completionTokens) - .totalTokenCount(totalTokens > 0 ? totalTokens : promptTokens + completionTokens) - .build(); + GenerateContentResponseUsageMetadata.Builder builder = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokens) + .candidatesTokenCount(completionTokens) + .totalTokenCount(totalTokens > 0 ? totalTokens : promptTokens + completionTokens); + + if (promptAudioTokens > 0) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(promptAudioTokens) + .build())); + } + if (completionAudioTokens > 0) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(completionAudioTokens) + .build())); + } + return builder.build(); } } catch (Exception e) { logger.warn("Failed to parse token usage from Ollama response", e); diff --git a/core/src/main/java/com/google/adk/models/SarvamBaseLM.java b/core/src/main/java/com/google/adk/models/SarvamBaseLM.java index 487dad652..1bf3e1559 100644 --- a/core/src/main/java/com/google/adk/models/SarvamBaseLM.java +++ b/core/src/main/java/com/google/adk/models/SarvamBaseLM.java @@ -13,6 +13,8 @@ import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.MediaModality; +import com.google.genai.types.ModalityTokenCount; import com.google.genai.types.Part; import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Flowable; @@ -183,6 +185,8 @@ private Flowable generateContentStream(LlmRequest llmRequest) { final AtomicBoolean streamCompleted = new AtomicBoolean(false); final AtomicInteger inputTokens = new AtomicInteger(0); final AtomicInteger outputTokens = new AtomicInteger(0); + final AtomicInteger promptAudioTokens = new AtomicInteger(0); + final AtomicInteger completionAudioTokens = new AtomicInteger(0); return Flowable.generate( () -> @@ -208,7 +212,9 @@ private Flowable generateContentStream(LlmRequest llmRequest) { functionCallName, functionCallArgs, inputTokens.get(), - outputTokens.get()); + outputTokens.get(), + promptAudioTokens.get(), + completionAudioTokens.get()); emitter.onComplete(); return; } @@ -226,7 +232,9 @@ private Flowable generateContentStream(LlmRequest llmRequest) { functionCallName, functionCallArgs, inputTokens.get(), - outputTokens.get()); + outputTokens.get(), + promptAudioTokens.get(), + completionAudioTokens.get()); emitter.onComplete(); return; } @@ -248,6 +256,18 @@ private Flowable generateContentStream(LlmRequest llmRequest) { JSONObject usage = chunk.getJSONObject("usage"); inputTokens.set(usage.optInt("prompt_tokens", 0)); outputTokens.set(usage.optInt("completion_tokens", 0)); + if (usage.has("prompt_tokens_details")) { + JSONObject pDetails = usage.optJSONObject("prompt_tokens_details"); + if (pDetails != null && pDetails.has("audio_tokens")) { + promptAudioTokens.set(pDetails.getInt("audio_tokens")); + } + } + if (usage.has("completion_tokens_details")) { + JSONObject cDetails = usage.optJSONObject("completion_tokens_details"); + if (cDetails != null && cDetails.has("audio_tokens")) { + completionAudioTokens.set(cDetails.getInt("audio_tokens")); + } + } } JSONArray choices = chunk.optJSONArray("choices"); @@ -308,10 +328,13 @@ private void emitFinalStreamResponse( StringBuilder functionCallName, StringBuilder functionCallArgs, int promptTokens, - int completionTokens) { + int completionTokens, + int promptAudioTokens, + int completionAudioTokens) { GenerateContentResponseUsageMetadata usageMetadata = - buildUsageMetadata(promptTokens, completionTokens); + buildUsageMetadata( + promptTokens, completionTokens, promptAudioTokens, completionAudioTokens); if (inFunctionCall.get() && functionCallName.length() > 0) { try { @@ -657,11 +680,35 @@ private GenerateContentResponseUsageMetadata extractUsageMetadata(JSONObject res promptTokens, completionTokens, totalTokens); - return GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(promptTokens) - .candidatesTokenCount(completionTokens) - .totalTokenCount(totalTokens) - .build(); + GenerateContentResponseUsageMetadata.Builder builder = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokens) + .candidatesTokenCount(completionTokens) + .totalTokenCount(totalTokens); + + if (usage.has("prompt_tokens_details")) { + JSONObject pDetails = usage.optJSONObject("prompt_tokens_details"); + if (pDetails != null && pDetails.has("audio_tokens")) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(pDetails.getInt("audio_tokens")) + .build())); + } + } + if (usage.has("completion_tokens_details")) { + JSONObject cDetails = usage.optJSONObject("completion_tokens_details"); + if (cDetails != null && cDetails.has("audio_tokens")) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(cDetails.getInt("audio_tokens")) + .build())); + } + } + return builder.build(); } } catch (Exception e) { logger.warn("Failed to parse token usage from Sarvam response", e); @@ -670,14 +717,32 @@ private GenerateContentResponseUsageMetadata extractUsageMetadata(JSONObject res } private GenerateContentResponseUsageMetadata buildUsageMetadata( - int promptTokens, int completionTokens) { + int promptTokens, int completionTokens, int promptAudioTokens, int completionAudioTokens) { int totalTokens = promptTokens + completionTokens; if (totalTokens > 0 || promptTokens > 0 || completionTokens > 0) { - return GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(promptTokens) - .candidatesTokenCount(completionTokens) - .totalTokenCount(totalTokens) - .build(); + GenerateContentResponseUsageMetadata.Builder builder = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokens) + .candidatesTokenCount(completionTokens) + .totalTokenCount(totalTokens); + + if (promptAudioTokens > 0) { + builder.promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(promptAudioTokens) + .build())); + } + if (completionAudioTokens > 0) { + builder.candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(MediaModality.Known.AUDIO) + .tokenCount(completionAudioTokens) + .build())); + } + return builder.build(); } return null; } From ee8c00f07f8d66c6acb5e8a3fb69d23336f87bbc Mon Sep 17 00:00:00 2001 From: Sandeep Belgavi Date: Thu, 18 Jun 2026 20:23:50 +0530 Subject: [PATCH 5/5] docs: add testing instructions for audio token capture Co-authored-by: Cursor --- tutorials/live-audio-single-agent/README.md | 29 +++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tutorials/live-audio-single-agent/README.md b/tutorials/live-audio-single-agent/README.md index 86c155374..225e6d486 100644 --- a/tutorials/live-audio-single-agent/README.md +++ b/tutorials/live-audio-single-agent/README.md @@ -59,6 +59,35 @@ Once running, you can interact with the agent through: - Process the request and call the `getWeather` tool - Respond with audio (automatically transcribed via outputAudioTranscription) +## Testing Audio Token Capture + +This tutorial also includes samples to test the capture of audio token metrics from LLM responses. The ADK provides a unified `GenerateContentResponseUsageMetadata` object that exposes `promptTokensDetails` and `candidatesTokensDetails` for audio modalities across supported providers (Gemini, OpenAI, Bedrock, Ollama, Sarvam). + +### 1. Run the Live Audio Tokens Test +This test sends an audio instruction to a Gemini model and prints out the token usage breakdown (including audio tokens). + +```shell +# Ensure your API key is set +export GEMINI_API_KEY={YOUR-KEY} + +# Run the test +mvn compile exec:java -pl tutorials/live-audio-single-agent -Dexec.mainClass="com.google.adk.tutorials.LiveAudioTokensTest" +``` + +### 2. Run the List Models Utility +This utility lists all available Gemini models that support audio or live capabilities. + +```shell +mvn compile exec:java -pl tutorials/live-audio-single-agent -Dexec.mainClass="com.google.adk.tutorials.ListModels" +``` + +### 3. Run the Unit Tests +To verify the token parsing logic across different model providers (e.g., OpenAI, Bedrock, Sarvam, Ollama): + +```shell +mvn test -pl core -Dtest=ChatCompletionsResponseAudioTokensTest +``` + ## Learn More See https://google.github.io/adk-docs/get-started/quickstart/#java for more information.