Skip to content

Commit 9c771ee

Browse files
wikaaaaacopybara-github
wikaaaaa
authored and
copybara-github
committed
feat: apply context-compaction summaries when building LLM request contents
PiperOrigin-RevId: 936017629
1 parent 3e7b94e commit 9c771ee

2 files changed

Lines changed: 318 additions & 2 deletions

File tree

‎core/src/commonMain/kotlin/com/google/adk/kt/processors/HistoryRewriterProcessor.kt‎

Lines changed: 81 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -51,17 +51,27 @@ internal class HistoryRewriterProcessor {
5151
shouldIncludeEventInContext(currentBranch, it)
5252
}
5353

54-
// Process events
54+
// Process events. Compaction events are kept here (they carry their summary in
55+
// actions.compaction rather than content) so processCompactionEvents can expand them below.
5556
val filteredEvents = rawFilteredEvents.mapNotNull { event ->
5657
when {
58+
event.actions.compaction != null -> event
5759
event.content == null -> null
5860
isOtherAgentReply(agentName, event) -> presentOtherAgentMessage(event)
5961
else -> event
6062
}
6163
}
6264

65+
// Replace each compaction event with its summary and drop the raw events it covers.
66+
val eventsWithCompactionApplied =
67+
if (filteredEvents.any { it.actions.compaction != null }) {
68+
processCompactionEvents(filteredEvents)
69+
} else {
70+
filteredEvents
71+
}
72+
6373
// Rearrange for latest function response (merge scenarios) and async function responses
64-
return filteredEvents
74+
return eventsWithCompactionApplied
6575
.let { rearrangeEventsForLatestFunctionResponse(it) }
6676
.let { rearrangeEventsForAsyncFunctionResponsesInHistory(it) }
6777
.mapNotNull { event ->
@@ -70,6 +80,72 @@ internal class HistoryRewriterProcessor {
7080
}
7181
}
7282

83+
/**
84+
* Processes events by applying compaction. Identifies compacted ranges and filters out events
85+
* that are covered by compaction summaries.
86+
*
87+
* @param events The list of events to process.
88+
* @return The list of events with compaction applied.
89+
*/
90+
private fun processCompactionEvents(events: List<Event>): List<Event> {
91+
// Extract all compaction ranges from the events.
92+
val compactionRanges = events.mapIndexedNotNull { index, event ->
93+
event.actions.compaction?.let { CompactionRange(index, it.startTimestamp, it.endTimestamp) }
94+
}
95+
val coveredIndices = coveredCompactionRangeIndices(compactionRanges)
96+
val keptCompactionRanges = compactionRanges.filter { it.index !in coveredIndices }
97+
98+
data class Item(val timestamp: Long, val index: Int, val event: Event)
99+
100+
val finalItems = mutableListOf<Item>()
101+
102+
// Pass 1: append all kept compaction events.
103+
for (range in keptCompactionRanges) {
104+
val compaction = events[range.index].actions.compaction!!
105+
finalItems.add(
106+
Item(
107+
compaction.endTimestamp,
108+
range.index,
109+
events[range.index].copy(
110+
author = Role.MODEL,
111+
content = compaction.compactedContent,
112+
timestamp = compaction.endTimestamp,
113+
),
114+
)
115+
)
116+
}
117+
118+
// Pass 2: append raw (non-compaction) events that don't fall into a kept compaction range.
119+
finalItems +=
120+
events
121+
.withIndex()
122+
.filter { (_, event) -> event.actions.compaction == null }
123+
.filter { (_, event) -> keptCompactionRanges.none { event.timestamp in it.start..it.end } }
124+
.map { (index, event) -> Item(event.timestamp, index, event) }
125+
126+
return finalItems.sortedWith(compareBy({ it.timestamp }, { it.index })).map { it.event }
127+
}
128+
129+
/**
130+
* Returns the indices of [ranges] that are fully contained by another range. When two ranges are
131+
* identical only the later one is kept; partially overlapping ranges (neither containing the
132+
* other) are both kept.
133+
*/
134+
private fun coveredCompactionRangeIndices(ranges: List<CompactionRange>): Set<Int> =
135+
ranges.filter { range -> ranges.any { it.covers(range) } }.map { it.index }.toSet()
136+
137+
private data class CompactionRange(val index: Int, val start: Long, val end: Long) {
138+
/**
139+
* True if this range fully contains [other] -- strictly larger on at least one side, or
140+
* identical but appearing later (so equal ranges keep only the most recent).
141+
*/
142+
fun covers(other: CompactionRange): Boolean =
143+
index != other.index &&
144+
start <= other.start &&
145+
end >= other.end &&
146+
(start < other.start || end > other.end || index > other.index)
147+
}
148+
73149
/**
74150
* Returns the suffix of [events] that belongs to the current turn.
75151
*
@@ -193,6 +269,9 @@ internal class HistoryRewriterProcessor {
193269
* Parts with only thoughts are also considered empty.
194270
*/
195271
private fun containsEmptyContent(event: Event): Boolean {
272+
// Compaction events carry their summary in actions.compaction rather than content; keep them so
273+
// processCompactionEvents can expand them into summary content.
274+
if (event.actions.compaction != null) return false
196275

197276
val hasContent =
198277
event.content != null &&

‎core/src/commonTest/kotlin/com/google/adk/kt/processors/ContentsProcessorTest.kt‎

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import com.google.adk.kt.sessions.InMemorySessionService
2525
import com.google.adk.kt.sessions.SessionKey
2626
import com.google.adk.kt.testing.DummyAgent
2727
import com.google.adk.kt.testing.DummyModel
28+
import com.google.adk.kt.testing.compactionEvent
2829
import com.google.adk.kt.testing.modelMessage
2930
import com.google.adk.kt.testing.testSession
3031
import com.google.adk.kt.testing.userMessage
@@ -1183,6 +1184,242 @@ class ContentsProcessorTest {
11831184
assertThat(result).isEmpty()
11841185
}
11851186

1187+
@Test
1188+
fun process_compactionEvent_replacesCoveredEventsWithSummary() = runTest {
1189+
val processor = ContentsProcessor()
1190+
var request = LlmRequest(contents = emptyList())
1191+
val context =
1192+
createLlmAgentTestContext(
1193+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1194+
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
1195+
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
1196+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 4L, summary = "summary"),
1197+
)
1198+
1199+
request = processor.process(context, request)
1200+
1201+
// u1(ts=1) and m1(ts=2) fall in [1,2] -> replaced by the summary (at ts=2); u2(ts=3) is kept.
1202+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1203+
.containsExactly("summary", "u2")
1204+
.inOrder()
1205+
}
1206+
1207+
@Test
1208+
fun process_nestedCompactions_keepsOnlyOuterSummary() = runTest {
1209+
val processor = ContentsProcessor()
1210+
var request = LlmRequest(contents = emptyList())
1211+
val context =
1212+
createLlmAgentTestContext(
1213+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1214+
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
1215+
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
1216+
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 4L),
1217+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 5L, summary = "inner"),
1218+
compactionEvent(startTs = 1L, endTs = 4L, timestamp = 6L, summary = "outer"),
1219+
Event(author = "user", content = userMessage("u3"), timestamp = 7L),
1220+
)
1221+
1222+
request = processor.process(context, request)
1223+
1224+
// [1,2] is contained in [1,4], so only "outer" survives (covering u1..m2); u3 is kept.
1225+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1226+
.containsExactly("outer", "u3")
1227+
.inOrder()
1228+
}
1229+
1230+
@Test
1231+
fun process_partiallyOverlappingCompactions_keepsBoth() = runTest {
1232+
val processor = ContentsProcessor()
1233+
var request = LlmRequest(contents = emptyList())
1234+
val context =
1235+
createLlmAgentTestContext(
1236+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1237+
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
1238+
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
1239+
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 4L),
1240+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 5L, summary = "first"),
1241+
compactionEvent(startTs = 2L, endTs = 4L, timestamp = 6L, summary = "second"),
1242+
Event(author = "user", content = userMessage("u3"), timestamp = 7L),
1243+
)
1244+
1245+
request = processor.process(context, request)
1246+
1247+
// [1,2] and [2,4] overlap but neither contains the other, so both summaries are kept.
1248+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1249+
.containsExactly("first", "second", "u3")
1250+
.inOrder()
1251+
}
1252+
1253+
@Test
1254+
fun process_noCompaction_returnsEventsUnchanged() = runTest {
1255+
val processor = ContentsProcessor()
1256+
var request = LlmRequest(contents = emptyList())
1257+
val context =
1258+
createLlmAgentTestContext(
1259+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1260+
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
1261+
Event(author = "user", content = userMessage("u2"), timestamp = 3L),
1262+
)
1263+
1264+
request = processor.process(context, request)
1265+
1266+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1267+
.containsExactly("u1", "m1", "u2")
1268+
.inOrder()
1269+
}
1270+
1271+
@Test
1272+
fun process_noCompaction_preservesOriginalEventOrder() = runTest {
1273+
val processor = ContentsProcessor()
1274+
var request = LlmRequest(contents = emptyList())
1275+
val context =
1276+
createLlmAgentTestContext(
1277+
Event(author = "user", content = userMessage("first"), timestamp = 2L),
1278+
Event(author = "testAgent", content = modelMessage("second"), timestamp = 1L),
1279+
)
1280+
1281+
request = processor.process(context, request)
1282+
1283+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1284+
.containsExactly("first", "second")
1285+
.inOrder()
1286+
}
1287+
1288+
@Test
1289+
fun process_compactionAtBeginning_keepsLaterEvents() = runTest {
1290+
val processor = ContentsProcessor()
1291+
var request = LlmRequest(contents = emptyList())
1292+
val context =
1293+
createLlmAgentTestContext(
1294+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 2L, summary = "summary"),
1295+
Event(author = "user", content = userMessage("u3"), timestamp = 3L),
1296+
Event(author = "testAgent", content = modelMessage("m4"), timestamp = 4L),
1297+
)
1298+
1299+
request = processor.process(context, request)
1300+
1301+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1302+
.containsExactly("summary", "u3", "m4")
1303+
.inOrder()
1304+
}
1305+
1306+
@Test
1307+
fun process_compactionAtEnd_keepsEarlierRawEvents() = runTest {
1308+
val processor = ContentsProcessor()
1309+
var request = LlmRequest(contents = emptyList())
1310+
val context =
1311+
createLlmAgentTestContext(
1312+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1313+
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 2L),
1314+
Event(author = "user", content = userMessage("u3"), timestamp = 3L),
1315+
compactionEvent(startTs = 2L, endTs = 3L, timestamp = 4L, summary = "summary"),
1316+
)
1317+
1318+
request = processor.process(context, request)
1319+
1320+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1321+
.containsExactly("u1", "summary")
1322+
.inOrder()
1323+
}
1324+
1325+
@Test
1326+
fun process_twoAdjacentCompactions_keepBothSummaries() = runTest {
1327+
val processor = ContentsProcessor()
1328+
var request = LlmRequest(contents = emptyList())
1329+
val context =
1330+
createLlmAgentTestContext(
1331+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1332+
Event(author = "testAgent", content = modelMessage("m2"), timestamp = 2L),
1333+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 2L, summary = "summary1to2"),
1334+
Event(author = "user", content = userMessage("u3"), timestamp = 3L),
1335+
Event(author = "testAgent", content = modelMessage("m4"), timestamp = 4L),
1336+
compactionEvent(startTs = 3L, endTs = 4L, timestamp = 4L, summary = "summary3to4"),
1337+
Event(author = "user", content = userMessage("u5"), timestamp = 5L),
1338+
)
1339+
1340+
request = processor.process(context, request)
1341+
1342+
// [1,2] and [3,4] each replace their range; u5 (ts=5) is uncovered and kept.
1343+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1344+
.containsExactly("summary1to2", "summary3to4", "u5")
1345+
.inOrder()
1346+
}
1347+
1348+
@Test
1349+
fun process_multipleCompactions_replaceRangesAndKeepRawEventsBetween() = runTest {
1350+
val processor = ContentsProcessor()
1351+
var request = LlmRequest(contents = emptyList())
1352+
val context =
1353+
createLlmAgentTestContext(
1354+
Event(author = "user", content = userMessage("e1"), timestamp = 1L),
1355+
Event(author = "user", content = userMessage("e2"), timestamp = 2L),
1356+
Event(author = "user", content = userMessage("e3"), timestamp = 3L),
1357+
Event(author = "user", content = userMessage("e4"), timestamp = 4L),
1358+
compactionEvent(startTs = 1L, endTs = 4L, timestamp = 4L, summary = "summary1to4"),
1359+
Event(author = "user", content = userMessage("e5"), timestamp = 5L),
1360+
Event(author = "user", content = userMessage("e6"), timestamp = 6L),
1361+
Event(author = "user", content = userMessage("e7"), timestamp = 7L),
1362+
Event(author = "user", content = userMessage("e8"), timestamp = 8L),
1363+
Event(author = "user", content = userMessage("e9"), timestamp = 9L),
1364+
compactionEvent(startTs = 6L, endTs = 9L, timestamp = 9L, summary = "summary6to9"),
1365+
Event(author = "user", content = userMessage("e10"), timestamp = 10L),
1366+
)
1367+
1368+
request = processor.process(context, request)
1369+
1370+
// [1,4] and [6,9] are replaced by their summaries; e5 (gap) and e10 (after) are kept.
1371+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1372+
.containsExactly("summary1to4", "e5", "summary6to9", "e10")
1373+
.inOrder()
1374+
}
1375+
1376+
@Test
1377+
fun process_compactionAppendedLate_keepsNewerEvents() = runTest {
1378+
val processor = ContentsProcessor()
1379+
var request = LlmRequest(contents = emptyList())
1380+
val context =
1381+
createLlmAgentTestContext(
1382+
Event(author = "user", content = userMessage("e1"), timestamp = 1L),
1383+
Event(author = "user", content = userMessage("e2"), timestamp = 2L),
1384+
Event(author = "user", content = userMessage("e3"), timestamp = 3L),
1385+
Event(author = "user", content = userMessage("u4"), timestamp = 4L),
1386+
Event(author = "testAgent", content = modelMessage("m5"), timestamp = 5L),
1387+
compactionEvent(startTs = 1L, endTs = 3L, timestamp = 6L, summary = "summary1to3"),
1388+
)
1389+
1390+
request = processor.process(context, request)
1391+
1392+
// The compaction covers [1,3] but was appended at ts=6; u4,m5 (after the range) survive, and
1393+
// the summary is positioned at its end timestamp (3) -- ahead of them -- not at the append
1394+
// time.
1395+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1396+
.containsExactly("summary1to3", "u4", "m5")
1397+
.inOrder()
1398+
}
1399+
1400+
@Test
1401+
fun process_duplicateRangeCompactions_keepsOnlyMostRecentSummary() = runTest {
1402+
val processor = ContentsProcessor()
1403+
var request = LlmRequest(contents = emptyList())
1404+
val context =
1405+
createLlmAgentTestContext(
1406+
Event(author = "user", content = userMessage("u1"), timestamp = 1L),
1407+
Event(author = "testAgent", content = modelMessage("m1"), timestamp = 2L),
1408+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 3L, summary = "old_summary"),
1409+
compactionEvent(startTs = 1L, endTs = 2L, timestamp = 4L, summary = "new_summary"),
1410+
Event(author = "user", content = userMessage("u2"), timestamp = 5L),
1411+
)
1412+
1413+
request = processor.process(context, request)
1414+
1415+
// Two compactions cover the identical range [1,2]; the tie-break keeps only the later one
1416+
// ("new_summary") and drops the earlier ("old_summary"). Without the tie-break both would be
1417+
// marked as covering each other and dropped, leaving u1/m1 unsummarized.
1418+
assertThat(request.contents.map { it.parts.firstOrNull()?.text })
1419+
.containsExactly("new_summary", "u2")
1420+
.inOrder()
1421+
}
1422+
11861423
// Helpers
11871424

11881425
private suspend fun createTestContext(

0 commit comments

Comments
 (0)