From 871d849fc00b02092fbc53ac131aa6d11bd68262 Mon Sep 17 00:00:00 2001 From: "vaidyanath.b" Date: Fri, 19 Dec 2025 13:07:43 +0530 Subject: [PATCH 01/11] add : version file --- VERSION | 1 + 1 file changed, 1 insertion(+) create mode 100644 VERSION diff --git a/VERSION b/VERSION new file mode 100644 index 000000000..87a0f102e --- /dev/null +++ b/VERSION @@ -0,0 +1 @@ +0.5.1-SNAPSHOT \ No newline at end of file From b0d8f0d81dc41785ea2931644568c44b04ed36b4 Mon Sep 17 00:00:00 2001 From: "vaidyanath.b" Date: Tue, 27 Jan 2026 19:36:45 +0530 Subject: [PATCH 02/11] removed logs --- core/src/main/java/com/google/adk/utils/PostgresDBHelper.java | 4 ---- 1 file changed, 4 deletions(-) diff --git a/core/src/main/java/com/google/adk/utils/PostgresDBHelper.java b/core/src/main/java/com/google/adk/utils/PostgresDBHelper.java index b9a2d639c..6936965c3 100644 --- a/core/src/main/java/com/google/adk/utils/PostgresDBHelper.java +++ b/core/src/main/java/com/google/adk/utils/PostgresDBHelper.java @@ -77,10 +77,6 @@ private Connection createConnection() throws SQLException { String password = System.getenv(PASSWORD); String host = System.getenv(DB_URL); - logger.info("DB_URL env var: {}", host); - logger.info("USER env var: {}", userName); - logger.info("PASSWORD env var: {}", password != null ? "***SET***" : "NULL"); - if (host == null) { host = PropertiesHelper.getInstance().getValue("db_url"); } From f0a6d5251a3da96a8b56bfed940613f158f4520e Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Jun 2026 19:33:04 +0000 Subject: [PATCH 03/11] feat: track token usage for live (BIDI) sessions via plugin callbacks Wire onEventCallback and afterRunCallback into Runner.runLiveImpl so the plugin lifecycle runs for live/bidi sessions, matching the non-live path. Previously live events were appended to the session but never passed through any plugin hook, so token usage emitted by Gemini Live (including audio modality breakdowns) could not be captured by plugins. Add LiveTokenTrackingPlugin which observes usageMetadata on each event and logs per-invocation totals on run completion. Live reports cumulative counts, so the plugin keeps the latest value per field rather than summing. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2 --- .../adk/plugins/LiveTokenTrackingPlugin.java | 175 ++++++++++++++++++ .../java/com/google/adk/runner/Runner.java | 19 +- .../plugins/LiveTokenTrackingPluginTest.java | 135 ++++++++++++++ 3 files changed, 328 insertions(+), 1 deletion(-) create mode 100644 core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java create mode 100644 core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java diff --git a/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java b/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java new file mode 100644 index 000000000..6bb224d81 --- /dev/null +++ b/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java @@ -0,0 +1,175 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.plugins; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.ModalityTokenCount; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Plugin that tracks token usage emitted during a run, including bidirectional (BIDI) live sessions + * such as Gemini Live audio. + * + *

Token usage for a live session arrives on dedicated {@code usageMetadata} events (separate + * from the audio/content events) via {@link #onEventCallback}. Gemini Live reports these counts as + * running totals for the session rather than per-event deltas, so this plugin keeps the latest + * value seen for each field rather than summing across events. When the run completes, {@link + * #afterRunCallback} logs the final totals and releases the per-invocation state. + * + *

Register it on the runner like any other plugin to get per-session token accounting for both + * live and non-live runs. + */ +public final class LiveTokenTrackingPlugin extends BasePlugin { + + private static final Logger logger = LoggerFactory.getLogger(LiveTokenTrackingPlugin.class); + + private final Map usageByInvocation = new ConcurrentHashMap<>(); + + public LiveTokenTrackingPlugin() { + super("live_token_tracking_plugin"); + } + + public LiveTokenTrackingPlugin(String name) { + super(name); + } + + @Override + public Maybe onEventCallback(InvocationContext invocationContext, Event event) { + event + .usageMetadata() + .ifPresent( + usageMetadata -> + usageByInvocation + .computeIfAbsent(invocationContext.invocationId(), unused -> new Usage()) + .update(usageMetadata)); + // Return empty so the original event flows through unchanged; this plugin only observes. + return Maybe.empty(); + } + + @Override + public Completable afterRunCallback(InvocationContext invocationContext) { + return Completable.fromRunnable( + () -> { + Usage usage = usageByInvocation.remove(invocationContext.invocationId()); + if (usage == null) { + return; + } + logger.info("Token usage for invocation {}: {}", invocationContext.invocationId(), usage); + }); + } + + /** + * Returns the latest token usage observed for the given invocation, or {@code null} if none has + * been recorded. Intended for tests and programmatic inspection before the run completes. + */ + public Usage usageFor(String invocationId) { + return usageByInvocation.get(invocationId); + } + + /** + * Mutable accumulator holding the latest token counts seen for a single invocation. Live sessions + * report cumulative totals, so each field keeps the most recent non-empty value. + */ + public static final class Usage { + private Integer promptTokenCount; + private Integer candidatesTokenCount; + private Integer totalTokenCount; + private Integer thoughtsTokenCount; + private Integer cachedContentTokenCount; + private final Map promptTokensByModality = new LinkedHashMap<>(); + private final Map candidatesTokensByModality = new LinkedHashMap<>(); + + synchronized void update(GenerateContentResponseUsageMetadata usageMetadata) { + usageMetadata.promptTokenCount().ifPresent(value -> promptTokenCount = value); + usageMetadata.candidatesTokenCount().ifPresent(value -> candidatesTokenCount = value); + usageMetadata.totalTokenCount().ifPresent(value -> totalTokenCount = value); + usageMetadata.thoughtsTokenCount().ifPresent(value -> thoughtsTokenCount = value); + usageMetadata.cachedContentTokenCount().ifPresent(value -> cachedContentTokenCount = value); + usageMetadata + .promptTokensDetails() + .ifPresent(details -> mergeModalityTokens(promptTokensByModality, details)); + usageMetadata + .candidatesTokensDetails() + .ifPresent(details -> mergeModalityTokens(candidatesTokensByModality, details)); + } + + private static void mergeModalityTokens( + Map target, Iterable details) { + for (ModalityTokenCount detail : details) { + if (detail.modality().isEmpty() || detail.tokenCount().isEmpty()) { + continue; + } + target.put(detail.modality().get().toString(), detail.tokenCount().get()); + } + } + + public synchronized Integer promptTokenCount() { + return promptTokenCount; + } + + public synchronized Integer candidatesTokenCount() { + return candidatesTokenCount; + } + + public synchronized Integer totalTokenCount() { + return totalTokenCount; + } + + public synchronized Integer thoughtsTokenCount() { + return thoughtsTokenCount; + } + + public synchronized Integer cachedContentTokenCount() { + return cachedContentTokenCount; + } + + public synchronized Map promptTokensByModality() { + return new LinkedHashMap<>(promptTokensByModality); + } + + public synchronized Map candidatesTokensByModality() { + return new LinkedHashMap<>(candidatesTokensByModality); + } + + @Override + public synchronized String toString() { + return "Usage{" + + "promptTokenCount=" + + promptTokenCount + + ", candidatesTokenCount=" + + candidatesTokenCount + + ", totalTokenCount=" + + totalTokenCount + + ", thoughtsTokenCount=" + + thoughtsTokenCount + + ", cachedContentTokenCount=" + + cachedContentTokenCount + + ", promptTokensByModality=" + + promptTokensByModality + + ", candidatesTokensByModality=" + + candidatesTokensByModality + + '}'; + } + } +} diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 043f56fa3..8746ef0d8 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -749,7 +749,24 @@ protected Flowable runLiveImpl( updatedInvocationContext .agent() .runLive(updatedInvocationContext) - .doOnNext(event -> this.sessionService.appendEvent(session, event))) + .doOnNext(event -> this.sessionService.appendEvent(session, event)) + // Run onEventCallback for each live event so plugins can observe or + // replace it (e.g. token-usage tracking). Mirrors the non-live runImpl + // path; the persisted event above is unaffected by the callback result. + .concatMapSingle( + event -> + updatedInvocationContext + .pluginManager() + .onEventCallback(updatedInvocationContext, event) + .defaultIfEmpty(event)) + // Run afterRunCallback once the live run completes so plugins can flush + // or log aggregates (e.g. total token usage for the session). + .concatWith( + Completable.defer( + () -> + updatedInvocationContext + .pluginManager() + .afterRunCallback(updatedInvocationContext)))) .doOnError( throwable -> { Span span = Span.current(); diff --git a/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java b/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java new file mode 100644 index 000000000..ed3dd3810 --- /dev/null +++ b/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java @@ -0,0 +1,135 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.adk.plugins; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.Mockito.when; + +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.MediaModality; +import com.google.genai.types.ModalityTokenCount; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mock; +import org.mockito.junit.MockitoJUnit; +import org.mockito.junit.MockitoRule; + +@RunWith(JUnit4.class) +public class LiveTokenTrackingPluginTest { + + @Rule public MockitoRule mockitoRule = MockitoJUnit.rule(); + + private static final String INVOCATION_ID = "invocation-1"; + + private final LiveTokenTrackingPlugin plugin = new LiveTokenTrackingPlugin(); + @Mock private InvocationContext mockInvocationContext; + + @Before + public void setUp() { + when(mockInvocationContext.invocationId()).thenReturn(INVOCATION_ID); + } + + private static Event eventWithUsage(GenerateContentResponseUsageMetadata usageMetadata) { + return Event.builder() + .id(Event.generateEventId()) + .author("model") + .usageMetadata(usageMetadata) + .build(); + } + + @Test + public void onEventCallback_keepsLatestCumulativeTotals() { + plugin + .onEventCallback( + mockInvocationContext, + eventWithUsage( + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(10) + .candidatesTokenCount(5) + .totalTokenCount(15) + .build())) + .blockingGet(); + plugin + .onEventCallback( + mockInvocationContext, + eventWithUsage( + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(10) + .candidatesTokenCount(20) + .totalTokenCount(30) + .build())) + .blockingGet(); + + LiveTokenTrackingPlugin.Usage usage = plugin.usageFor(INVOCATION_ID); + assertThat(usage).isNotNull(); + // Live reports running totals, so the latest values win rather than summing. + assertThat(usage.promptTokenCount()).isEqualTo(10); + assertThat(usage.candidatesTokenCount()).isEqualTo(20); + assertThat(usage.totalTokenCount()).isEqualTo(30); + } + + @Test + public void onEventCallback_capturesAudioModalityBreakdown() { + plugin + .onEventCallback( + mockInvocationContext, + eventWithUsage( + GenerateContentResponseUsageMetadata.builder() + .promptTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(new MediaModality(MediaModality.Known.AUDIO)) + .tokenCount(42) + .build())) + .build())) + .blockingGet(); + + LiveTokenTrackingPlugin.Usage usage = plugin.usageFor(INVOCATION_ID); + assertThat(usage).isNotNull(); + assertThat(usage.promptTokensByModality()).containsEntry("AUDIO", 42); + } + + @Test + public void onEventCallback_passesEventThroughUnchanged() { + Event event = + eventWithUsage(GenerateContentResponseUsageMetadata.builder().totalTokenCount(7).build()); + + // Empty result means the original event flows through unmodified downstream. + assertThat(plugin.onEventCallback(mockInvocationContext, event).blockingGet()).isNull(); + } + + @Test + public void afterRunCallback_releasesInvocationState() { + plugin + .onEventCallback( + mockInvocationContext, + eventWithUsage( + GenerateContentResponseUsageMetadata.builder().totalTokenCount(99).build())) + .blockingGet(); + assertThat(plugin.usageFor(INVOCATION_ID)).isNotNull(); + + plugin.afterRunCallback(mockInvocationContext).blockingAwait(); + + assertThat(plugin.usageFor(INVOCATION_ID)).isNull(); + } +} From fb8b5fe371ccd695fb4199f8db67b4783298c473 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Jun 2026 19:44:51 +0000 Subject: [PATCH 04/11] docs: add live BIDI token-tracking harness example Runnable example that drives a Gemini Live BIDI session with LiveTokenTrackingPlugin registered, demonstrating that usageMetadata (including audio modality breakdown) reaches onEventCallback and that afterRunCallback receives the aggregated per-session totals. Requires GOOGLE_API_KEY; model selectable via -Dmodel. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2 --- .../adk/tutorials/LiveTokenPluginHarness.java | 110 ++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java new file mode 100644 index 000000000..22d532802 --- /dev/null +++ b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java @@ -0,0 +1,110 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.tutorials; + +import com.google.adk.agents.LiveRequestQueue; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.plugins.LiveTokenTrackingPlugin; +import com.google.adk.runner.Runner; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.Modality; +import com.google.genai.types.Part; + +/** + * Manual harness that drives a live BIDI session and verifies that {@link LiveTokenTrackingPlugin} + * receives token usage through the plugin callbacks. + * + *

Requires GOOGLE_API_KEY in the environment. Optional system properties: {@code -Dmodel=...} to + * pick the live model. + */ +public class LiveTokenPluginHarness { + public static void main(String[] args) { + String model = System.getProperty("model", "gemini-2.0-flash-live-001"); + + LlmAgent agent = + LlmAgent.builder() + .name("audio_agent") + .model(model) + .instruction("You are a helpful assistant. Briefly introduce yourself.") + .build(); + + LiveTokenTrackingPlugin tokenPlugin = new LiveTokenTrackingPlugin(); + + // Printer plugin registered BEFORE tokenPlugin so its afterRunCallback reads the aggregated + // usage before tokenPlugin's own afterRunCallback clears the state. Proves the hook fires with + // accumulated data (the tutorial's slf4j-simple binding suppresses the plugin's info log). + com.google.adk.plugins.BasePlugin printer = + new com.google.adk.plugins.BasePlugin("usage_printer") { + @Override + public io.reactivex.rxjava3.core.Completable afterRunCallback( + com.google.adk.agents.InvocationContext invocationContext) { + System.out.println( + "[afterRunCallback] aggregated usage = " + + tokenPlugin.usageFor(invocationContext.invocationId())); + return io.reactivex.rxjava3.core.Completable.complete(); + } + }; + + Runner runner = + Runner.builder() + .agent(agent) + .appName("token_harness") + .plugins(printer, tokenPlugin) + .build(); + + RunConfig runConfig = + RunConfig.builder() + .autoCreateSession(true) + .streamingMode(RunConfig.StreamingMode.BIDI) + .responseModalities(ImmutableList.of(new Modality(Modality.Known.AUDIO))) + .build(); + + Content userMessage = + Content.builder() + .role("user") + .parts(ImmutableList.of(Part.fromText("Please introduce yourself in one sentence."))) + .build(); + + LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); + liveRequestQueue.content(userMessage); + + System.out.println("Model: " + model); + System.out.println("Starting live BIDI session..."); + + runner + .runLive("user1", "session1", liveRequestQueue, runConfig) + .doOnNext( + event -> { + if (event.usageMetadata().isPresent()) { + System.out.println("Event carried usageMetadata: " + event.usageMetadata().get()); + } + if ("audio_agent".equals(event.author()) && event.turnComplete().orElse(false)) { + liveRequestQueue.close(); + } + }) + .doOnError(Throwable::printStackTrace) + .blockingSubscribe(); + + System.out.println("\n=== Plugin-captured usage (invocation lookups) ==="); + System.out.println( + "NOTE: afterRunCallback clears state on completion; the per-event log above is the live" + + " proof the plugin saw the data."); + System.out.println("Done."); + System.exit(0); + } +} From a5a74b6f444c6759cd447ae0ce2c8a1c4eebe68a Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Jun 2026 19:53:24 +0000 Subject: [PATCH 05/11] test: extend live token harness to probe per-turn vs cumulative usage Log every usageMetadata event with a sequence number and run two turns in one session. Empirically Gemini Live emits exactly one usageMetadata event per turn, and each event reports that turn's own totals (total = prompt + candidates) rather than a session-cumulative running total. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2 --- .../adk/tutorials/LiveTokenPluginHarness.java | 44 +++++++++++++++---- 1 file changed, 35 insertions(+), 9 deletions(-) diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java index 22d532802..0f0071051 100644 --- a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java +++ b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java @@ -74,32 +74,58 @@ public io.reactivex.rxjava3.core.Completable afterRunCallback( .responseModalities(ImmutableList.of(new Modality(Modality.Known.AUDIO))) .build(); + String prompt = + System.getProperty( + "prompt", "Count slowly from one to twenty out loud, saying each number on its own."); Content userMessage = - Content.builder() - .role("user") - .parts(ImmutableList.of(Part.fromText("Please introduce yourself in one sentence."))) - .build(); + Content.builder().role("user").parts(ImmutableList.of(Part.fromText(prompt))).build(); LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); liveRequestQueue.content(userMessage); System.out.println("Model: " + model); - System.out.println("Starting live BIDI session..."); + System.out.println("Prompt (turn 1): " + prompt); + System.out.println("Starting live BIDI session (2 turns)..."); + java.util.concurrent.atomic.AtomicInteger usageSeq = + new java.util.concurrent.atomic.AtomicInteger(); + java.util.concurrent.atomic.AtomicInteger turn = + new java.util.concurrent.atomic.AtomicInteger(1); runner .runLive("user1", "session1", liveRequestQueue, runConfig) .doOnNext( event -> { - if (event.usageMetadata().isPresent()) { - System.out.println("Event carried usageMetadata: " + event.usageMetadata().get()); - } + event + .usageMetadata() + .ifPresent( + u -> + System.out.printf( + "usageMetadata #%d (turn %d) total=%s prompt=%s candidates=%s%n", + usageSeq.incrementAndGet(), + turn.get(), + u.totalTokenCount().orElse(null), + u.promptTokenCount().orElse(null), + u.candidatesTokenCount().orElse(null))); if ("audio_agent".equals(event.author()) && event.turnComplete().orElse(false)) { - liveRequestQueue.close(); + if (turn.get() == 1) { + turn.set(2); + String prompt2 = "Now say the days of the week out loud, one by one."; + System.out.println("Prompt (turn 2): " + prompt2); + liveRequestQueue.content( + Content.builder() + .role("user") + .parts(ImmutableList.of(Part.fromText(prompt2))) + .build()); + } else { + liveRequestQueue.close(); + } } }) .doOnError(Throwable::printStackTrace) .blockingSubscribe(); + System.out.println("Total usageMetadata events seen: " + usageSeq.get()); + System.out.println("\n=== Plugin-captured usage (invocation lookups) ==="); System.out.println( "NOTE: afterRunCallback clears state on completion; the per-event log above is the live" From 6ef37cc3aa9b9d48896815b5322d5520e9e489e5 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 20 Jun 2026 19:58:00 +0000 Subject: [PATCH 06/11] fix: sum per-turn token usage instead of keeping latest value A live two-turn session showed Gemini Live emits one usageMetadata event per turn, each reporting that turn's own usage (total = prompt + candidates) rather than a session-cumulative running total. Keeping the latest value therefore dropped all earlier turns and undercounted multi-turn sessions. Accumulate by summing scalar counts and per-modality token counts across events so the aggregate reflects true session totals. Verified end-to-end against a live session: per-turn totals 892 and 1044 sum to 1936. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01FWJMsUqoQgsxLScGRk7uB2 --- .../adk/plugins/LiveTokenTrackingPlugin.java | 49 +++++++++++++------ .../plugins/LiveTokenTrackingPluginTest.java | 44 +++++++++++------ 2 files changed, 62 insertions(+), 31 deletions(-) diff --git a/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java b/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java index 6bb224d81..8ebab451c 100644 --- a/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java +++ b/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java @@ -32,9 +32,9 @@ * such as Gemini Live audio. * *

Token usage for a live session arrives on dedicated {@code usageMetadata} events (separate - * from the audio/content events) via {@link #onEventCallback}. Gemini Live reports these counts as - * running totals for the session rather than per-event deltas, so this plugin keeps the latest - * value seen for each field rather than summing across events. When the run completes, {@link + * from the audio/content events) via {@link #onEventCallback}. Gemini Live emits one such event per + * turn, each reporting that turn's own usage rather than a session-cumulative running total, so + * this plugin sums the per-turn values to obtain session totals. When the run completes, {@link * #afterRunCallback} logs the final totals and releases the per-invocation state. * *

Register it on the runner like any other plugin to get per-session token accounting for both @@ -80,16 +80,19 @@ public Completable afterRunCallback(InvocationContext invocationContext) { } /** - * Returns the latest token usage observed for the given invocation, or {@code null} if none has - * been recorded. Intended for tests and programmatic inspection before the run completes. + * Returns the accumulated token usage for the given invocation, or {@code null} if none has been + * recorded. Intended for tests and programmatic inspection before the run completes (after which + * {@link #afterRunCallback} releases the state). */ public Usage usageFor(String invocationId) { return usageByInvocation.get(invocationId); } /** - * Mutable accumulator holding the latest token counts seen for a single invocation. Live sessions - * report cumulative totals, so each field keeps the most recent non-empty value. + * Mutable accumulator that sums token counts across all usageMetadata events seen for a single + * invocation. Gemini Live emits one usageMetadata event per turn, each reporting that turn's own + * usage (not a session-cumulative running total), so per-turn values are summed to obtain the + * session totals. */ public static final class Usage { private Integer promptTokenCount; @@ -101,26 +104,40 @@ public static final class Usage { private final Map candidatesTokensByModality = new LinkedHashMap<>(); synchronized void update(GenerateContentResponseUsageMetadata usageMetadata) { - usageMetadata.promptTokenCount().ifPresent(value -> promptTokenCount = value); - usageMetadata.candidatesTokenCount().ifPresent(value -> candidatesTokenCount = value); - usageMetadata.totalTokenCount().ifPresent(value -> totalTokenCount = value); - usageMetadata.thoughtsTokenCount().ifPresent(value -> thoughtsTokenCount = value); - usageMetadata.cachedContentTokenCount().ifPresent(value -> cachedContentTokenCount = value); + usageMetadata + .promptTokenCount() + .ifPresent(value -> promptTokenCount = sum(promptTokenCount, value)); + usageMetadata + .candidatesTokenCount() + .ifPresent(value -> candidatesTokenCount = sum(candidatesTokenCount, value)); + usageMetadata + .totalTokenCount() + .ifPresent(value -> totalTokenCount = sum(totalTokenCount, value)); + usageMetadata + .thoughtsTokenCount() + .ifPresent(value -> thoughtsTokenCount = sum(thoughtsTokenCount, value)); + usageMetadata + .cachedContentTokenCount() + .ifPresent(value -> cachedContentTokenCount = sum(cachedContentTokenCount, value)); usageMetadata .promptTokensDetails() - .ifPresent(details -> mergeModalityTokens(promptTokensByModality, details)); + .ifPresent(details -> addModalityTokens(promptTokensByModality, details)); usageMetadata .candidatesTokensDetails() - .ifPresent(details -> mergeModalityTokens(candidatesTokensByModality, details)); + .ifPresent(details -> addModalityTokens(candidatesTokensByModality, details)); + } + + private static Integer sum(Integer existing, int addend) { + return existing == null ? addend : existing + addend; } - private static void mergeModalityTokens( + private static void addModalityTokens( Map target, Iterable details) { for (ModalityTokenCount detail : details) { if (detail.modality().isEmpty() || detail.tokenCount().isEmpty()) { continue; } - target.put(detail.modality().get().toString(), detail.tokenCount().get()); + target.merge(detail.modality().get().toString(), detail.tokenCount().get(), Integer::sum); } } diff --git a/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java b/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java index ed3dd3810..85191acf3 100644 --- a/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java +++ b/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java @@ -58,15 +58,17 @@ private static Event eventWithUsage(GenerateContentResponseUsageMetadata usageMe } @Test - public void onEventCallback_keepsLatestCumulativeTotals() { + public void onEventCallback_sumsPerTurnTotalsAcrossEvents() { + // Mirrors observed Gemini Live behavior: one usageMetadata event per turn, each reporting that + // turn's own usage. Session totals are the sum of the per-turn values. plugin .onEventCallback( mockInvocationContext, eventWithUsage( GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(10) - .candidatesTokenCount(5) - .totalTokenCount(15) + .promptTokenCount(185) + .candidatesTokenCount(653) + .totalTokenCount(838) .build())) .blockingGet(); plugin @@ -74,39 +76,51 @@ public void onEventCallback_keepsLatestCumulativeTotals() { mockInvocationContext, eventWithUsage( GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(10) - .candidatesTokenCount(20) - .totalTokenCount(30) + .promptTokenCount(942) + .candidatesTokenCount(241) + .totalTokenCount(1183) .build())) .blockingGet(); LiveTokenTrackingPlugin.Usage usage = plugin.usageFor(INVOCATION_ID); assertThat(usage).isNotNull(); - // Live reports running totals, so the latest values win rather than summing. - assertThat(usage.promptTokenCount()).isEqualTo(10); - assertThat(usage.candidatesTokenCount()).isEqualTo(20); - assertThat(usage.totalTokenCount()).isEqualTo(30); + assertThat(usage.promptTokenCount()).isEqualTo(185 + 942); + assertThat(usage.candidatesTokenCount()).isEqualTo(653 + 241); + assertThat(usage.totalTokenCount()).isEqualTo(838 + 1183); } @Test - public void onEventCallback_capturesAudioModalityBreakdown() { + public void onEventCallback_sumsAudioModalityBreakdownAcrossEvents() { plugin .onEventCallback( mockInvocationContext, eventWithUsage( GenerateContentResponseUsageMetadata.builder() - .promptTokensDetails( + .candidatesTokensDetails( ImmutableList.of( ModalityTokenCount.builder() .modality(new MediaModality(MediaModality.Known.AUDIO)) - .tokenCount(42) + .tokenCount(653) + .build())) + .build())) + .blockingGet(); + plugin + .onEventCallback( + mockInvocationContext, + eventWithUsage( + GenerateContentResponseUsageMetadata.builder() + .candidatesTokensDetails( + ImmutableList.of( + ModalityTokenCount.builder() + .modality(new MediaModality(MediaModality.Known.AUDIO)) + .tokenCount(241) .build())) .build())) .blockingGet(); LiveTokenTrackingPlugin.Usage usage = plugin.usageFor(INVOCATION_ID); assertThat(usage).isNotNull(); - assertThat(usage.promptTokensByModality()).containsEntry("AUDIO", 42); + assertThat(usage.candidatesTokensByModality()).containsEntry("AUDIO", 653 + 241); } @Test From 3d49867148ae5c6831703cf8d0b1af32a7898cfb Mon Sep 17 00:00:00 2001 From: "alfred.jimmy" Date: Tue, 23 Jun 2026 18:36:20 +0530 Subject: [PATCH 07/11] token tracking for bidi --- contrib/sarvam-ai/pom.xml | 2 +- .../google/adk/models/GeminiLlmConnection.java | 14 ++++++++++++++ .../adk/models/GeminiLlmConnectionTest.java | 16 ++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/contrib/sarvam-ai/pom.xml b/contrib/sarvam-ai/pom.xml index 4b6a6b069..22a28d0bc 100644 --- a/contrib/sarvam-ai/pom.xml +++ b/contrib/sarvam-ai/pom.xml @@ -20,7 +20,7 @@ com.google.adk google-adk-parent - 1.4.1-SNAPSHOT + 1.5.1-SNAPSHOT ../../pom.xml diff --git a/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java b/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java index 20b78921a..670015e6f 100644 --- a/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java +++ b/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java @@ -164,6 +164,20 @@ static Optional convertToServerResponse(LiveServerMessage message) } else if (message.setupComplete().isPresent()) { logger.debug("Received setup complete."); return Optional.empty(); + } else if (message.sessionResumptionUpdate().isPresent()) { + logger.debug("Received session resumption update: {}", message.sessionResumptionUpdate().get()); + return Optional.empty(); + } else if (message.goAway().isPresent()) { + logger.debug("Received go away: {}", message.goAway().get()); + return Optional.empty(); + } else if (message.voiceActivityDetectionSignal().isPresent()) { + logger.debug( + "Received voice activity detection signal: {}", + message.voiceActivityDetectionSignal().get()); + return Optional.empty(); + } else if (message.voiceActivity().isPresent()) { + logger.debug("Received voice activity: {}", message.voiceActivity().get()); + return Optional.empty(); } else if (message.usageMetadata().isEmpty()) { logger.warn("Received unknown or empty server message: {}", message.toJson()); builder diff --git a/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java b/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java index 6a95e1532..43c5bbf3f 100644 --- a/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java @@ -190,6 +190,22 @@ public void convertToServerResponse_withSetupComplete_returnsEmpty() { assertThat(result.isPresent()).isFalse(); } + @Test + public void convertToServerResponse_withSessionResumptionUpdate_returnsEmpty() { + LiveServerMessage message = + LiveServerMessage.builder() + .sessionResumptionUpdate( + com.google.genai.types.LiveServerSessionResumptionUpdate.builder() + .newHandle("handle-123") + .resumable(true) + .build()) + .build(); + + Optional result = GeminiLlmConnection.convertToServerResponse(message); + + assertThat(result.isPresent()).isFalse(); + } + @Test public void convertToServerResponse_withUnknownMessage_returnsErrorResponse() { LiveServerMessage message = LiveServerMessage.builder().build(); From cfe33b9880f3a60a5f12a9942b08cb01459484a3 Mon Sep 17 00:00:00 2001 From: "alfred.jimmy" Date: Tue, 23 Jun 2026 18:53:58 +0530 Subject: [PATCH 08/11] added MapDB Support and Ollama basic integration basis https://github.com/TrueGeometry/adk-ollama-mapdb.git --- core/pom.xml | 10 + .../com/google/adk/models/OllamaBaseLM.java | 280 +++++++++++ .../com/google/adk/runner/MapDbRunner.java | 24 + .../adk/sessions/MapDbSessionService.java | 475 ++++++++++++++++++ 4 files changed, 789 insertions(+) create mode 100644 core/src/main/java/com/google/adk/models/OllamaBaseLM.java create mode 100644 core/src/main/java/com/google/adk/runner/MapDbRunner.java create mode 100644 core/src/main/java/com/google/adk/sessions/MapDbSessionService.java diff --git a/core/pom.xml b/core/pom.xml index 20d3978c4..0a2049552 100644 --- a/core/pom.xml +++ b/core/pom.xml @@ -217,6 +217,16 @@ arrow-memory-netty 17.0.0 + + org.json + json + 20180813 + + + org.mapdb + mapdb + 3.0.8 + diff --git a/core/src/main/java/com/google/adk/models/OllamaBaseLM.java b/core/src/main/java/com/google/adk/models/OllamaBaseLM.java new file mode 100644 index 000000000..bb21ddeca --- /dev/null +++ b/core/src/main/java/com/google/adk/models/OllamaBaseLM.java @@ -0,0 +1,280 @@ +/* + * Click nbfs://nbhost/SystemFileSystem/Templates/Licenses/license-default.txt to change this license + * Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template + */ +package com.google.adk.models; + +import com.google.common.collect.ImmutableList; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.Part; + +import io.reactivex.rxjava3.core.Flowable; +import java.io.BufferedReader; +import java.io.DataOutputStream; +import java.io.IOException; +import java.io.InputStreamReader; +import java.net.HttpURLConnection; +import java.net.MalformedURLException; +import java.net.URL; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.logging.Level; +import java.util.stream.Collectors; +import org.json.JSONArray; +import org.json.JSONObject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * + * @author ryzen + */ +public class OllamaBaseLM extends BaseLlm { + + public static String OLLAMA_EP = "http://localhost:11434";//"http://192.168.1.8:11434";// "https://eb28-122-176-48-130.ngrok-free.app";// + + private static final Logger logger = LoggerFactory.getLogger(Claude.class); + + public OllamaBaseLM(String model) { + + super(model); + } + + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + + String systemText = ""; + Optional configOpt = llmRequest.config(); + if (configOpt.isPresent()) { + Optional systemInstructionOpt = configOpt.get().systemInstruction(); + if (systemInstructionOpt.isPresent()) { + String extractedSystemText + = systemInstructionOpt.get().parts().orElse(ImmutableList.of()).stream() + .filter(p -> p.text().isPresent()) + .map(p -> p.text().get()) + .collect(Collectors.joining("\n")); + if (!extractedSystemText.isEmpty()) { + systemText = extractedSystemText; + } + } + } + + String toolSupportedModel =this.model();// "devstral";//"llama3.2:3b-instruct-q2_K";//"llama3.2"; // The 1b doesn't support tool + //Introduce agent to create Ontology + //String agentresponse = agentManager.sendMessageOllama(noteMaker.getName(), toolSupportedModel, "Temperature in Bangalore?"); + + //agentresponse = agentManager.sendMessageOllama(noteMaker.getName(), toolSupportedModel, Ontology_Prompt + "\n" + template_JSON); + //Search the Ontology + String userQuestion = "I want to know 8 detils, What are parts of a car ?"; + JSONArray messagesToSend = new JSONArray();//Order is important + + JSONObject llmMessageJson1 = new JSONObject(); + llmMessageJson1.put("role", "system"); + llmMessageJson1.put("content", systemText); + messagesToSend.put(llmMessageJson1);//Agent system prompt is always added + + JSONObject userMessageJson = new JSONObject(); + userMessageJson.put("role", "user"); + userMessageJson.put("content", llmRequest.contents().get(0).text());//Do better eork here + messagesToSend.put(userMessageJson);//Agent system prompt is always added + + JSONObject agentresponse = callLLMChat(userQuestion, toolSupportedModel, messagesToSend, null); + + String llmResponse = agentresponse.getJSONObject("message").getString("content"); + + LlmResponse.Builder responseBuilder = LlmResponse.builder(); + List parts = new ArrayList<>(); + Part part = ollamaContentBlockToPart(agentresponse.getJSONObject("message")); + parts.add(part); + + responseBuilder.content( + Content.builder().role("model").parts(ImmutableList.copyOf(parts)).build()); + + logger.debug("Ollama response: {}", llmResponse); + + return Flowable.just(responseBuilder.build()); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + throw new UnsupportedOperationException("Not supported yet."); // Generated from nbfs://nbhost/SystemFileSystem/Templates/Classes/Code/GeneratedMethodBody + } + + private void updateTypeString(Map valueDict) { + if (valueDict == null) { + return; + } + if (valueDict.containsKey("type")) { + valueDict.put("type", ((String) valueDict.get("type")).toLowerCase()); + } + + if (valueDict.containsKey("items")) { + updateTypeString((Map) valueDict.get("items")); + + if (valueDict.get("items") instanceof Map + && ((Map) valueDict.get("items")).containsKey("properties")) { + Map properties + = (Map) ((Map) valueDict.get("items")).get("properties"); + if (properties != null) { + for (Object value : properties.values()) { + if (value instanceof Map) { + updateTypeString((Map) value); + } + } + } + } + } + } + + private Part ollamaContentBlockToPart(JSONObject blockJson) { + // Check for tool_calls first, as the example with tool_calls had empty content + if (blockJson.has("tool_calls")) { + JSONArray toolCalls = blockJson.optJSONArray("tool_calls"); // Use optJSONArray for null safety + if (toolCalls != null && toolCalls.length() > 0) { + // Based on the provided structure and LangChain4j Part, + // we typically handle one function call per Part. + // We will process the first tool call in the array. + JSONObject toolCall = toolCalls.optJSONObject(0); // Use optJSONObject for null safety + + if (toolCall != null && toolCall.has("function")) { + JSONObject function = toolCall.optJSONObject("function"); // Use optJSONObject for null safety + + if (function != null && function.has("name") && function.has("arguments")) { + String name = function.optString("name", null); // Use optString for null safety + JSONObject argsJson = function.optJSONObject("arguments"); // Use optJSONObject for null safety + + if (name != null && argsJson != null) { + // Convert JSONObject arguments to Map + // Assuming org.json.JSONObject.toMap() is available + Map args = argsJson.toMap(); + + // Build the FunctionCall Part + // The provided JSON does not include an 'id' for the tool call, so omitting it. + FunctionCall functionCall = FunctionCall.builder() + .name(name) + .args(args) + .build(); + + return Part.builder().functionCall(functionCall).build(); + } + } + } + // If tool_calls array is present but malformed or empty, + // it might fall through to check content or throw. + // Based on original code, falling through to unsupported might be appropriate + // if no valid tool call was found despite the key being present. + } + } + + // If no valid tool_calls were processed, check for text content + if (blockJson.has("content")) { + Object content = blockJson.opt("content"); // Use opt for null safety + if (content instanceof String) { + String text = (String) content; + // Return a text Part, even if the string is empty (matches empty content example) + return Part.builder().text(text).build(); + } + // If 'content' key exists but value is not a String, might be unsupported. + } + + // If neither usable tool_calls nor String content was found + // This covers cases like malformed JSON matching the structure, + // or structures not covered (e.g., image parts, other types). + throw new UnsupportedOperationException("Unsupported content block format or missing required fields: " + blockJson.toString()); + } + + /** + * Use prompt parameter to moderate the questions is prompt!=null, using the + * generate "options": { "num_ctx": 4096 } + * + * @param prompt + * @param model + * @param messages + * @param tools + * @return + */ + public static JSONObject callLLMChat(String prompt, String model, JSONArray messages, JSONArray tools) { + JSONObject responseJ = new JSONObject(); + try { + + // API endpoint URL + String apiUrl = OLLAMA_EP + "/api/chat"; + + // Constructing the JSON payload + JSONObject payload = new JSONObject(); + payload.put("model", model); + payload.put("stream", false); + + JSONObject options = new JSONObject(); + options.put("num_ctx", 4096); + +// JSONArray messages = new JSONArray(); +// JSONObject message = new JSONObject(); +// message.put("role", "user"); +// message.put("content", prompt); +// messages.put(message); + payload.put("messages", messages); + if (tools != null) { + payload.put("tools", tools); + } + payload.put("options", options); + + // Convert payload to string + String jsonString = payload.toString(); + //System.out.println(payload.toString(1)); + + // Create URL object + URL url = new URL(apiUrl); + + // Open connection + HttpURLConnection connection = (HttpURLConnection) url.openConnection(); + + // Set request method + connection.setRequestMethod("POST"); + + // Set headers + connection.setRequestProperty("Content-Type", "application/json"); + + // Enable output and set content length + connection.setDoOutput(true); + connection.setFixedLengthStreamingMode(jsonString.getBytes().length); + + // Write JSON data to output stream + try (DataOutputStream outputStream = new DataOutputStream(connection.getOutputStream())) { + outputStream.writeBytes(jsonString); + outputStream.flush(); + } + + // Read response + int responseCode = connection.getResponseCode(); + System.out.println("Response Code: " + responseCode); + + // Read response body + try (BufferedReader reader = new BufferedReader(new InputStreamReader(connection.getInputStream()))) { + StringBuilder response = new StringBuilder(); + String line; + while ((line = reader.readLine()) != null) { + response.append(line); + } + System.out.println("Response Body: " + response.toString()); + + responseJ = new JSONObject(response.toString()); + + } + + // Close connection + connection.disconnect(); + + } catch (MalformedURLException ex) { + java.util.logging.Logger.getLogger(OllamaBaseLM.class.getName()).log(Level.SEVERE, null, ex); + } catch (IOException ex) { + java.util.logging.Logger.getLogger(OllamaBaseLM.class.getName()).log(Level.SEVERE, null, ex); + } + return responseJ; + } + +} diff --git a/core/src/main/java/com/google/adk/runner/MapDbRunner.java b/core/src/main/java/com/google/adk/runner/MapDbRunner.java new file mode 100644 index 000000000..9bbb18c41 --- /dev/null +++ b/core/src/main/java/com/google/adk/runner/MapDbRunner.java @@ -0,0 +1,24 @@ +/* + * Click nbfs://nbhost/SystemFileSystem/Templates/Licenses/license-default.txt to change this license + * Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template + */ +package com.google.adk.runner; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.sessions.MapDbSessionService; +import java.io.IOException; + +/** The class for the in-memory GenAi runner, using in-memory artifact and session services. */ +public class MapDbRunner extends Runner { + + public MapDbRunner(BaseAgent agent) throws IOException { + // TODO: Change the default appName to InMemoryRunner to align with adk python. + // Check the dev UI in case we break something there. + this(agent, /* appName= */ agent.name()); + } + + public MapDbRunner(BaseAgent agent, String appName) throws IOException { + super(agent, appName, new InMemoryArtifactService(), new MapDbSessionService(appName)); + } +} diff --git a/core/src/main/java/com/google/adk/sessions/MapDbSessionService.java b/core/src/main/java/com/google/adk/sessions/MapDbSessionService.java new file mode 100644 index 000000000..c937bb006 --- /dev/null +++ b/core/src/main/java/com/google/adk/sessions/MapDbSessionService.java @@ -0,0 +1,475 @@ +/* + * Click nbfs://nbhost/SystemFileSystem/Templates/Licenses/license-default.txt to change this license + * Click nbfs://nbhost/SystemFileSystem/Templates/Classes/Class.java to edit this template + */ +package com.google.adk.sessions; + +/** + * + * @author manoj.kumar + */ + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.google.adk.events.Event; +import com.google.adk.events.EventActions; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.io.File; +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.logging.Level; +import java.util.stream.Collectors; +import org.jspecify.annotations.Nullable; +import org.mapdb.DB; +import org.mapdb.DBMaker; +import org.mapdb.Serializer; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A MapDB implementation of {@link BaseSessionService} for persistent storage. + * Stores sessions, user state, and app state in a MapDB file. + * + *

Note: Requires Session, Event, and all objects stored in state maps to be + * serializable by MapDB's Serializer.java. + * State merging (app/user state prefixed with {@code _app_} / {@code _user_}) occurs + * during retrieval operations ({@code getSession}, {@code createSession}). + */ +public final class MapDbSessionService implements BaseSessionService, AutoCloseable { + + private static final Logger logger = LoggerFactory.getLogger(MapDbSessionService.class); + private final DB db; + // Key: sessionId -> Value: Session + private final ConcurrentMap sessionsMap; + // Key: appName:userId -> Value: Map (user state) + private final ConcurrentMap> userStateMap; + // Key: appName -> Value: Map (app state) + private final ConcurrentMap> appStateMap; + + private static final String SESSIONS_MAP_NAME = "sessions"; + private static final String USER_STATE_MAP_NAME = "userState"; + private static final String APP_STATE_MAP_NAME = "appState"; + + /** + * Creates a new instance of the MapDB session service. + * + * @param filePath The path to the MapDB database file. + * @throws IOException if the database file cannot be opened or created. + */ + public MapDbSessionService(String filePath) throws IOException { + Objects.requireNonNull(filePath, "filePath cannot be null"); + + // Configure MapDB - use a file, enable transactions, enable MVStore for concurrency/durability + this.db = DBMaker.fileDB(new File(filePath)) + .transactionEnable() // Use transactions for ACID properties + .executorEnable() // Optional: use separate thread pool for background tasks + .closeOnJvmShutdown() // Ensure database is closed on JVM shutdown + .make(); + + // Get or create maps using Serializer.java (requires Serializable objects) + this.sessionsMap = db.hashMap(SESSIONS_MAP_NAME, Serializer.STRING, Serializer.JAVA) + .createOrOpen(); + this.userStateMap = db.hashMap(USER_STATE_MAP_NAME, Serializer.STRING, Serializer.JAVA) + .createOrOpen(); + this.appStateMap = db.hashMap(APP_STATE_MAP_NAME, Serializer.STRING, Serializer.JAVA) + .createOrOpen(); + + logger.info("MapDbSessionService initialized with file: {}", filePath); + } + + @Override + public Single createSession( + String appName, + String userId, + @Nullable ConcurrentMap state, + @Nullable String sessionId) { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); + + String resolvedSessionId = + Optional.ofNullable(sessionId) + .map(String::trim) + .filter(s -> !s.isEmpty()) + .orElseGet(() -> UUID.randomUUID().toString()); + + // Ensure state map and events list are mutable for the new session + ConcurrentMap initialState = + (state == null) ? new ConcurrentHashMap<>() : new ConcurrentHashMap<>(state); + List initialEvents = new ArrayList<>(); + + // Build the Session object (assumes Session.builder creates a mutable state/events) + Session newSession = Session.builder(resolvedSessionId) .appName(appName) .userId(userId) .state(initialState) // Store initial state in session + .events(initialEvents) .lastUpdateTime(Instant.now()) .build(); + + logger.info( newSession.toJson()); + // Store the new session + sessionsMap.put(resolvedSessionId, newSession.toJson() ); + db.commit(); // Commit the change + + // Create a mutable copy for the return value and merge global state + Session returnCopy = copySession(newSession); + // Merge global state into the copy before returning + return Single.just(mergeWithGlobalState(appName, userId, returnCopy)); + } + + @Override + public Maybe getSession( + String appName, String userId, String sessionId, Optional configOpt) { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); + Objects.requireNonNull(sessionId, "sessionId cannot be null"); + Objects.requireNonNull(configOpt, "configOpt cannot be null"); + + ObjectMapper objectMapper = new ObjectMapper(); + + // Retrieve the session by ID + Session storedSession = null; + try { + storedSession = objectMapper.readValue( sessionsMap.get(sessionId), Session.class); + } catch (JsonProcessingException ex) { + java.util.logging.Logger.getLogger(MapDbSessionService.class.getName()).log(Level.SEVERE, null, ex); + } + + // Also check appName and userId match, although sessionId is the primary key + if (storedSession == null || !appName.equals(storedSession.appName()) || !userId.equals(storedSession.userId())) { + return Maybe.empty(); + } + + // Create a mutable copy to apply filters and merge state + Session sessionCopy = copySession(storedSession); + + // Apply filtering based on config directly to the mutable list in the copy + GetSessionConfig config = configOpt.orElse(GetSessionConfig.builder().build()); + List eventsInCopy = sessionCopy.events(); // Assumes events() returns mutable list + + config + .numRecentEvents() + .ifPresent( + num -> { + if (!eventsInCopy.isEmpty() && num < eventsInCopy.size()) { + // Keep the last 'num' events by removing older ones + // Create sublist view (modifications affect original list) + + List eventsToRemove = eventsInCopy.subList(0, eventsInCopy.size() - num); + eventsToRemove.clear(); // Clear the sublist view, modifying eventsInCopy + } + }); + + // Only apply timestamp filter if numRecentEvents was not applied + if (!config.numRecentEvents().isPresent() && config.afterTimestamp().isPresent()) { + Instant threshold = config.afterTimestamp().get(); + + eventsInCopy.removeIf( + event -> getEventTimestampEpochSeconds(event) < threshold.getEpochSecond()); + } + + // Merge global state into the potentially filtered copy and return + return Maybe.just(mergeWithGlobalState(appName, userId, sessionCopy)); + } + + // Helper to get event timestamp as epoch seconds (adapt based on Event.timestamp() actual type) + private long getEventTimestampEpochSeconds(Event event) { + // Assuming Event.timestamp() returns a value compatible with epoch seconds + // If it returns Instant, use event.timestamp().getEpochSecond() + return event.timestamp(); + } + + @Override + public Single listSessions(String appName, String userId) { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); + + // Assume sessionsMap, appName, and userId are already defined + // Assume sessionsMap is available here (Map) + System.out.println("Printing details for all sessions:"); + sessionsMap.forEach((sessionId, sessiont) -> { + ObjectMapper objectMapper = new ObjectMapper(); + Session session = null; + try { + session = objectMapper.readValue(sessiont, Session.class); + } catch (JsonProcessingException ex) { + java.util.logging.Logger.getLogger(MapDbSessionService.class.getName()).log(Level.SEVERE, null, ex); + } + System.out.println("Session ID: " + sessionId + + ", App Name: " + session.appName() + + ", User ID: " + session.userId()); + }); + // Iterate through all sessions and filter by appName and userId + List sessionCopies = sessionsMap.values().stream() + // .filter(session -> appName.equals(session.appName()) && userId.equals(session.userId())) + .map(this::copySessionMetadata) // Create metadata copies + .collect(Collectors.toCollection(ArrayList::new)); + + return Single.just(ListSessionsResponse.builder().sessions(sessionCopies).build()); + } + + @Override + public Completable deleteSession(String appName, String userId, String sessionId) { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); + Objects.requireNonNull(sessionId, "sessionId cannot be null"); + + // Check if the session exists and belongs to the correct app/user before deleting + ObjectMapper objectMapper = new ObjectMapper(); + Session storedSession = null; + try { + storedSession = objectMapper.readValue( sessionsMap.get(sessionId), Session.class); + } catch (JsonProcessingException ex) { + java.util.logging.Logger.getLogger(MapDbSessionService.class.getName()).log(Level.SEVERE, null, ex); + } +; + if (storedSession != null && appName.equals(storedSession.appName()) && userId.equals(storedSession.userId())) { + sessionsMap.remove(sessionId); + // Note: This implementation, like the InMemory one, does NOT delete + // associated user/app state when a session is deleted. + db.commit(); // Commit the change + } else { + logger.warn("Attempted to delete session {} for user {} in app {}, but it was not found or did not match criteria.", sessionId, userId, appName); + } + return Completable.complete(); // Operation completes even if session wasn't found + } + + @Override + public Single listEvents(String appName, String userId, String sessionId) { + Objects.requireNonNull(appName, "appName cannot be null"); + Objects.requireNonNull(userId, "userId cannot be null"); + Objects.requireNonNull(sessionId, "sessionId cannot be null"); + + // Retrieve the session by ID + ObjectMapper objectMapper = new ObjectMapper(); + Session storedSession = null; + try { + storedSession = objectMapper.readValue(sessionsMap.get(sessionId), Session.class); + } catch (JsonProcessingException ex) { + java.util.logging.Logger.getLogger(MapDbSessionService.class.getName()).log(Level.SEVERE, null, ex); + } +; + + // Also check appName and userId match + if (storedSession == null || !appName.equals(storedSession.appName()) || !userId.equals(storedSession.userId())) { + return Single.just(ListEventsResponse.builder().build()); + } + + // Return a copy of the events list (ImmutableList is safe) + ImmutableList eventsCopy = ImmutableList.copyOf(storedSession.events()); // Assumes events() returns a List + return Single.just(ListEventsResponse.builder().events(eventsCopy).build()); + } + + @CanIgnoreReturnValue + @Override + public Single appendEvent(Session session, Event event) { + Objects.requireNonNull(session, "session cannot be null"); + Objects.requireNonNull(event, "event cannot be null"); + Objects.requireNonNull(session.appName(), "session.appName cannot be null"); + Objects.requireNonNull(session.userId(), "session.userId cannot be null"); + Objects.requireNonNull(session.id(), "session.id cannot be null"); + + String appName = session.appName(); + String userId = session.userId(); + String sessionId = session.id(); + + // Retrieve the *actual* stored session from MapDB + // We need to modify the stored session's event list and possibly state + ObjectMapper objectMapper = new ObjectMapper(); + Session storedSession = null; + try { + storedSession = objectMapper.readValue(sessionsMap.get(sessionId), Session.class); + } catch (JsonProcessingException ex) { + java.util.logging.Logger.getLogger(MapDbSessionService.class.getName()).log(Level.SEVERE, null, ex); + } +; + + if (storedSession == null) { + logger.warn( + String.format( + "appendEvent called for session %s which is not found in MapDbSessionService", + sessionId)); + // Should we create it? The InMemory implementation just logs and does nothing. + // Let's follow that behavior for now. + return Single.error(new IllegalArgumentException("Session not found: " + sessionId)); + } + + // --- Update User/App State --- + EventActions actions = event.actions(); + if (actions != null) { + Map stateDelta = actions.stateDelta(); + if (stateDelta != null && !stateDelta.isEmpty()) { + stateDelta.forEach( + (key, value) -> { + if (key.startsWith(State.APP_PREFIX)) { + String appStateKey = key.substring(State.APP_PREFIX.length()); + // Get, modify, and re-put the app state map + Map currentAppState = appStateMap.computeIfAbsent(appName, k -> new ConcurrentHashMap<>()); + currentAppState.put(appStateKey, value); + appStateMap.put(appName, currentAppState); // Re-put to ensure persistence + } else if (key.startsWith(State.USER_PREFIX)) { + String userStateKey = key.substring(State.USER_PREFIX.length()); + // Get, modify, and re-put the user state map + Map currentUserState = userStateMap.computeIfAbsent( + appName + ":" + userId, k -> new ConcurrentHashMap<>()); + currentUserState.put(userStateKey, value); + userStateMap.put(appName + ":" + userId, currentUserState); // Re-put to ensure persistence + } + }); + // Commit state changes + db.commit(); + } + } + + // --- Append Event to Stored Session --- + // Get the mutable events list from the stored session + List storedEvents = storedSession.events(); // Assumes events() returns mutable list + if (storedEvents != null) { + storedEvents.add(event); // Append the event + + // Update the last update time + storedSession.lastUpdateTime(getInstantFromEvent(event)); + + // Put the modified session back into the map + sessionsMap.put(sessionId, storedSession.toJson()); + + // Commit the session changes + db.commit(); + + // The event should also be added to the *passed-in* session object, as per BaseSessionService contract + // (though the stored session is the persistent one) + BaseSessionService.super.appendEvent(session, event); + + return Single.just(event); + } else { + // This case should ideally not happen if Session is constructed correctly + logger.error("Stored session {} events list is null!", sessionId); + return Single.error(new IllegalStateException("Stored session events list is null")); + } + } + + /** Converts an event's timestamp to an Instant. Adapt based on actual Event structure. */ + // TODO: have Event.timestamp() return Instant directly + private Instant getInstantFromEvent(Event event) { + // Assuming Event.timestamp() returns a double representing epoch seconds + double epochSeconds = event.timestamp(); + long seconds = (long) epochSeconds; + long nanos = (long) ((epochSeconds - seconds) * 1_000_000_000L); + return Instant.ofEpochSecond(seconds, nanos); + } + + /** + * Creates a shallow copy of the session, but with deep copies of the mutable state map and events + * list. Assumes Session provides necessary getters and a suitable constructor/setters that result + * in mutable collections. + * + * @param original The session to copy. + * @return A new Session instance with copied data, including mutable collections. + */ + private Session copySession(Session original) { + // Assumes original.state() and original.events() return collections that + // can be copied into new mutable ones (ConcurrentHashMap, ArrayList). + // Assumes Session.builder can accept these mutable copies. + return Session.builder(original.id()) + .appName(original.appName()) + .userId(original.userId()) + // Create mutable copies of the state map and events list + .state(new ConcurrentHashMap<>(original.state())) + .events(new ArrayList<>(original.events())) + .lastUpdateTime(original.lastUpdateTime()) + .build(); + } + + /** + * Creates a copy of the session containing only metadata fields (ID, appName, userId, timestamp). + * State and Events are explicitly *not* copied. + * + * @param original The session whose metadata to copy. + * @return A new Session instance with only metadata fields populated. + */ + private Session copySessionMetadata(Session original) { + return Session.builder(original.id()) + .appName(original.appName()) + .userId(original.userId()) + .lastUpdateTime(original.lastUpdateTime()) + // Explicitly set state and events to empty/null for metadata copy + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) // Or ImmutableList.of() or null if builder handles null + .build(); + } + + private Session copySessionMetadata(String Session_original) { + ObjectMapper objectMapper = new ObjectMapper(); + Session original = null; + try { + original = objectMapper.readValue(Session_original, Session.class); + } catch (JsonProcessingException ex) { + java.util.logging.Logger.getLogger(MapDbSessionService.class.getName()).log(Level.SEVERE, null, ex); + } + + return Session.builder(original.id()) + .appName(original.appName()) + .userId(original.userId()) + .lastUpdateTime(original.lastUpdateTime()) + // Explicitly set state and events to empty/null for metadata copy + .state(new ConcurrentHashMap<>()) + .events(new ArrayList<>()) // Or ImmutableList.of() or null if builder handles null + .build(); + } + + /** + * Merges the app-specific and user-specific state (stored separately) into the provided + * *mutable* session's state map. + * + * @param appName The application name. + * @param userId The user ID. + * @param session The mutable session whose state map will be augmented. + * @return The same session instance passed in, now with merged state. + */ + @CanIgnoreReturnValue + private Session mergeWithGlobalState(String appName, String userId, Session session) { + Map sessionState = session.state(); // Assumes session.state() returns a mutable map + + // Merge App State + Map currentAppState = appStateMap.get(appName); + if (currentAppState != null) { + currentAppState.forEach((key, value) -> sessionState.put(State.APP_PREFIX + key, value)); + } + + + // Merge User State + Map currentUserState = userStateMap.get(appName + ":" + userId); + if (currentUserState != null) { + currentUserState.forEach((key, value) -> sessionState.put(State.USER_PREFIX + key, value)); + } + + return session; + } + + /** Closes the MapDB database connection. Should be called on application shutdown. */ + @Override + public void close() throws IOException { + if (db != null && !db.isClosed()) { + logger.info("Closing MapDbSessionService database."); + db.close(); + } + } + + // Add a finalize method as a safety net, though try-with-resources and shutdown hook are preferred + @Override + protected void finalize() throws Throwable { + try { + close(); + } finally { + super.finalize(); + } + } +} \ No newline at end of file From bcb703fd24707ed70da0c1964b6a87b20e3ec121 Mon Sep 17 00:00:00 2001 From: "alfred.jimmy" Date: Wed, 24 Jun 2026 18:00:04 +0530 Subject: [PATCH 09/11] removed liveTokenTrackingPlugin since it was folded inside sdk --- .../adk/models/GeminiLlmConnection.java | 3 +- .../adk/plugins/LiveTokenTrackingPlugin.java | 192 ------------------ .../plugins/LiveTokenTrackingPluginTest.java | 149 -------------- .../adk/tutorials/LiveTokenPluginHarness.java | 136 ------------- 4 files changed, 2 insertions(+), 478 deletions(-) delete mode 100644 core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java delete mode 100644 core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java delete mode 100644 tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java diff --git a/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java b/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java index 670015e6f..fbc9bf3e1 100644 --- a/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java +++ b/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java @@ -165,7 +165,8 @@ static Optional convertToServerResponse(LiveServerMessage message) logger.debug("Received setup complete."); return Optional.empty(); } else if (message.sessionResumptionUpdate().isPresent()) { - logger.debug("Received session resumption update: {}", message.sessionResumptionUpdate().get()); + logger.debug( + "Received session resumption update: {}", message.sessionResumptionUpdate().get()); return Optional.empty(); } else if (message.goAway().isPresent()) { logger.debug("Received go away: {}", message.goAway().get()); diff --git a/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java b/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java deleted file mode 100644 index 8ebab451c..000000000 --- a/core/src/main/java/com/google/adk/plugins/LiveTokenTrackingPlugin.java +++ /dev/null @@ -1,192 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.adk.plugins; - -import com.google.adk.agents.InvocationContext; -import com.google.adk.events.Event; -import com.google.genai.types.GenerateContentResponseUsageMetadata; -import com.google.genai.types.ModalityTokenCount; -import io.reactivex.rxjava3.core.Completable; -import io.reactivex.rxjava3.core.Maybe; -import java.util.LinkedHashMap; -import java.util.Map; -import java.util.concurrent.ConcurrentHashMap; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Plugin that tracks token usage emitted during a run, including bidirectional (BIDI) live sessions - * such as Gemini Live audio. - * - *

Token usage for a live session arrives on dedicated {@code usageMetadata} events (separate - * from the audio/content events) via {@link #onEventCallback}. Gemini Live emits one such event per - * turn, each reporting that turn's own usage rather than a session-cumulative running total, so - * this plugin sums the per-turn values to obtain session totals. When the run completes, {@link - * #afterRunCallback} logs the final totals and releases the per-invocation state. - * - *

Register it on the runner like any other plugin to get per-session token accounting for both - * live and non-live runs. - */ -public final class LiveTokenTrackingPlugin extends BasePlugin { - - private static final Logger logger = LoggerFactory.getLogger(LiveTokenTrackingPlugin.class); - - private final Map usageByInvocation = new ConcurrentHashMap<>(); - - public LiveTokenTrackingPlugin() { - super("live_token_tracking_plugin"); - } - - public LiveTokenTrackingPlugin(String name) { - super(name); - } - - @Override - public Maybe onEventCallback(InvocationContext invocationContext, Event event) { - event - .usageMetadata() - .ifPresent( - usageMetadata -> - usageByInvocation - .computeIfAbsent(invocationContext.invocationId(), unused -> new Usage()) - .update(usageMetadata)); - // Return empty so the original event flows through unchanged; this plugin only observes. - return Maybe.empty(); - } - - @Override - public Completable afterRunCallback(InvocationContext invocationContext) { - return Completable.fromRunnable( - () -> { - Usage usage = usageByInvocation.remove(invocationContext.invocationId()); - if (usage == null) { - return; - } - logger.info("Token usage for invocation {}: {}", invocationContext.invocationId(), usage); - }); - } - - /** - * Returns the accumulated token usage for the given invocation, or {@code null} if none has been - * recorded. Intended for tests and programmatic inspection before the run completes (after which - * {@link #afterRunCallback} releases the state). - */ - public Usage usageFor(String invocationId) { - return usageByInvocation.get(invocationId); - } - - /** - * Mutable accumulator that sums token counts across all usageMetadata events seen for a single - * invocation. Gemini Live emits one usageMetadata event per turn, each reporting that turn's own - * usage (not a session-cumulative running total), so per-turn values are summed to obtain the - * session totals. - */ - public static final class Usage { - private Integer promptTokenCount; - private Integer candidatesTokenCount; - private Integer totalTokenCount; - private Integer thoughtsTokenCount; - private Integer cachedContentTokenCount; - private final Map promptTokensByModality = new LinkedHashMap<>(); - private final Map candidatesTokensByModality = new LinkedHashMap<>(); - - synchronized void update(GenerateContentResponseUsageMetadata usageMetadata) { - usageMetadata - .promptTokenCount() - .ifPresent(value -> promptTokenCount = sum(promptTokenCount, value)); - usageMetadata - .candidatesTokenCount() - .ifPresent(value -> candidatesTokenCount = sum(candidatesTokenCount, value)); - usageMetadata - .totalTokenCount() - .ifPresent(value -> totalTokenCount = sum(totalTokenCount, value)); - usageMetadata - .thoughtsTokenCount() - .ifPresent(value -> thoughtsTokenCount = sum(thoughtsTokenCount, value)); - usageMetadata - .cachedContentTokenCount() - .ifPresent(value -> cachedContentTokenCount = sum(cachedContentTokenCount, value)); - usageMetadata - .promptTokensDetails() - .ifPresent(details -> addModalityTokens(promptTokensByModality, details)); - usageMetadata - .candidatesTokensDetails() - .ifPresent(details -> addModalityTokens(candidatesTokensByModality, details)); - } - - private static Integer sum(Integer existing, int addend) { - return existing == null ? addend : existing + addend; - } - - private static void addModalityTokens( - Map target, Iterable details) { - for (ModalityTokenCount detail : details) { - if (detail.modality().isEmpty() || detail.tokenCount().isEmpty()) { - continue; - } - target.merge(detail.modality().get().toString(), detail.tokenCount().get(), Integer::sum); - } - } - - public synchronized Integer promptTokenCount() { - return promptTokenCount; - } - - public synchronized Integer candidatesTokenCount() { - return candidatesTokenCount; - } - - public synchronized Integer totalTokenCount() { - return totalTokenCount; - } - - public synchronized Integer thoughtsTokenCount() { - return thoughtsTokenCount; - } - - public synchronized Integer cachedContentTokenCount() { - return cachedContentTokenCount; - } - - public synchronized Map promptTokensByModality() { - return new LinkedHashMap<>(promptTokensByModality); - } - - public synchronized Map candidatesTokensByModality() { - return new LinkedHashMap<>(candidatesTokensByModality); - } - - @Override - public synchronized String toString() { - return "Usage{" - + "promptTokenCount=" - + promptTokenCount - + ", candidatesTokenCount=" - + candidatesTokenCount - + ", totalTokenCount=" - + totalTokenCount - + ", thoughtsTokenCount=" - + thoughtsTokenCount - + ", cachedContentTokenCount=" - + cachedContentTokenCount - + ", promptTokensByModality=" - + promptTokensByModality - + ", candidatesTokensByModality=" - + candidatesTokensByModality - + '}'; - } - } -} diff --git a/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java b/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java deleted file mode 100644 index 85191acf3..000000000 --- a/core/src/test/java/com/google/adk/plugins/LiveTokenTrackingPluginTest.java +++ /dev/null @@ -1,149 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.adk.plugins; - -import static com.google.common.truth.Truth.assertThat; -import static org.mockito.Mockito.when; - -import com.google.adk.agents.InvocationContext; -import com.google.adk.events.Event; -import com.google.common.collect.ImmutableList; -import com.google.genai.types.GenerateContentResponseUsageMetadata; -import com.google.genai.types.MediaModality; -import com.google.genai.types.ModalityTokenCount; -import org.junit.Before; -import org.junit.Rule; -import org.junit.Test; -import org.junit.runner.RunWith; -import org.junit.runners.JUnit4; -import org.mockito.Mock; -import org.mockito.junit.MockitoJUnit; -import org.mockito.junit.MockitoRule; - -@RunWith(JUnit4.class) -public class LiveTokenTrackingPluginTest { - - @Rule public MockitoRule mockitoRule = MockitoJUnit.rule(); - - private static final String INVOCATION_ID = "invocation-1"; - - private final LiveTokenTrackingPlugin plugin = new LiveTokenTrackingPlugin(); - @Mock private InvocationContext mockInvocationContext; - - @Before - public void setUp() { - when(mockInvocationContext.invocationId()).thenReturn(INVOCATION_ID); - } - - private static Event eventWithUsage(GenerateContentResponseUsageMetadata usageMetadata) { - return Event.builder() - .id(Event.generateEventId()) - .author("model") - .usageMetadata(usageMetadata) - .build(); - } - - @Test - public void onEventCallback_sumsPerTurnTotalsAcrossEvents() { - // Mirrors observed Gemini Live behavior: one usageMetadata event per turn, each reporting that - // turn's own usage. Session totals are the sum of the per-turn values. - plugin - .onEventCallback( - mockInvocationContext, - eventWithUsage( - GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(185) - .candidatesTokenCount(653) - .totalTokenCount(838) - .build())) - .blockingGet(); - plugin - .onEventCallback( - mockInvocationContext, - eventWithUsage( - GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(942) - .candidatesTokenCount(241) - .totalTokenCount(1183) - .build())) - .blockingGet(); - - LiveTokenTrackingPlugin.Usage usage = plugin.usageFor(INVOCATION_ID); - assertThat(usage).isNotNull(); - assertThat(usage.promptTokenCount()).isEqualTo(185 + 942); - assertThat(usage.candidatesTokenCount()).isEqualTo(653 + 241); - assertThat(usage.totalTokenCount()).isEqualTo(838 + 1183); - } - - @Test - public void onEventCallback_sumsAudioModalityBreakdownAcrossEvents() { - plugin - .onEventCallback( - mockInvocationContext, - eventWithUsage( - GenerateContentResponseUsageMetadata.builder() - .candidatesTokensDetails( - ImmutableList.of( - ModalityTokenCount.builder() - .modality(new MediaModality(MediaModality.Known.AUDIO)) - .tokenCount(653) - .build())) - .build())) - .blockingGet(); - plugin - .onEventCallback( - mockInvocationContext, - eventWithUsage( - GenerateContentResponseUsageMetadata.builder() - .candidatesTokensDetails( - ImmutableList.of( - ModalityTokenCount.builder() - .modality(new MediaModality(MediaModality.Known.AUDIO)) - .tokenCount(241) - .build())) - .build())) - .blockingGet(); - - LiveTokenTrackingPlugin.Usage usage = plugin.usageFor(INVOCATION_ID); - assertThat(usage).isNotNull(); - assertThat(usage.candidatesTokensByModality()).containsEntry("AUDIO", 653 + 241); - } - - @Test - public void onEventCallback_passesEventThroughUnchanged() { - Event event = - eventWithUsage(GenerateContentResponseUsageMetadata.builder().totalTokenCount(7).build()); - - // Empty result means the original event flows through unmodified downstream. - assertThat(plugin.onEventCallback(mockInvocationContext, event).blockingGet()).isNull(); - } - - @Test - public void afterRunCallback_releasesInvocationState() { - plugin - .onEventCallback( - mockInvocationContext, - eventWithUsage( - GenerateContentResponseUsageMetadata.builder().totalTokenCount(99).build())) - .blockingGet(); - assertThat(plugin.usageFor(INVOCATION_ID)).isNotNull(); - - plugin.afterRunCallback(mockInvocationContext).blockingAwait(); - - assertThat(plugin.usageFor(INVOCATION_ID)).isNull(); - } -} diff --git a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java b/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java deleted file mode 100644 index 0f0071051..000000000 --- a/tutorials/live-audio-single-agent/src/main/java/com/google/adk/tutorials/LiveTokenPluginHarness.java +++ /dev/null @@ -1,136 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.google.adk.tutorials; - -import com.google.adk.agents.LiveRequestQueue; -import com.google.adk.agents.LlmAgent; -import com.google.adk.agents.RunConfig; -import com.google.adk.plugins.LiveTokenTrackingPlugin; -import com.google.adk.runner.Runner; -import com.google.common.collect.ImmutableList; -import com.google.genai.types.Content; -import com.google.genai.types.Modality; -import com.google.genai.types.Part; - -/** - * Manual harness that drives a live BIDI session and verifies that {@link LiveTokenTrackingPlugin} - * receives token usage through the plugin callbacks. - * - *

Requires GOOGLE_API_KEY in the environment. Optional system properties: {@code -Dmodel=...} to - * pick the live model. - */ -public class LiveTokenPluginHarness { - public static void main(String[] args) { - String model = System.getProperty("model", "gemini-2.0-flash-live-001"); - - LlmAgent agent = - LlmAgent.builder() - .name("audio_agent") - .model(model) - .instruction("You are a helpful assistant. Briefly introduce yourself.") - .build(); - - LiveTokenTrackingPlugin tokenPlugin = new LiveTokenTrackingPlugin(); - - // Printer plugin registered BEFORE tokenPlugin so its afterRunCallback reads the aggregated - // usage before tokenPlugin's own afterRunCallback clears the state. Proves the hook fires with - // accumulated data (the tutorial's slf4j-simple binding suppresses the plugin's info log). - com.google.adk.plugins.BasePlugin printer = - new com.google.adk.plugins.BasePlugin("usage_printer") { - @Override - public io.reactivex.rxjava3.core.Completable afterRunCallback( - com.google.adk.agents.InvocationContext invocationContext) { - System.out.println( - "[afterRunCallback] aggregated usage = " - + tokenPlugin.usageFor(invocationContext.invocationId())); - return io.reactivex.rxjava3.core.Completable.complete(); - } - }; - - Runner runner = - Runner.builder() - .agent(agent) - .appName("token_harness") - .plugins(printer, tokenPlugin) - .build(); - - RunConfig runConfig = - RunConfig.builder() - .autoCreateSession(true) - .streamingMode(RunConfig.StreamingMode.BIDI) - .responseModalities(ImmutableList.of(new Modality(Modality.Known.AUDIO))) - .build(); - - String prompt = - System.getProperty( - "prompt", "Count slowly from one to twenty out loud, saying each number on its own."); - Content userMessage = - Content.builder().role("user").parts(ImmutableList.of(Part.fromText(prompt))).build(); - - LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); - liveRequestQueue.content(userMessage); - - System.out.println("Model: " + model); - System.out.println("Prompt (turn 1): " + prompt); - System.out.println("Starting live BIDI session (2 turns)..."); - - java.util.concurrent.atomic.AtomicInteger usageSeq = - new java.util.concurrent.atomic.AtomicInteger(); - java.util.concurrent.atomic.AtomicInteger turn = - new java.util.concurrent.atomic.AtomicInteger(1); - runner - .runLive("user1", "session1", liveRequestQueue, runConfig) - .doOnNext( - event -> { - event - .usageMetadata() - .ifPresent( - u -> - System.out.printf( - "usageMetadata #%d (turn %d) total=%s prompt=%s candidates=%s%n", - usageSeq.incrementAndGet(), - turn.get(), - u.totalTokenCount().orElse(null), - u.promptTokenCount().orElse(null), - u.candidatesTokenCount().orElse(null))); - if ("audio_agent".equals(event.author()) && event.turnComplete().orElse(false)) { - if (turn.get() == 1) { - turn.set(2); - String prompt2 = "Now say the days of the week out loud, one by one."; - System.out.println("Prompt (turn 2): " + prompt2); - liveRequestQueue.content( - Content.builder() - .role("user") - .parts(ImmutableList.of(Part.fromText(prompt2))) - .build()); - } else { - liveRequestQueue.close(); - } - } - }) - .doOnError(Throwable::printStackTrace) - .blockingSubscribe(); - - System.out.println("Total usageMetadata events seen: " + usageSeq.get()); - - System.out.println("\n=== Plugin-captured usage (invocation lookups) ==="); - System.out.println( - "NOTE: afterRunCallback clears state on completion; the per-event log above is the live" - + " proof the plugin saw the data."); - System.out.println("Done."); - System.exit(0); - } -} From 628dbc23970e7fe3fc0eda54d17aa27745bb3942 Mon Sep 17 00:00:00 2001 From: "alfred.jimmy" Date: Mon, 29 Jun 2026 15:29:24 +0530 Subject: [PATCH 10/11] removed appending events --- .../java/com/google/adk/runner/Runner.java | 23 +++++++++++-------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 8746ef0d8..34a33be95 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -749,24 +749,27 @@ protected Flowable runLiveImpl( updatedInvocationContext .agent() .runLive(updatedInvocationContext) - .doOnNext(event -> this.sessionService.appendEvent(session, event)) + // .doOnNext(event -> this.sessionService.appendEvent(session, event)) // + // Commented out to prevent storing live events in DB // Run onEventCallback for each live event so plugins can observe or // replace it (e.g. token-usage tracking). Mirrors the non-live runImpl // path; the persisted event above is unaffected by the callback result. .concatMapSingle( - event -> - updatedInvocationContext - .pluginManager() - .onEventCallback(updatedInvocationContext, event) - .defaultIfEmpty(event)) + event -> { + return updatedInvocationContext + .pluginManager() + .onEventCallback(updatedInvocationContext, event) + .defaultIfEmpty(event); + }) // Run afterRunCallback once the live run completes so plugins can flush // or log aggregates (e.g. total token usage for the session). .concatWith( Completable.defer( - () -> - updatedInvocationContext - .pluginManager() - .afterRunCallback(updatedInvocationContext)))) + () -> { + return updatedInvocationContext + .pluginManager() + .afterRunCallback(updatedInvocationContext); + }))) .doOnError( throwable -> { Span span = Span.current(); From 651e934a6458feb1a72ff008a354d320e954295c Mon Sep 17 00:00:00 2001 From: "alfred.jimmy" Date: Mon, 29 Jun 2026 15:37:25 +0530 Subject: [PATCH 11/11] readded appending events --- core/src/main/java/com/google/adk/runner/Runner.java | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index 34a33be95..de92209fb 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -749,11 +749,7 @@ protected Flowable runLiveImpl( updatedInvocationContext .agent() .runLive(updatedInvocationContext) - // .doOnNext(event -> this.sessionService.appendEvent(session, event)) // - // Commented out to prevent storing live events in DB - // Run onEventCallback for each live event so plugins can observe or - // replace it (e.g. token-usage tracking). Mirrors the non-live runImpl - // path; the persisted event above is unaffected by the callback result. + .doOnNext(event -> this.sessionService.appendEvent(session, event)) .concatMapSingle( event -> { return updatedInvocationContext