diff --git a/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java b/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java
index 4159ad2cd..2945c9a85 100644
--- a/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java
+++ b/core/src/main/java/com/google/adk/summarizer/TailRetentionEventCompactor.java
@@ -22,7 +22,6 @@
import com.google.adk.events.EventCompaction;
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
-import com.google.common.collect.Lists;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import io.reactivex.rxjava3.core.Completable;
import io.reactivex.rxjava3.core.Maybe;
@@ -41,7 +40,10 @@
*
Keeps the {@code retentionSize} most recent events raw.
* Compacts all events that never compacted and older than the retained tail, including the
* most recent compaction event, into a new summary event.
- * Triggers compaction only if the prompt token count exceeds the {@code tokenThreshold}.
+ * Triggers compaction only if the estimated prompt token count exceeds the {@code
+ * tokenThreshold}: the last reported prompt token count plus about one token per four
+ * characters of the response that reported it and the events after it. Without a count
+ * reported since the latest compaction, all events since then are estimated this way.
* The new summary event is generated by the {@link BaseEventSummarizer}.
* Appends this new summary event to the end of the event stream.
*
@@ -140,10 +142,11 @@ public Completable compact(Session session, BaseSessionService sessionService) {
* @param events The list of events to process.
*/
private Maybe> getCompactionEvents(List events) {
- Optional count = getLatestPromptTokenCount(events);
+ Optional count = getEstimatedPromptTokenCount(events);
if (count.isPresent() && count.get() <= tokenThreshold) {
logger.debug(
- "Skipping compaction. Prompt token count {} is within threshold {}",
+ "Skipping compaction. Prompt token count {} (last reported plus estimate) is within"
+ + " threshold {}",
count.get(),
tokenThreshold);
return Maybe.empty();
@@ -226,13 +229,28 @@ private int estimateTokenCount(List events) {
return events.stream().mapToInt(event -> event.stringifyContent().length()).sum() / 4;
}
- private Optional getLatestPromptTokenCount(List events) {
- return Lists.reverse(events).stream()
- .map(Event::usageMetadata)
- .flatMap(Optional::stream)
- .map(GenerateContentResponseUsageMetadata::promptTokenCount)
- .flatMap(Optional::stream)
- .findFirst();
+ /**
+ * Returns the latest reported prompt token count plus an estimate for the events it does not
+ * cover, or empty if there is no such count after the latest compaction.
+ */
+ private Optional getEstimatedPromptTokenCount(List events) {
+ ListIterator iter = events.listIterator(events.size());
+ while (iter.hasPrevious()) {
+ Event event = iter.previous();
+ if (isCompactEvent(event)) {
+ // A count reported before a compaction describes a prompt that the compaction replaced.
+ return Optional.empty();
+ }
+ Optional promptTokenCount =
+ event.usageMetadata().flatMap(GenerateContentResponseUsageMetadata::promptTokenCount);
+ if (promptTokenCount.isPresent()) {
+ // The count excludes this response and every event appended after it.
+ return Optional.of(
+ promptTokenCount.get()
+ + estimateTokenCount(events.subList(iter.nextIndex(), events.size())));
+ }
+ }
+ return Optional.empty();
}
private static boolean isCompactEvent(Event event) {
diff --git a/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java b/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java
index 3ceba5641..a5017beb6 100644
--- a/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java
+++ b/core/src/test/java/com/google/adk/flows/llmflows/CompactionTest.java
@@ -31,7 +31,10 @@
import com.google.adk.summarizer.BaseEventSummarizer;
import com.google.adk.summarizer.EventsCompactionConfig;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
+import com.google.genai.types.Part;
import io.reactivex.rxjava3.core.Maybe;
import io.reactivex.rxjava3.core.Single;
import java.util.Optional;
@@ -82,13 +85,9 @@ public void processRequest_withConfig_triggersCompaction() {
when(context.eventsCompactionConfig()).thenReturn(Optional.of(config));
// Setup events with usage > 100 to trigger compaction
- Event event1 = mock(Event.class);
- Event event2 = mock(Event.class);
- Event event3 = mock(Event.class);
- when(event3.usageMetadata())
- .thenReturn(
- Optional.of(
- GenerateContentResponseUsageMetadata.builder().promptTokenCount(200).build()));
+ Event event1 = Event.builder().timestamp(1).build();
+ Event event2 = Event.builder().timestamp(2).build();
+ Event event3 = eventWithPromptTokenCount(3, 200);
session =
Session.builder("test-session").events(ImmutableList.of(event1, event2, event3)).build();
@@ -117,11 +116,7 @@ public void processRequest_withConfig_skipsCompactionIfBelowThreshold() {
when(context.eventsCompactionConfig()).thenReturn(Optional.of(config));
// Setup events with usage 200 (below 500)
- Event event3 = mock(Event.class);
- when(event3.usageMetadata())
- .thenReturn(
- Optional.of(
- GenerateContentResponseUsageMetadata.builder().promptTokenCount(200).build()));
+ Event event3 = eventWithPromptTokenCount(3, 200);
session = Session.builder("test-session").events(ImmutableList.of(event3)).build();
when(context.session()).thenReturn(session);
@@ -137,6 +132,41 @@ public void processRequest_withConfig_skipsCompactionIfBelowThreshold() {
verify(sessionService, never()).appendEvent(any(), any());
}
+ @Test
+ public void processRequest_withConfig_toolResponseAfterTokenUsage_triggersCompaction() {
+ // Threshold 300: the reported count of 200 alone is within it, but the function response
+ // appended after it (an 800-character result, over 200 tokens) pushes the estimate above.
+ EventsCompactionConfig config = new EventsCompactionConfig(5, 1, summarizer, 300, 2);
+ when(context.eventsCompactionConfig()).thenReturn(Optional.of(config));
+
+ Event event1 = Event.builder().timestamp(1).build();
+ Event event2 = eventWithPromptTokenCount(2, 200);
+ Event event3 =
+ Event.builder()
+ .timestamp(3)
+ .content(
+ Content.fromParts(
+ Part.fromFunctionResponse("tool", ImmutableMap.of("result", "x".repeat(800)))))
+ .build();
+
+ session =
+ Session.builder("test-session").events(ImmutableList.of(event1, event2, event3)).build();
+ when(context.session()).thenReturn(session);
+
+ Event summaryEvent = mock(Event.class);
+ when(summarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
+ when(sessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));
+
+ Compaction compaction = new Compaction();
+ compaction
+ .processRequest(context, request)
+ .test()
+ .assertNoErrors()
+ .assertValue(r -> r.updatedRequest() == request);
+
+ verify(sessionService).appendEvent(eq(session), eq(summaryEvent));
+ }
+
@Test
public void processRequest_withConfig_nullRetentionSize_doesNothing() {
// Setup config with retentionSize = null
@@ -154,4 +184,14 @@ public void processRequest_withConfig_nullRetentionSize_doesNothing() {
verify(sessionService, never()).appendEvent(any(), any());
verify(context, never()).session();
}
+
+ private static Event eventWithPromptTokenCount(long timestamp, int promptTokenCount) {
+ return Event.builder()
+ .timestamp(timestamp)
+ .usageMetadata(
+ GenerateContentResponseUsageMetadata.builder()
+ .promptTokenCount(promptTokenCount)
+ .build())
+ .build();
+ }
}
diff --git a/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java b/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java
index 7a4a3ddb9..c3d8840a4 100644
--- a/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java
+++ b/core/src/test/java/com/google/adk/summarizer/TailRetentionEventCompactorTest.java
@@ -30,6 +30,7 @@
import com.google.adk.sessions.BaseSessionService;
import com.google.adk.sessions.Session;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
import com.google.genai.types.Content;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import com.google.genai.types.Part;
@@ -160,6 +161,118 @@ public void compaction_happensWhenTokenUsageAboveThreshold() {
verify(mockSessionService).appendEvent(eq(session), eq(summaryEvent));
}
+ @Test
+ public void compaction_happensWhenEventsAfterTokenUsagePushAboveThreshold() {
+ // Threshold is 300, the last model call used 200 prompt tokens, and a tool response of
+ // 800 characters (about 200 tokens) arrived after it: 200 + 200 > 300.
+ EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
+ ImmutableList events =
+ ImmutableList.of(
+ createEvent(1, "Event1"),
+ withUsage(createEvent(2, "Call"), 200),
+ createEvent(3, "x".repeat(800)));
+ Session session = Session.builder("id").events(events).build();
+ Event summaryEvent = createEvent(4, "Summary");
+
+ when(mockSummarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
+ when(mockSessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));
+
+ compactor.compact(session, mockSessionService).blockingSubscribe();
+
+ verify(mockSummarizer).summarizeEvents(eventListCaptor.capture());
+ assertThat(
+ eventListCaptor.getValue().stream().map(TailRetentionEventCompactorTest::getPromptText))
+ .containsExactly("Event1", "Call")
+ .inOrder();
+ verify(mockSessionService).appendEvent(eq(session), eq(summaryEvent));
+ }
+
+ @Test
+ public void compaction_happensWhenToolResponseAfterTokenUsagePushesAboveThreshold() {
+ // Threshold is 300, the last model call used 200 prompt tokens, and a function response with
+ // an 800-character result (over 200 tokens once stringified) arrived after it.
+ EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
+ ImmutableList events =
+ ImmutableList.of(
+ createEvent(1, "Event1"),
+ withUsage(createEvent(2, "Call"), 200),
+ createToolResponseEvent(3, "x".repeat(800)));
+ Session session = Session.builder("id").events(events).build();
+ Event summaryEvent = createEvent(4, "Summary");
+
+ when(mockSummarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
+ when(mockSessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));
+
+ compactor.compact(session, mockSessionService).blockingSubscribe();
+
+ verify(mockSummarizer).summarizeEvents(eventListCaptor.capture());
+ assertThat(
+ eventListCaptor.getValue().stream().map(TailRetentionEventCompactorTest::getPromptText))
+ .containsExactly("Event1", "Call")
+ .inOrder();
+ verify(mockSessionService).appendEvent(eq(session), eq(summaryEvent));
+ }
+
+ @Test
+ public void compaction_skippedWhenEventsAfterTokenUsageStayBelowThreshold() {
+ // Threshold is 300, the last model call used 200 prompt tokens, and a tool response of
+ // 200 characters (about 50 tokens) arrived after it: 200 + 50 <= 300.
+ EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
+ ImmutableList events =
+ ImmutableList.of(
+ createEvent(1, "Event1"),
+ withUsage(createEvent(2, "Call"), 200),
+ createEvent(3, "x".repeat(200)));
+ Session session = Session.builder("id").events(events).build();
+
+ compactor.compact(session, mockSessionService).blockingSubscribe();
+
+ verify(mockSummarizer, never()).summarizeEvents(any());
+ verify(mockSessionService, never()).appendEvent(any(), any());
+ }
+
+ @Test
+ public void compaction_countsTheResponseThatReportedTokenUsage() {
+ // promptTokenCount excludes the response that reports it. Its 40 characters (~10 tokens) are
+ // what push 200 + 100 (the 400-character tool response) past the threshold of 300.
+ EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
+ ImmutableList events =
+ ImmutableList.of(
+ createEvent(1, "Event1"),
+ withUsage(createEvent(2, "y".repeat(40)), 200),
+ createEvent(3, "x".repeat(400)));
+ Session session = Session.builder("id").events(events).build();
+ Event summaryEvent = createEvent(4, "Summary");
+
+ when(mockSummarizer.summarizeEvents(any())).thenReturn(Maybe.just(summaryEvent));
+ when(mockSessionService.appendEvent(any(), any())).thenReturn(Single.just(summaryEvent));
+
+ compactor.compact(session, mockSessionService).blockingSubscribe();
+
+ verify(mockSummarizer).summarizeEvents(any());
+ }
+
+ @Test
+ public void compaction_ignoresTokenUsageReportedBeforeLatestCompaction() {
+ // The 1000-token count describes the prompt that the compaction replaced. What remains
+ // ("Summary", "Retain1", "Retain2") is ~5 estimated tokens, within the threshold of 300.
+ EventCompactor compactor = new TailRetentionEventCompactor(mockSummarizer, 1, 300);
+ ImmutableList events =
+ ImmutableList.of(
+ createEvent(1, "Event1"),
+ withUsage(createEvent(2, "Call"), 1000),
+ createEvent(3, "Retain1"),
+ createCompactedEvent(
+ /* startTimestamp= */ 1, /* endTimestamp= */ 2, "Summary", /* eventTimestamp= */ 4),
+ createEvent(5, "Retain2"));
+ Session session = Session.builder("id").events(events).build();
+
+ compactor.compact(session, mockSessionService).blockingSubscribe();
+
+ verify(mockSummarizer, never()).summarizeEvents(any());
+ verify(mockSessionService, never()).appendEvent(any(), any());
+ }
+
@Test
public void compact_notEnoughEvents_doesNothing() {
ImmutableList events =
@@ -330,6 +443,14 @@ private static Event createEvent(long timestamp, String text) {
.build();
}
+ private static Event createToolResponseEvent(long timestamp, String result) {
+ return Event.builder()
+ .timestamp(timestamp)
+ .content(
+ Content.fromParts(Part.fromFunctionResponse("tool", ImmutableMap.of("result", result))))
+ .build();
+ }
+
private static String getPromptText(Event event) {
return event
.content()