Skip to content

Commit 4feb461

Browse files
wikaaaaacopybara-github
wikaaaaa
authored and
copybara-github
committed
feat: run token-threshold compaction before model calls
PiperOrigin-RevId: 947021398
1 parent a667bcf commit 4feb461

8 files changed

Lines changed: 739 additions & 2 deletions

File tree

‎core/src/commonMain/kotlin/com/google/adk/kt/agents/InvocationContext.kt‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ import com.google.adk.kt.memory.MemoryService
3131
import com.google.adk.kt.plugins.PluginManager
3232
import com.google.adk.kt.sessions.Session
3333
import com.google.adk.kt.sessions.SessionService
34+
import com.google.adk.kt.summarizer.EventsCompactionConfig
3435
import com.google.adk.kt.telemetry.EMPTY_JSON
3536
import com.google.adk.kt.telemetry.Span
3637
import com.google.adk.kt.telemetry.TelemetryAttributes
@@ -122,6 +123,14 @@ data class InvocationContext(
122123
*/
123124
val resumabilityConfig: ResumabilityConfig? = null,
124125

126+
/**
127+
* Optional event-compaction configuration for this invocation.
128+
*
129+
* Threaded from the runner's [com.google.adk.kt.apps.App] so intra-invocation request processors
130+
* (e.g. token-threshold compaction) can read it. `null` when no compaction is configured.
131+
*/
132+
val eventsCompactionConfig: EventsCompactionConfig? = null,
133+
125134
// State
126135
/** The user content that started this invocation. Readonly. */
127136
val userContent: Content? = null,

‎core/src/commonMain/kotlin/com/google/adk/kt/agents/LlmAgent.kt‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import com.google.adk.kt.logging.LoggerFactory
3030
import com.google.adk.kt.models.Model
3131
import com.google.adk.kt.processors.AgentTransferProcessor
3232
import com.google.adk.kt.processors.BasicRequestProcessor
33+
import com.google.adk.kt.processors.CompactionRequestProcessor
3334
import com.google.adk.kt.processors.ContentsProcessor
3435
import com.google.adk.kt.processors.InstructionsProcessor
3536
import com.google.adk.kt.processors.LlmRequestProcessor
@@ -188,6 +189,9 @@ class LlmAgent(
188189
BasicRequestProcessor(),
189190
RequestConfirmationProcessor(),
190191
InstructionsProcessor(),
192+
// Compaction should run before contents so compacted events are reflected in the model
193+
// request context.
194+
CompactionRequestProcessor(),
191195
ContentsProcessor(),
192196
AgentTransferProcessor(),
193197
OutputSchemaProcessor(),

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

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.adk.kt.processors
17+
18+
import com.google.adk.kt.agents.InvocationContext
19+
import com.google.adk.kt.events.Event
20+
import com.google.adk.kt.models.LlmRequest
21+
import com.google.adk.kt.summarizer.TokenThresholdEventCompactor
22+
23+
/**
24+
* Runs token-threshold event compaction before the conversation history is assembled for a model
25+
* call.
26+
*
27+
* When the invocation's [InvocationContext.eventsCompactionConfig] has token-threshold compaction
28+
* configured and the most recently observed prompt token count has reached the threshold, this
29+
* appends a compaction summary event to the session. Because it runs before [ContentsProcessor],
30+
* the freshly appended summary is reflected in the contents built for the request. The [LlmRequest]
31+
* itself is returned unchanged.
32+
*/
33+
internal class CompactionRequestProcessor : LlmRequestProcessor {
34+
override suspend fun process(
35+
context: InvocationContext,
36+
request: LlmRequest,
37+
emitEvent: suspend (Event) -> Unit,
38+
): LlmRequest {
39+
val config = context.eventsCompactionConfig ?: return request
40+
if (!config.hasTokenThresholdConfig()) return request
41+
val sessionService = context.sessionService ?: return request
42+
TokenThresholdEventCompactor(config, context.agent.name, context.branch)
43+
.compact(context.session, sessionService)
44+
return request
45+
}
46+
}

‎core/src/commonMain/kotlin/com/google/adk/kt/runners/AbstractRunner.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -543,6 +543,7 @@ abstract class AbstractRunner : Runner {
543543
userContent = newMessage,
544544
pluginManager = pluginManager,
545545
resumabilityConfig = resumabilityConfig,
546+
eventsCompactionConfig = app?.eventsCompactionConfig,
546547
)
547548
.let {
548549
// Run callbacks and append user message to session
@@ -598,6 +599,7 @@ abstract class AbstractRunner : Runner {
598599
userContent = userMessage,
599600
pluginManager = pluginManager,
600601
resumabilityConfig = resumabilityConfig,
602+
eventsCompactionConfig = app?.eventsCompactionConfig,
601603
)
602604

603605
val currentContext =

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

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.adk.kt.processors
17+
18+
import com.google.adk.kt.agents.InvocationContext
19+
import com.google.adk.kt.events.Event
20+
import com.google.adk.kt.models.LlmRequest
21+
import com.google.adk.kt.sessions.GetSessionConfig
22+
import com.google.adk.kt.sessions.ListEventsResponse
23+
import com.google.adk.kt.sessions.ListSessionsResponse
24+
import com.google.adk.kt.sessions.Session
25+
import com.google.adk.kt.sessions.SessionKey
26+
import com.google.adk.kt.sessions.SessionService
27+
import com.google.adk.kt.summarizer.EventSummarizer
28+
import com.google.adk.kt.summarizer.EventsCompactionConfig
29+
import com.google.adk.kt.testing.DummyAgent
30+
import com.google.adk.kt.testing.compactionEvent
31+
import com.google.adk.kt.testing.modelEventWithPromptTokens
32+
import com.google.adk.kt.testing.testSession
33+
import com.google.adk.kt.testing.userEvent
34+
import kotlin.test.Test
35+
import kotlin.test.assertEquals
36+
import kotlin.test.assertSame
37+
import kotlin.test.assertTrue
38+
import kotlinx.coroutines.test.runTest
39+
40+
class CompactionRequestProcessorTest {
41+
42+
@Test
43+
fun process_noCompactionConfig_doesNotCompact() = runTest {
44+
val sessionService = RecordingSessionService()
45+
val session = testSession()
46+
session.events.add(modelEventWithPromptTokens(500, invocationId = "inv_1", timestamp = 100L))
47+
val context =
48+
InvocationContext(
49+
session = session,
50+
agent = DummyAgent(name = "agent"),
51+
sessionService = sessionService,
52+
eventsCompactionConfig = null,
53+
)
54+
55+
val unused = CompactionRequestProcessor().process(context, LlmRequest()) {}
56+
57+
assertTrue(sessionService.appended.isEmpty())
58+
}
59+
60+
@Test
61+
fun process_slidingWindowOnlyConfig_doesNotCompact() = runTest {
62+
val sessionService = RecordingSessionService()
63+
val session = testSession()
64+
session.events.add(userEvent("u1", invocationId = "inv_1", timestamp = 100L))
65+
session.events.add(modelEventWithPromptTokens(500, invocationId = "inv_1", timestamp = 110L))
66+
val context =
67+
InvocationContext(
68+
session = session,
69+
agent = DummyAgent(name = "agent"),
70+
sessionService = sessionService,
71+
// No token-threshold fields: the token compactor must not run.
72+
eventsCompactionConfig =
73+
EventsCompactionConfig(
74+
compactionInterval = 2,
75+
overlapSize = 0,
76+
summarizer = RecordingSummarizer(),
77+
),
78+
)
79+
80+
val unused = CompactionRequestProcessor().process(context, LlmRequest()) {}
81+
82+
assertTrue(sessionService.appended.isEmpty())
83+
}
84+
85+
@Test
86+
fun process_overThreshold_appendsCompactionEvent() = runTest {
87+
val sessionService = RecordingSessionService()
88+
val summarizer = RecordingSummarizer(returning = compactionEvent(startTs = 100L, endTs = 100L))
89+
val session = testSession()
90+
session.events.add(userEvent("u1", invocationId = "inv_1", timestamp = 100L))
91+
session.events.add(modelEventWithPromptTokens(500, invocationId = "inv_2", timestamp = 200L))
92+
val context =
93+
InvocationContext(
94+
session = session,
95+
agent = DummyAgent(name = "agent"),
96+
sessionService = sessionService,
97+
eventsCompactionConfig =
98+
EventsCompactionConfig(
99+
tokenThreshold = 100,
100+
eventRetentionSize = 1,
101+
summarizer = summarizer,
102+
),
103+
)
104+
val request = LlmRequest()
105+
106+
val result = CompactionRequestProcessor().process(context, request) {}
107+
108+
// The processor only appends a compaction event to the session; it never modifies the request.
109+
assertSame(request, result)
110+
assertEquals(1, sessionService.appended.size)
111+
assertTrue(sessionService.appended.single().actions.compaction != null)
112+
assertTrue(session.events.any { it.actions.compaction != null })
113+
}
114+
115+
@Test
116+
fun process_belowThreshold_doesNotCompact() = runTest {
117+
val sessionService = RecordingSessionService()
118+
val session = testSession()
119+
session.events.add(userEvent("u1", invocationId = "inv_1", timestamp = 100L))
120+
// Most recent prompt reported only 50 tokens; below the threshold of 100.
121+
session.events.add(modelEventWithPromptTokens(50, invocationId = "inv_2", timestamp = 200L))
122+
val context =
123+
InvocationContext(
124+
session = session,
125+
agent = DummyAgent(name = "agent"),
126+
sessionService = sessionService,
127+
eventsCompactionConfig =
128+
EventsCompactionConfig(
129+
tokenThreshold = 100,
130+
eventRetentionSize = 1,
131+
summarizer = RecordingSummarizer(),
132+
),
133+
)
134+
135+
val unused = CompactionRequestProcessor().process(context, LlmRequest()) {}
136+
137+
assertTrue(sessionService.appended.isEmpty())
138+
}
139+
140+
// ----- helpers -----
141+
142+
private class RecordingSummarizer(private val returning: Event? = null) : EventSummarizer {
143+
val calls: MutableList<List<Event>> = mutableListOf()
144+
145+
override suspend fun summarizeEvents(events: List<Event>): Event? {
146+
calls.add(events.toList())
147+
return returning
148+
}
149+
}
150+
151+
private class RecordingSessionService : SessionService {
152+
val appended: MutableList<Event> = mutableListOf()
153+
154+
override suspend fun appendEvent(session: Session, event: Event): Event {
155+
appended.add(event)
156+
return super.appendEvent(session, event)
157+
}
158+
159+
override suspend fun createSession(key: SessionKey, state: Map<String, Any>?): Session =
160+
error("not used")
161+
162+
override suspend fun getSession(key: SessionKey, config: GetSessionConfig?): Session? =
163+
error("not used")
164+
165+
override suspend fun listSessions(appName: String, userId: String): ListSessionsResponse =
166+
error("not used")
167+
168+
override suspend fun deleteSession(key: SessionKey) = error("not used")
169+
170+
override suspend fun listEvents(key: SessionKey): ListEventsResponse = error("not used")
171+
}
172+
}

0 commit comments

Comments
 (0)