diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index ed9b2106..3997eb9e 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -55,6 +55,7 @@ public class InvocationContext { private final Map endOfAgents; private final ResumabilityConfig resumabilityConfig; @Nullable private final EventsCompactionConfig eventsCompactionConfig; + @Nullable private final ContextCacheConfig contextCacheConfig; private final InvocationCostManager invocationCostManager; private Optional branch; @@ -79,6 +80,7 @@ protected InvocationContext(Builder builder) { this.endOfAgents = builder.endOfAgents; this.resumabilityConfig = builder.resumabilityConfig; this.eventsCompactionConfig = builder.eventsCompactionConfig; + this.contextCacheConfig = builder.contextCacheConfig; this.invocationCostManager = builder.invocationCostManager; } @@ -364,6 +366,11 @@ public Optional eventsCompactionConfig() { return Optional.ofNullable(eventsCompactionConfig); } + /** Returns the context cache configuration for the current agent run. */ + public Optional contextCacheConfig() { + return Optional.ofNullable(contextCacheConfig); + } + /** Returns whether to pause the invocation right after this [event]. */ public boolean shouldPauseInvocation(Event event) { if (!isResumable()) { @@ -436,6 +443,7 @@ private Builder(InvocationContext context) { this.endOfAgents = new ConcurrentHashMap<>(context.endOfAgents); this.resumabilityConfig = context.resumabilityConfig; this.eventsCompactionConfig = context.eventsCompactionConfig; + this.contextCacheConfig = context.contextCacheConfig; this.invocationCostManager = context.invocationCostManager; } @@ -456,6 +464,7 @@ private Builder(InvocationContext context) { private Map endOfAgents = new ConcurrentHashMap<>(); private ResumabilityConfig resumabilityConfig = new ResumabilityConfig(); @Nullable private EventsCompactionConfig eventsCompactionConfig; + @Nullable private ContextCacheConfig contextCacheConfig; private InvocationCostManager invocationCostManager = new InvocationCostManager(); /** @@ -692,6 +701,18 @@ public Builder eventsCompactionConfig(@Nullable EventsCompactionConfig eventsCom return this; } + /** + * Sets the context cache configuration for the current agent run. + * + * @param contextCacheConfig the context cache configuration. + * @return this builder instance for chaining. + */ + @CanIgnoreReturnValue + public Builder contextCacheConfig(@Nullable ContextCacheConfig contextCacheConfig) { + this.contextCacheConfig = contextCacheConfig; + return this; + } + /** * Builds the {@link InvocationContext} instance. * @@ -728,6 +749,7 @@ public boolean equals(Object o) { && Objects.equals(endOfAgents, that.endOfAgents) && Objects.equals(resumabilityConfig, that.resumabilityConfig) && Objects.equals(eventsCompactionConfig, that.eventsCompactionConfig) + && Objects.equals(contextCacheConfig, that.contextCacheConfig) && Objects.equals(invocationCostManager, that.invocationCostManager); } @@ -751,6 +773,7 @@ public int hashCode() { endOfAgents, resumabilityConfig, eventsCompactionConfig, + contextCacheConfig, invocationCostManager); } } diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index e543f7d6..17f37803 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -18,6 +18,7 @@ import com.google.adk.agents.ActiveStreamingTool; import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.ContextCacheConfig; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LiveRequestQueue; import com.google.adk.agents.LlmAgent; @@ -75,6 +76,7 @@ public class Runner { private final PluginManager pluginManager; private final ResumabilityConfig resumabilityConfig; @Nullable private final EventsCompactionConfig eventsCompactionConfig; + @Nullable private final ContextCacheConfig contextCacheConfig; /** Builder for {@link Runner}. */ public static class Builder { @@ -138,6 +140,7 @@ public Runner build() { List buildPlugins; ResumabilityConfig buildResumabilityConfig; EventsCompactionConfig buildEventsCompactionConfig; + ContextCacheConfig buildContextCacheConfig; if (this.app != null) { if (this.agent != null) { @@ -154,12 +157,14 @@ public Runner build() { ? this.app.resumabilityConfig() : new ResumabilityConfig(); buildEventsCompactionConfig = this.app.eventsCompactionConfig(); + buildContextCacheConfig = this.app.contextCacheConfig(); } else { buildAgent = this.agent; buildAppName = this.appName; buildPlugins = this.plugins; buildResumabilityConfig = new ResumabilityConfig(); buildEventsCompactionConfig = null; + buildContextCacheConfig = null; } if (buildAgent == null) { @@ -182,7 +187,8 @@ public Runner build() { memoryService, buildPlugins, buildResumabilityConfig, - buildEventsCompactionConfig); + buildEventsCompactionConfig, + buildContextCacheConfig); } } @@ -257,6 +263,7 @@ public Runner( memoryService, plugins, resumabilityConfig, + null, null); } @@ -274,7 +281,8 @@ protected Runner( @Nullable BaseMemoryService memoryService, List plugins, ResumabilityConfig resumabilityConfig, - @Nullable EventsCompactionConfig eventsCompactionConfig) { + @Nullable EventsCompactionConfig eventsCompactionConfig, + @Nullable ContextCacheConfig contextCacheConfig) { this.agent = agent; this.appName = appName; this.artifactService = artifactService; @@ -283,6 +291,7 @@ protected Runner( this.pluginManager = new PluginManager(plugins); this.resumabilityConfig = resumabilityConfig; this.eventsCompactionConfig = createEventsCompactionConfig(agent, eventsCompactionConfig); + this.contextCacheConfig = contextCacheConfig; } /** @@ -636,6 +645,7 @@ private InvocationContext.Builder newInvocationContextBuilder(Session session) { .session(session) .resumabilityConfig(this.resumabilityConfig) .eventsCompactionConfig(this.eventsCompactionConfig) + .contextCacheConfig(this.contextCacheConfig) .agent(this.findAgentToRun(session, rootAgent)); }