diff --git a/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java b/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java index 49e73401f..a4d34c385 100644 --- a/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java +++ b/core/src/main/java/com/google/adk/agents/ContextCacheConfig.java @@ -15,34 +15,51 @@ */ package com.google.adk.agents; +import static com.google.common.base.Preconditions.checkArgument; + import com.google.errorprone.annotations.InlineMe; import com.google.genai.types.HttpOptions; import java.time.Duration; import org.jspecify.annotations.Nullable; /** - * Configuration for context caching across all agents in an app. - * - *

This configuration enables and controls context caching behavior for all LLM agents in an app. - * When this config is present on an app, context caching is enabled for all agents. When absent - * (null), context caching is disabled. - * - *

Context caching can significantly reduce costs and improve response times by reusing - * previously processed context across multiple requests. + * Configures context caching for all agents in an app. Without it, nothing is cached. * - * @param cacheIntervals Maximum number of invocations to reuse the same cache before refreshing it. - * Defaults to 10. - * @param ttl Time-to-live for cache. Defaults to 1800 seconds (30 minutes). - * @param minTokens Minimum estimated request tokens required to enable caching. This compares - * against the estimated total tokens of the request (system instruction + tools + contents). - * Context cache storage may have cost. Set higher to avoid caching small requests where - * overhead may exceed benefits. Defaults to 0. - * @param createHttpOptions HTTP options, such as a timeout, for the call that creates a cache; null - * uses the client's defaults. Defaults to null. + * @param cacheIntervals maximum number of invocations that reuse the same cache before it is + * refreshed; must be from 1 to 100 + * @param ttl time-to-live of the cache; must be a whole number of seconds, at least one + * @param minTokens minimum prompt token count of the agent's previous request needed to create a + * cache; raise it to skip small requests, where cache storage can cost more than it saves; must + * not be negative + * @param createHttpOptions HTTP options, such as a timeout, for the call that creates a cache, or + * {@code null} to use the client's defaults; if that call fails, the request is sent without a + * cache */ public record ContextCacheConfig( int cacheIntervals, Duration ttl, int minTokens, @Nullable HttpOptions createHttpOptions) { + /** + * Creates a config. + * + * @throws IllegalArgumentException if a value is out of range + */ + public ContextCacheConfig { + checkArgument( + cacheIntervals >= 1 && cacheIntervals <= 100, + "cacheIntervals must be between 1 and 100, but was %s.", + cacheIntervals); + checkArgument( + ttl.getSeconds() >= 1 && ttl.getNano() == 0, + "ttl must be a whole number of seconds, at least one, but was %s.", + ttl); + checkArgument(minTokens >= 0, "minTokens must not be negative, but was %s.", minTokens); + } + + /** + * Creates a config with the defaults: {@code cacheIntervals} of 10, {@code ttl} of 30 minutes, + * {@code minTokens} of 0, and the client's default HTTP options ({@code createHttpOptions} is + * {@code null}). + */ public ContextCacheConfig() { this(10, Duration.ofMinutes(30), 0); } @@ -63,7 +80,7 @@ public int maxInvocations() { return cacheIntervals(); } - /** Returns TTL as string format for cache creation. */ + /** Returns the TTL in the {@code "s"} form that the Gemini API uses. */ public String getTtlString() { return ttl.toSeconds() + "s"; } diff --git a/core/src/main/java/com/google/adk/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index 78063683f..63ba27bde 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -278,11 +278,26 @@ public Gemini build() { @Override public Flowable generateContent(LlmRequest llmRequest, boolean stream) { - llmRequest = + LlmRequest preparedRequest = GeminiUtil.prepareGenenerateContentRequest( llmRequest, !apiClient.vertexAI(), /* stripThoughts= */ false); + String effectiveModelName = preparedRequest.model().orElse(model()); + if (preparedRequest.cacheConfig().isEmpty()) { + return sendRequest(effectiveModelName, preparedRequest, stream, /* cacheMetadata= */ null); + } + return new GeminiContextCacheManager(apiClient, effectiveModelName) + .handleContextCaching(preparedRequest) + .flatMapPublisher( + cached -> sendRequest(effectiveModelName, cached.request(), stream, cached.metadata())); + } + + /** Sends the request, and puts {@code cacheMetadata}, if any, on the final response. */ + private Flowable sendRequest( + String effectiveModelName, + LlmRequest llmRequest, + boolean stream, + @Nullable CacheMetadata cacheMetadata) { GenerateContentConfig config = withoutAutomaticContinuation(llmRequest.config().orElse(null)); - String effectiveModelName = llmRequest.model().orElse(model()); logger.trace("Request Contents: {}", llmRequest.contents()); logger.trace("Request Config: {}", config); @@ -295,18 +310,29 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre return Flowable.defer( () -> streamResumingPauses( - effectiveModelName, new GeminiContinuation(contents, config), streamFuture)); + effectiveModelName, + new GeminiContinuation(contents, config), + streamFuture, + cacheMetadata)); } else { logger.debug("Sending generateContent request to model {}", effectiveModelName); CompletableFuture responseFuture = apiClient.async.models.generateContent(effectiveModelName, contents, config); return Flowable.defer( - () -> - generateResumingPauses( - effectiveModelName, new GeminiContinuation(contents, config), responseFuture)); + () -> + generateResumingPauses( + effectiveModelName, new GeminiContinuation(contents, config), responseFuture)) + .map(response -> withCacheMetadata(response, cacheMetadata)); } } + private static LlmResponse withCacheMetadata( + LlmResponse response, @Nullable CacheMetadata cacheMetadata) { + return cacheMetadata == null + ? response + : response.toBuilder().cacheMetadata(cacheMetadata).build(); + } + @VisibleForTesting static Flowable processRawResponses(Flowable rawResponses) { return Flowable.defer(() -> new StreamingResponseAggregator().process(rawResponses)); @@ -387,7 +413,8 @@ private Optional completeOrResume( private Flowable streamResumingPauses( String modelName, GeminiContinuation continuation, - CompletableFuture> firstStream) { + CompletableFuture> firstStream, + @Nullable CacheMetadata cacheMetadata) { Deque>> streams = new ArrayDeque<>(); streams.add(firstStream); StreamingResponseAggregator aggregator = new StreamingResponseAggregator(); @@ -396,7 +423,11 @@ private Flowable streamResumingPauses( .concatMap(aggregator::processRawResponse) .concatWith( Flowable.defer( - () -> aggregator.processFinalResponse().map(continuation::withSummedUsage))); + () -> + aggregator + .processFinalResponse() + .map(continuation::withSummedUsage) + .map(response -> withCacheMetadata(response, cacheMetadata)))); } /** diff --git a/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java b/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java index 855cdcd8d..f83040a75 100644 --- a/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java +++ b/core/src/test/java/com/google/adk/agents/ContextCacheConfigTest.java @@ -16,6 +16,7 @@ package com.google.adk.agents; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.common.collect.ImmutableMap; import com.google.genai.types.HttpOptions; @@ -41,7 +42,8 @@ public void defaultConstructor_usesDefaults() { @Test public void threeArgConstructor_leavesHttpOptionsUnset() { - ContextCacheConfig config = new ContextCacheConfig(5, TTL, 100); + ContextCacheConfig config = + new ContextCacheConfig(/* cacheIntervals= */ 5, TTL, /* minTokens= */ 100); assertThat(config.cacheIntervals()).isEqualTo(5); assertThat(config.createHttpOptions()).isNull(); @@ -51,14 +53,104 @@ public void threeArgConstructor_leavesHttpOptionsUnset() { public void constructor_keepsHttpOptions() { HttpOptions httpOptions = HttpOptions.builder().timeout(10_000).build(); - assertThat(new ContextCacheConfig(5, TTL, 100, httpOptions).createHttpOptions()) + assertThat( + new ContextCacheConfig(/* cacheIntervals= */ 5, TTL, /* minTokens= */ 100, httpOptions) + .createHttpOptions()) .isEqualTo(httpOptions); } @Test @SuppressWarnings({"deprecation", "InlineMeInliner"}) // Covers the deprecated alias. public void maxInvocations_returnsCacheIntervals() { - assertThat(new ContextCacheConfig(7, TTL, 0).maxInvocations()).isEqualTo(7); + assertThat( + new ContextCacheConfig(/* cacheIntervals= */ 7, TTL, /* minTokens= */ 0) + .maxInvocations()) + .isEqualTo(7); + } + + @Test + public void constructor_cacheIntervalsAtBounds_isAccepted() { + assertThat( + new ContextCacheConfig(/* cacheIntervals= */ 1, TTL, /* minTokens= */ 0) + .cacheIntervals()) + .isEqualTo(1); + assertThat( + new ContextCacheConfig(/* cacheIntervals= */ 100, TTL, /* minTokens= */ 0) + .cacheIntervals()) + .isEqualTo(100); + } + + @Test + public void constructor_cacheIntervalsBelowOne_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(/* cacheIntervals= */ 0, TTL, /* minTokens= */ 0)); + } + + @Test + public void constructor_cacheIntervalsAboveHundred_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(/* cacheIntervals= */ 101, TTL, /* minTokens= */ 0)); + } + + @Test + public void constructor_zeroTtl_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(/* cacheIntervals= */ 10, Duration.ZERO, /* minTokens= */ 0)); + } + + @Test + public void constructor_negativeTtl_throws() { + assertThrows( + IllegalArgumentException.class, + () -> + new ContextCacheConfig( + /* cacheIntervals= */ 10, Duration.ofSeconds(-1), /* minTokens= */ 0)); + } + + @Test + public void constructor_oneSecondTtl_isAccepted() { + assertThat( + new ContextCacheConfig( + /* cacheIntervals= */ 10, Duration.ofSeconds(1), /* minTokens= */ 0) + .ttl()) + .isEqualTo(Duration.ofSeconds(1)); + } + + @Test + public void constructor_ttlUnderOneSecond_throws() { + assertThrows( + IllegalArgumentException.class, + () -> + new ContextCacheConfig( + /* cacheIntervals= */ 10, Duration.ofMillis(500), /* minTokens= */ 0)); + } + + @Test + public void constructor_fractionalTtl_throws() { + assertThrows( + IllegalArgumentException.class, + () -> + new ContextCacheConfig( + /* cacheIntervals= */ 10, Duration.ofMillis(1500), /* minTokens= */ 0)); + } + + @Test + public void constructor_negativeMinTokens_throws() { + assertThrows( + IllegalArgumentException.class, + () -> new ContextCacheConfig(/* cacheIntervals= */ 10, TTL, /* minTokens= */ -1)); + } + + @Test + public void getTtlString_returnsSeconds() { + assertThat( + new ContextCacheConfig( + /* cacheIntervals= */ 10, Duration.ofMinutes(5), /* minTokens= */ 0) + .getTtlString()) + .isEqualTo("300s"); } @Test @@ -74,7 +166,9 @@ public void toString_hidesHttpOptionValues() { HttpOptions httpOptions = HttpOptions.builder().headers(ImmutableMap.of("x-goog-api-key", "secret")).build(); - assertThat(new ContextCacheConfig(5, TTL, 100, httpOptions).toString()) + assertThat( + new ContextCacheConfig(/* cacheIntervals= */ 5, TTL, /* minTokens= */ 100, httpOptions) + .toString()) .isEqualTo( "ContextCacheConfig(cacheIntervals=5, ttl=1800s, minTokens=100," + " createHttpOptions=set)"); diff --git a/core/src/test/java/com/google/adk/models/GeminiContextCachingTest.java b/core/src/test/java/com/google/adk/models/GeminiContextCachingTest.java new file mode 100644 index 000000000..3a4bd9416 --- /dev/null +++ b/core/src/test/java/com/google/adk/models/GeminiContextCachingTest.java @@ -0,0 +1,225 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.google.adk.models; + +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Streams.stream; +import static com.google.common.truth.Truth.assertThat; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.adk.agents.ContextCacheConfig; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.apps.App; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.events.Event; +import com.google.adk.models.FakeGeminiApi.Kind; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.InMemorySessionService; +import com.google.common.collect.ImmutableList; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import java.time.Duration; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Runs agents on {@link Gemini} with context caching, against a fake Gemini API. */ +@RunWith(JUnit4.class) +public final class GeminiContextCachingTest { + + private static final String MODEL = "gemini-test-model"; + private static final String INSTRUCTION = "Answer every question briefly and politely."; + private static final ContextCacheConfig CACHE_CONFIG = + new ContextCacheConfig(/* cacheIntervals= */ 10, Duration.ofMinutes(30), /* minTokens= */ 0); + private static final String USER_ID = "user"; + private static final String CACHE_NAME = "cachedContents/greeting"; + private static final String FIRST_ANSWER = "Hi, how can I help?"; + private static final RunConfig STREAMING = + RunConfig.builder().streamingMode(RunConfig.StreamingMode.SSE).build(); + + private final FakeGeminiApi api = new FakeGeminiApi(); + + @Test + public void runAsync_secondTurn_sendsCacheNameAndOnlyTheNewMessage() { + api.queueAnswers(FIRST_ANSWER); + api.queueCacheNames(CACHE_NAME); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + run(runner, sessionId, "Hello"); + + run(runner, sessionId, "Tell me more"); + + // The cache holds the first turn ("Hello" and the first answer), so the second call names the + // cache and sends only the new message. + assertThat(texts(api.bodies(Kind.CREATE_CACHE).get(0))) + .containsExactly("Hello", FIRST_ANSWER) + .inOrder(); + JsonNode secondCall = lastGenerateBody(); + assertThat(secondCall.path("cachedContent").asText()).isEqualTo(CACHE_NAME); + assertThat(texts(secondCall)).containsExactly("Tell me more"); + } + + @Test + public void runAsync_thirdTurn_reusesCacheWithoutCreatingAnother() { + api.queueCacheNames(CACHE_NAME); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + run(runner, sessionId, "Hello"); + run(runner, sessionId, "Tell me more"); + + run(runner, sessionId, "And then?"); + + assertThat(api.bodies(Kind.CREATE_CACHE)).hasSize(1); + assertThat(lastGenerateBody().path("cachedContent").asText()).isEqualTo(CACHE_NAME); + } + + @Test + public void runAsync_streaming_putsCacheMetadataOnlyOnTheFinalEvent() { + api.queueCacheNames(CACHE_NAME); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + run(runner, sessionId, "Hello", STREAMING); + + ImmutableList events = run(runner, sessionId, "Tell me more", STREAMING); + + assertThat(partialEvents(events)).isNotEmpty(); + assertThat(eventsWithCacheMetadata(events)).containsExactly(getLast(events)); + } + + @Test + public void runAsync_turnThatCreatesCache_staysOnTheCallerThread() { + api.delayCacheCalls(Duration.ofMillis(300)); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + ImmutableList firstTurn = + eventThreads(runner, sessionId, "Hello", RunConfig.builder().build()); + + ImmutableList secondTurn = + eventThreads(runner, sessionId, "Tell me more", RunConfig.builder().build()); + + assertThat(api.kinds()).contains(Kind.CREATE_CACHE); + assertThat(firstTurn).isNotEmpty(); + assertThat(secondTurn).isEqualTo(firstTurn); + } + + @Test + public void runAsync_streamingTurnThatCreatesCache_staysOnTheCallerThread() { + api.delayCacheCalls(Duration.ofMillis(300)); + Runner runner = runner(CACHE_CONFIG); + String sessionId = newSession(runner); + ImmutableList firstTurn = eventThreads(runner, sessionId, "Hello", STREAMING); + + ImmutableList secondTurn = eventThreads(runner, sessionId, "Tell me more", STREAMING); + + assertThat(api.kinds()).contains(Kind.CREATE_CACHE); + assertThat(firstTurn).isNotEmpty(); + assertThat(secondTurn).isEqualTo(firstTurn); + } + + @Test + public void runAsync_withoutCacheConfig_sendsWholeHistoryWithoutCache() { + api.queueAnswers(FIRST_ANSWER); + Runner runner = runner(/* cacheConfig= */ null); + String sessionId = newSession(runner); + run(runner, sessionId, "Hello"); + + run(runner, sessionId, "Tell me more"); + + assertThat(api.kinds()).doesNotContain(Kind.CREATE_CACHE); + JsonNode secondCall = lastGenerateBody(); + assertThat(secondCall.has("cachedContent")).isFalse(); + assertThat(texts(secondCall)).containsExactly("Hello", FIRST_ANSWER, "Tell me more").inOrder(); + } + + private Runner runner(@Nullable ContextCacheConfig cacheConfig) { + LlmAgent agent = + LlmAgent.builder() + .name("assistant") + .model(new Gemini(MODEL, api.client())) + .instruction(INSTRUCTION) + .build(); + App.Builder app = App.builder().name("cache_app").rootAgent(agent); + if (cacheConfig != null) { + app.contextCacheConfig(cacheConfig); + } + return Runner.builder() + .app(app.build()) + .sessionService(new InMemorySessionService()) + .artifactService(new InMemoryArtifactService()) + .build(); + } + + private static String newSession(Runner runner) { + return runner.sessionService().createSession(runner.appName(), USER_ID).blockingGet().id(); + } + + /** Returns the request body of the latest model call the fake served. */ + private JsonNode lastGenerateBody() { + return getLast(api.bodies(Kind.GENERATE)); + } + + /** Returns the events that are partial (streamed chunks). */ + private static ImmutableList partialEvents(ImmutableList events) { + return events.stream() + .filter(event -> event.partial().orElse(false)) + .collect(toImmutableList()); + } + + /** Returns the events that carry cache metadata. */ + private static ImmutableList eventsWithCacheMetadata(ImmutableList events) { + return events.stream() + .filter(event -> event.cacheMetadata().isPresent()) + .collect(toImmutableList()); + } + + /** Returns the name of the thread each event of one turn was emitted on. */ + private static ImmutableList eventThreads( + Runner runner, String sessionId, String text, RunConfig runConfig) { + return ImmutableList.copyOf( + runner + .runAsync(USER_ID, sessionId, Content.fromParts(Part.fromText(text)), runConfig) + .map(event -> Thread.currentThread().getName()) + .toList() + .blockingGet()); + } + + @CanIgnoreReturnValue + private static ImmutableList run(Runner runner, String sessionId, String text) { + return run(runner, sessionId, text, RunConfig.builder().build()); + } + + @CanIgnoreReturnValue + private static ImmutableList run( + Runner runner, String sessionId, String text, RunConfig runConfig) { + return ImmutableList.copyOf( + runner + .runAsync(USER_ID, sessionId, Content.fromParts(Part.fromText(text)), runConfig) + .toList() + .blockingGet()); + } + + /** Returns the text of every part of every content in a request body, in order. */ + private static ImmutableList texts(JsonNode body) { + return stream(body.get("contents")) + .flatMap(content -> stream(content.get("parts"))) + .map(part -> part.path("text").asText()) + .collect(toImmutableList()); + } +}