diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedTool.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedTool.kt index 2585f57d4..3b1451333 100644 --- a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedTool.kt +++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedTool.kt @@ -19,10 +19,33 @@ abstract class CachedTool( return cache(CachedToolKey(input, seed)) { onCacheMissed(input) } } + /** + * Exposes the cache as a [Map] of [Input] to [Output] filtered by instance [seed] and + * [timeCachePolicy]. Removes expired cache entries. + */ + suspend fun getCache(): Map { + val lastTimeInCache = timeInMillis() - timeCachePolicy.inWholeMilliseconds + val withoutExpired = + cache.modify { cachedToolInfo -> + // Filter entries belonging to the current seed and have not expired + val validEntries = + cachedToolInfo + .filter { (key, value) -> + if (key.seed == seed) lastTimeInCache <= value.timestamp else true + } + .toMutableMap() + // Remove expired entries for the current seed only + cachedToolInfo.keys.removeAll { key -> key.seed == seed && !validEntries.containsKey(key) } + // Modifies state A, and returns state B + Pair(cachedToolInfo, validEntries) + } + return withoutExpired.map { it.key.value to it.value.value }.toMap() + } + abstract suspend fun onCacheMissed(input: Input): Output private suspend fun cache(input: CachedToolKey, block: suspend () -> Output): Output { - val cachedToolInfo = cache.get().get(input) + val cachedToolInfo = cache.get()[input] if (cachedToolInfo != null) { val lastTimeInCache = timeInMillis() - timeCachePolicy.inWholeMilliseconds if (lastTimeInCache > cachedToolInfo.timestamp) { @@ -32,7 +55,7 @@ abstract class CachedTool( } } val response = block() - cache.get().put(input, CachedToolValue(response, timeInMillis())) + cache.get()[input] = CachedToolValue(response, timeInMillis()) return response } }