From 9e854cdb3957a2bba4422129f65dc5dc0bceb472 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Javier=20P=C3=A9rez=20Pacheco?= <javi.pacheco@gmail.com>
Date: Tue, 11 Jun 2024 13:42:11 +0200
Subject: [PATCH] Cached Tool with seed (#759)

---
 .../functional/xef/llm/assistants/CachedTool.kt   | 15 ++++++++++-----
 .../xef/llm/assistants/CachedToolInfo.kt          |  3 ---
 2 files changed, 10 insertions(+), 8 deletions(-)
 delete mode 100644 core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedToolInfo.kt

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 583bf0aa6..2585f57d4 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
@@ -5,29 +5,34 @@ import arrow.fx.coroutines.timeInMillis
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.days
 
+data class CachedToolKey<K>(val value: K, val seed: String)
+
+data class CachedToolValue<V>(val value: V, val timestamp: Long)
+
 abstract class CachedTool<Input, Output>(
-  private val cache: Atomic<MutableMap<Input, CachedToolInfo<Output>>>,
+  private val cache: Atomic<MutableMap<CachedToolKey<Input>, CachedToolValue<Output>>>,
+  private val seed: String,
   private val timeCachePolicy: Duration = 1.days
 ) : Tool<Input, Output> {
 
   override suspend fun invoke(input: Input): Output {
-    return cache(input) { onCacheMissed(input) }
+    return cache(CachedToolKey(input, seed)) { onCacheMissed(input) }
   }
 
   abstract suspend fun onCacheMissed(input: Input): Output
 
-  private suspend fun cache(input: Input, block: suspend () -> Output): Output {
+  private suspend fun cache(input: CachedToolKey<Input>, block: suspend () -> Output): Output {
     val cachedToolInfo = cache.get().get(input)
     if (cachedToolInfo != null) {
       val lastTimeInCache = timeInMillis() - timeCachePolicy.inWholeMilliseconds
       if (lastTimeInCache > cachedToolInfo.timestamp) {
         cache.get().remove(input)
       } else {
-        return cachedToolInfo.response
+        return cachedToolInfo.value
       }
     }
     val response = block()
-    cache.get().put(input, CachedToolInfo(response, timeInMillis()))
+    cache.get().put(input, CachedToolValue(response, timeInMillis()))
     return response
   }
 }
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedToolInfo.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedToolInfo.kt
deleted file mode 100644
index 84dd49763..000000000
--- a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/CachedToolInfo.kt
+++ /dev/null
@@ -1,3 +0,0 @@
-package com.xebia.functional.xef.llm.assistants
-
-data class CachedToolInfo<Response>(val response: Response, val timestamp: Long)