Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion core/src/commonMain/kotlin/com/google/adk/kt/models/Gemini.kt
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ class Gemini(
// existing cache (dropping the cached prefix) and returns the metadata to attach to responses.
val cacheManager =
preparedRequest.cacheConfig?.let {
GeminiContextCacheManager(name, GenaiCacheClient(client.caches))
GeminiContextCacheManager(name, GenaiCacheClient(client.caches), cacheScopeOf(client))
}
val cacheResult = cacheManager?.handleContextCaching(preparedRequest)
val finalRequest = cacheResult?.request ?: preparedRequest
Expand Down Expand Up @@ -232,6 +232,20 @@ class Gemini(
}
}

/**
* The backend namespace that owns explicit cache resources, folded into the cache fingerprint so a
* cache is never reused across backends or projects. Includes the backend type and, for the
* enterprise backend, the project and location. The server base URL is not included because the SDK
* [Client] does not expose it.
*/
internal fun cacheScopeOf(client: Client): Map<String, String> = buildMap {
this["backend"] = if (client.enterprise) "vertex" else "gemini"
if (client.enterprise) {
client.project?.let { this["project"] = it }
client.location?.let { this["location"] = it }
}
}

/**
* Prepares an [LlmRequest] for the GenerateContent API.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,14 @@ import kotlinx.serialization.json.encodeToJsonElement
*
* @param modelName The model resource name used when creating caches.
* @param cacheClient The client used to create and delete Gemini caches.
* @param cacheScope The backend namespace that owns cache resources (e.g. backend type and, for the
* enterprise backend, project and location). Folded into the fingerprint so a cache is never
* reused across backends or projects.
*/
internal class GeminiContextCacheManager(
private val modelName: String,
private val cacheClient: CacheClient,
private val cacheScope: Map<String, String> = emptyMap(),
) {

/**
Expand Down Expand Up @@ -174,17 +178,19 @@ internal class GeminiContextCacheManager(
}

/**
* Generates a 16-character fingerprint over the cacheable state: model, system instruction,
* tools, tool config, and the first [cacheContentsCount] contents. Serializes them into a single
* canonical JSON object with every object's keys recursively sorted, so the fingerprint depends
* only on content, not on map insertion or property order, and has no field-boundary ambiguity.
* The model is included because explicit caches are model-specific.
* Generates a 16-character fingerprint over the cacheable state: model, cache scope, system
* instruction, tools, tool config, and the first [cacheContentsCount] contents. Serializes them
* into a single canonical JSON object with every object's keys recursively sorted, so the
* fingerprint depends only on content, not on map insertion or property order, and has no
* field-boundary ambiguity. The model and cache scope are included because explicit caches are
* model- and backend-specific.
*/
@OptIn(FrameworkInternalApi::class)
private fun generateFingerprint(request: LlmRequest, cacheContentsCount: Int): String {
val fields =
buildMap<String, JsonElement> {
this["model"] = JsonPrimitive(modelName)
this["cache_scope"] = adkJson.encodeToJsonElement(cacheScope)
request.config.systemInstruction?.let {
this["system_instruction"] = adkJson.encodeToJsonElement(it)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,37 @@ class GeminiContextCacheManagerTest {
assertNull(result.cacheMetadata?.cacheName)
}

@Test
fun handleContextCaching_differentBackend_doesNotReuseCache() = runTest {
val request = baseRequest()
// A cache fingerprinted under the Gemini backend must not validate under Vertex, because the
// cache scope (backend/project/location) is part of the fingerprint.
val geminiManager =
GeminiContextCacheManager("gemini-2.0-flash", FakeCacheClient(), mapOf("backend" to "gemini"))
val activeMetadata =
CacheMetadata(
fingerprint = fingerprintFor(geminiManager, request),
contentsCount = 1,
cacheName = "cache/gemini",
expireTime = Clock.System.now().toEpochMilliseconds() + 100_000,
invocationsUsed = 1,
)

val vertexFake = FakeCacheClient()
val vertexManager =
GeminiContextCacheManager(
"gemini-2.0-flash",
vertexFake,
mapOf("backend" to "vertex", "project" to "p", "location" to "us-central1"),
)
val result = vertexManager.handleContextCaching(request.copy(cacheMetadata = activeMetadata))

assertNull(result.request.config.cachedContent)
assertEquals(1, vertexFake.deleteCount)
assertEquals("cache/gemini", vertexFake.lastDeletedName)
assertNull(result.cacheMetadata?.cacheName)
}

@Test
fun handleContextCaching_validCache_stripsToolConfig() = runTest {
val fake = FakeCacheClient()
Expand Down
Loading