Skip to content

Commit 968a9a8

Browse files
google-genai-botcopybara-github
authored andcommitted
feat: Add ContextCacheConfig to InvocationContext
This change introduces ContextCacheConfig to the InvocationContext and Runner classes, allowing context cache settings to be configured per invocation. PiperOrigin-RevId: 868315679
1 parent 8acb1ea commit 968a9a8

2 files changed

Lines changed: 35 additions & 2 deletions

File tree

core/src/main/java/com/google/adk/agents/InvocationContext.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ public class InvocationContext {
5656
private final Map<String, Boolean> endOfAgents;
5757
private final ResumabilityConfig resumabilityConfig;
5858
@Nullable private final EventsCompactionConfig eventsCompactionConfig;
59+
@Nullable private final ContextCacheConfig contextCacheConfig;
5960
private final InvocationCostManager invocationCostManager;
6061
private final Map<String, Object> callbackContextData;
6162

@@ -81,6 +82,7 @@ protected InvocationContext(Builder builder) {
8182
this.endOfAgents = builder.endOfAgents;
8283
this.resumabilityConfig = builder.resumabilityConfig;
8384
this.eventsCompactionConfig = builder.eventsCompactionConfig;
85+
this.contextCacheConfig = builder.contextCacheConfig;
8486
this.invocationCostManager = builder.invocationCostManager;
8587
this.callbackContextData = builder.callbackContextData;
8688
}
@@ -400,6 +402,11 @@ public Optional<EventsCompactionConfig> eventsCompactionConfig() {
400402
return Optional.ofNullable(eventsCompactionConfig);
401403
}
402404

405+
/** Returns the context cache configuration for the current agent run. */
406+
public Optional<ContextCacheConfig> contextCacheConfig() {
407+
return Optional.ofNullable(contextCacheConfig);
408+
}
409+
403410
/** Returns whether to pause the invocation right after this [event]. */
404411
public boolean shouldPauseInvocation(Event event) {
405412
if (!isResumable()) {
@@ -472,6 +479,7 @@ private Builder(InvocationContext context) {
472479
this.endOfAgents = new ConcurrentHashMap<>(context.endOfAgents);
473480
this.resumabilityConfig = context.resumabilityConfig;
474481
this.eventsCompactionConfig = context.eventsCompactionConfig;
482+
this.contextCacheConfig = context.contextCacheConfig;
475483
this.invocationCostManager = context.invocationCostManager;
476484
this.callbackContextData = context.callbackContextData;
477485
}
@@ -493,6 +501,7 @@ private Builder(InvocationContext context) {
493501
private Map<String, Boolean> endOfAgents = new ConcurrentHashMap<>();
494502
private ResumabilityConfig resumabilityConfig = new ResumabilityConfig();
495503
@Nullable private EventsCompactionConfig eventsCompactionConfig;
504+
@Nullable private ContextCacheConfig contextCacheConfig;
496505
private InvocationCostManager invocationCostManager = new InvocationCostManager();
497506
private Map<String, Object> callbackContextData = new ConcurrentHashMap<>();
498507

@@ -730,6 +739,18 @@ public Builder eventsCompactionConfig(@Nullable EventsCompactionConfig eventsCom
730739
return this;
731740
}
732741

742+
/**
743+
* Sets the context cache configuration for the current agent run.
744+
*
745+
* @param contextCacheConfig the context cache configuration.
746+
* @return this builder instance for chaining.
747+
*/
748+
@CanIgnoreReturnValue
749+
public Builder contextCacheConfig(@Nullable ContextCacheConfig contextCacheConfig) {
750+
this.contextCacheConfig = contextCacheConfig;
751+
return this;
752+
}
753+
733754
/**
734755
* Sets the callback context data for the invocation.
735756
*
@@ -778,6 +799,7 @@ public boolean equals(Object o) {
778799
&& Objects.equals(endOfAgents, that.endOfAgents)
779800
&& Objects.equals(resumabilityConfig, that.resumabilityConfig)
780801
&& Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig)
802+
&& Objects.equals(contextCacheConfig, that.contextCacheConfig)
781803
&& Objects.equals(invocationCostManager, that.invocationCostManager)
782804
&& Objects.equals(callbackContextData, that.callbackContextData);
783805
}
@@ -802,6 +824,7 @@ public int hashCode() {
802824
endOfAgents,
803825
resumabilityConfig,
804826
eventsCompactionConfig,
827+
contextCacheConfig,
805828
invocationCostManager,
806829
callbackContextData);
807830
}

core/src/main/java/com/google/adk/runner/Runner.java

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919
import com.google.adk.agents.ActiveStreamingTool;
2020
import com.google.adk.agents.BaseAgent;
21+
import com.google.adk.agents.ContextCacheConfig;
2122
import com.google.adk.agents.InvocationContext;
2223
import com.google.adk.agents.LiveRequestQueue;
2324
import com.google.adk.agents.LlmAgent;
@@ -75,6 +76,7 @@ public class Runner {
7576
private final PluginManager pluginManager;
7677
private final ResumabilityConfig resumabilityConfig;
7778
@Nullable private final EventsCompactionConfig eventsCompactionConfig;
79+
@Nullable private final ContextCacheConfig contextCacheConfig;
7880

7981
/** Builder for {@link Runner}. */
8082
public static class Builder {
@@ -138,6 +140,7 @@ public Runner build() {
138140
List<? extends Plugin> buildPlugins;
139141
ResumabilityConfig buildResumabilityConfig;
140142
EventsCompactionConfig buildEventsCompactionConfig;
143+
ContextCacheConfig buildContextCacheConfig;
141144

142145
if (this.app != null) {
143146
if (this.agent != null) {
@@ -154,12 +157,14 @@ public Runner build() {
154157
? this.app.resumabilityConfig()
155158
: new ResumabilityConfig();
156159
buildEventsCompactionConfig = this.app.eventsCompactionConfig();
160+
buildContextCacheConfig = this.app.contextCacheConfig();
157161
} else {
158162
buildAgent = this.agent;
159163
buildAppName = this.appName;
160164
buildPlugins = this.plugins;
161165
buildResumabilityConfig = new ResumabilityConfig();
162166
buildEventsCompactionConfig = null;
167+
buildContextCacheConfig = null;
163168
}
164169

165170
if (buildAgent == null) {
@@ -182,7 +187,8 @@ public Runner build() {
182187
memoryService,
183188
buildPlugins,
184189
buildResumabilityConfig,
185-
buildEventsCompactionConfig);
190+
buildEventsCompactionConfig,
191+
buildContextCacheConfig);
186192
}
187193
}
188194

@@ -257,6 +263,7 @@ public Runner(
257263
memoryService,
258264
plugins,
259265
resumabilityConfig,
266+
null,
260267
null);
261268
}
262269

@@ -274,7 +281,8 @@ protected Runner(
274281
@Nullable BaseMemoryService memoryService,
275282
List<? extends Plugin> plugins,
276283
ResumabilityConfig resumabilityConfig,
277-
@Nullable EventsCompactionConfig eventsCompactionConfig) {
284+
@Nullable EventsCompactionConfig eventsCompactionConfig,
285+
@Nullable ContextCacheConfig contextCacheConfig) {
278286
this.agent = agent;
279287
this.appName = appName;
280288
this.artifactService = artifactService;
@@ -283,6 +291,7 @@ protected Runner(
283291
this.pluginManager = new PluginManager(plugins);
284292
this.resumabilityConfig = resumabilityConfig;
285293
this.eventsCompactionConfig = createEventsCompactionConfig(agent, eventsCompactionConfig);
294+
this.contextCacheConfig = contextCacheConfig;
286295
}
287296

288297
/**
@@ -644,6 +653,7 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session) {
644653
.session(session)
645654
.resumabilityConfig(this.resumabilityConfig)
646655
.eventsCompactionConfig(this.eventsCompactionConfig)
656+
.contextCacheConfig(this.contextCacheConfig)
647657
.agent(this.findAgentToRun(session, rootAgent));
648658
}
649659

0 commit comments

Comments
 (0)