Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -51,17 +51,27 @@ internal class HistoryRewriterProcessor {
shouldIncludeEventInContext(currentBranch, it)
}

// Process events
// Process events. Compaction events are kept here (they carry their summary in
// actions.compaction rather than content) so processCompactionEvents can expand them below.
val filteredEvents = rawFilteredEvents.mapNotNull { event ->
when {
event.actions.compaction != null -> event
event.content == null -> null
isOtherAgentReply(agentName, event) -> presentOtherAgentMessage(event)
else -> event
}
}

// Replace each compaction event with its summary and drop the raw events it covers.
val eventsWithCompactionApplied =
if (filteredEvents.any { it.actions.compaction != null }) {
processCompactionEvents(filteredEvents)
} else {
filteredEvents
}

// Rearrange for latest function response (merge scenarios) and async function responses
return filteredEvents
return eventsWithCompactionApplied
.let { rearrangeEventsForLatestFunctionResponse(it) }
.let { rearrangeEventsForAsyncFunctionResponsesInHistory(it) }
.mapNotNull { event ->
Expand All @@ -70,6 +80,72 @@ internal class HistoryRewriterProcessor {
}
}

/**
* Processes events by applying compaction. Identifies compacted ranges and filters out events
* that are covered by compaction summaries.
*
* @param events The list of events to process.
* @return The list of events with compaction applied.
*/
private fun processCompactionEvents(events: List<Event>): List<Event> {
// Extract all compaction ranges from the events.
val compactionRanges = events.mapIndexedNotNull { index, event ->
event.actions.compaction?.let { CompactionRange(index, it.startTimestamp, it.endTimestamp) }
}
val coveredIndices = coveredCompactionRangeIndices(compactionRanges)
val keptCompactionRanges = compactionRanges.filter { it.index !in coveredIndices }

data class Item(val timestamp: Long, val index: Int, val event: Event)

val finalItems = mutableListOf<Item>()

// Pass 1: append all kept compaction events.
for (range in keptCompactionRanges) {
val compaction = events[range.index].actions.compaction!!
finalItems.add(
Item(
compaction.endTimestamp,
range.index,
events[range.index].copy(
author = Role.MODEL,
content = compaction.compactedContent,
timestamp = compaction.endTimestamp,
),
)
)
}

// Pass 2: append raw (non-compaction) events that don't fall into a kept compaction range.
finalItems +=
events
.withIndex()
.filter { (_, event) -> event.actions.compaction == null }
.filter { (_, event) -> keptCompactionRanges.none { event.timestamp in it.start..it.end } }
.map { (index, event) -> Item(event.timestamp, index, event) }

return finalItems.sortedWith(compareBy({ it.timestamp }, { it.index })).map { it.event }
}

/**
* Returns the indices of [ranges] that are fully contained by another range. When two ranges are
* identical only the later one is kept; partially overlapping ranges (neither containing the
* other) are both kept.
*/
private fun coveredCompactionRangeIndices(ranges: List<CompactionRange>): Set<Int> =
ranges.filter { range -> ranges.any { it.covers(range) } }.map { it.index }.toSet()

private data class CompactionRange(val index: Int, val start: Long, val end: Long) {
/**
* True if this range fully contains [other] -- strictly larger on at least one side, or
* identical but appearing later (so equal ranges keep only the most recent).
*/
fun covers(other: CompactionRange): Boolean =
index != other.index &&
start <= other.start &&
end >= other.end &&
(start < other.start || end > other.end || index > other.index)
}

/**
* Returns the suffix of [events] that belongs to the current turn.
*
Expand Down Expand Up @@ -193,6 +269,9 @@ internal class HistoryRewriterProcessor {
* Parts with only thoughts are also considered empty.
*/
private fun containsEmptyContent(event: Event): Boolean {
// Compaction events carry their summary in actions.compaction rather than content; keep them so
// processCompactionEvents can expand them into summary content.
if (event.actions.compaction != null) return false

val hasContent =
event.content != null &&
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import com.google.adk.kt.sessions.InMemorySessionService
import com.google.adk.kt.sessions.SessionKey
import com.google.adk.kt.testing.DummyAgent
import com.google.adk.kt.testing.DummyModel
import com.google.adk.kt.testing.compactionEvent
import com.google.adk.kt.testing.modelMessage
import com.google.adk.kt.testing.testSession
import com.google.adk.kt.testing.userMessage
Expand Down Expand Up @@ -1183,6 +1184,242 @@ class ContentsProcessorTest {
assertThat(result).isEmpty()
}

@Test
fun process_compactionEvent_replacesCoveredEventsWithSummary() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 4L, summary = "summary"),
)

request = processor.process(context, request)

// u1(ts=1) and m1(ts=2) fall in [1,2] -> replaced by the summary (at ts=2); u2(ts=3) is kept.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("summary", "u2")
.inOrder()
}

@Test
fun process_nestedCompactions_keepsOnlyOuterSummary() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 4L),
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 5L, summary = "inner"),
compactionEvent(startTs = 1L, endTs = 4L, timestamp = 6L, summary = "outer"),
Event(author = "user", content = userMessage("u3"), timestamp = 7L),
)

request = processor.process(context, request)

// [1,2] is contained in [1,4], so only "outer" survives (covering u1..m2); u3 is kept.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("outer", "u3")
.inOrder()
}

@Test
fun process_partiallyOverlappingCompactions_keepsBoth() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 4L),
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 5L, summary = "first"),
compactionEvent(startTs = 2L, endTs = 4L, timestamp = 6L, summary = "second"),
Event(author = "user", content = userMessage("u3"), timestamp = 7L),
)

request = processor.process(context, request)

// [1,2] and [2,4] overlap but neither contains the other, so both summaries are kept.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("first", "second", "u3")
.inOrder()
}

@Test
fun process_noCompaction_returnsEventsUnchanged() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
)

request = processor.process(context, request)

assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("u1", "m1", "u2")
.inOrder()
}

@Test
fun process_noCompaction_preservesOriginalEventOrder() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("first"), timestamp = 2L),
Event(author = "testAgent", content = modelMessage("second"), timestamp = 1L),
)

request = processor.process(context, request)

assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("first", "second")
.inOrder()
}

@Test
fun process_compactionAtBeginning_keepsLaterEvents() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 2L, summary = "summary"),
Event(author = "user", content = userMessage("u3"), timestamp = 3L),
Event(author = "testAgent", content = modelMessage("m4"), timestamp = 4L),
)

request = processor.process(context, request)

assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("summary", "u3", "m4")
.inOrder()
}

@Test
fun process_compactionAtEnd_keepsEarlierRawEvents() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 2L),
Event(author = "user", content = userMessage("u3"), timestamp = 3L),
compactionEvent(startTs = 2L, endTs = 3L, timestamp = 4L, summary = "summary"),
)

request = processor.process(context, request)

assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("u1", "summary")
.inOrder()
}

@Test
fun process_twoAdjacentCompactions_keepBothSummaries() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 2L),
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 2L, summary = "summary1to2"),
Event(author = "user", content = userMessage("u3"), timestamp = 3L),
Event(author = "testAgent", content = modelMessage("m4"), timestamp = 4L),
compactionEvent(startTs = 3L, endTs = 4L, timestamp = 4L, summary = "summary3to4"),
Event(author = "user", content = userMessage("u5"), timestamp = 5L),
)

request = processor.process(context, request)

// [1,2] and [3,4] each replace their range; u5 (ts=5) is uncovered and kept.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("summary1to2", "summary3to4", "u5")
.inOrder()
}

@Test
fun process_multipleCompactions_replaceRangesAndKeepRawEventsBetween() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("e1"), timestamp = 1L),
Event(author = "user", content = userMessage("e2"), timestamp = 2L),
Event(author = "user", content = userMessage("e3"), timestamp = 3L),
Event(author = "user", content = userMessage("e4"), timestamp = 4L),
compactionEvent(startTs = 1L, endTs = 4L, timestamp = 4L, summary = "summary1to4"),
Event(author = "user", content = userMessage("e5"), timestamp = 5L),
Event(author = "user", content = userMessage("e6"), timestamp = 6L),
Event(author = "user", content = userMessage("e7"), timestamp = 7L),
Event(author = "user", content = userMessage("e8"), timestamp = 8L),
Event(author = "user", content = userMessage("e9"), timestamp = 9L),
compactionEvent(startTs = 6L, endTs = 9L, timestamp = 9L, summary = "summary6to9"),
Event(author = "user", content = userMessage("e10"), timestamp = 10L),
)

request = processor.process(context, request)

// [1,4] and [6,9] are replaced by their summaries; e5 (gap) and e10 (after) are kept.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("summary1to4", "e5", "summary6to9", "e10")
.inOrder()
}

@Test
fun process_compactionAppendedLate_keepsNewerEvents() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("e1"), timestamp = 1L),
Event(author = "user", content = userMessage("e2"), timestamp = 2L),
Event(author = "user", content = userMessage("e3"), timestamp = 3L),
Event(author = "user", content = userMessage("u4"), timestamp = 4L),
Event(author = "testAgent", content = modelMessage("m5"), timestamp = 5L),
compactionEvent(startTs = 1L, endTs = 3L, timestamp = 6L, summary = "summary1to3"),
)

request = processor.process(context, request)

// The compaction covers [1,3] but was appended at ts=6; u4,m5 (after the range) survive, and
// the summary is positioned at its end timestamp (3) -- ahead of them -- not at the append
// time.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("summary1to3", "u4", "m5")
.inOrder()
}

@Test
fun process_duplicateRangeCompactions_keepsOnlyMostRecentSummary() = runTest {
val processor = ContentsProcessor()
var request = LlmRequest(contents = emptyList())
val context =
createLlmAgentTestContext(
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 3L, summary = "old_summary"),
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 4L, summary = "new_summary"),
Event(author = "user", content = userMessage("u2"), timestamp = 5L),
)

request = processor.process(context, request)

// Two compactions cover the identical range [1,2]; the tie-break keeps only the later one
// ("new_summary") and drops the earlier ("old_summary"). Without the tie-break both would be
// marked as covering each other and dropped, leaving u1/m1 unsummarized.
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
.containsExactly("new_summary", "u2")
.inOrder()
}

// Helpers

private suspend fun createTestContext(
Expand Down
Loading