feat(community): Add TTL support to UpstashRedisCache (#7422)

langchain-ai · Dec 24, 2024 · 1200ec9 · 1200ec9
1 parent da80dce
commit 1200ec9
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/examples/src/cache/chat_models/upstash_redis.ts b/examples/src/cache/chat_models/upstash_redis.ts
@@ -7,6 +7,7 @@ const cache = new UpstashRedisCache({
     url: "UPSTASH_REDIS_REST_URL",
     token: "UPSTASH_REDIS_REST_TOKEN",
   },
+  ttl: 3600,
 });
 
 const model = new ChatOpenAI({ cache });
diff --git a/examples/src/cache/upstash_redis.ts b/examples/src/cache/upstash_redis.ts
@@ -7,6 +7,7 @@ const cache = new UpstashRedisCache({
     url: "UPSTASH_REDIS_REST_URL",
     token: "UPSTASH_REDIS_REST_TOKEN",
   },
+  ttl: 3600,
 });
 
 const model = new OpenAI({ cache });
diff --git a/libs/langchain-community/src/caches/upstash_redis.ts b/libs/langchain-community/src/caches/upstash_redis.ts
@@ -18,6 +18,10 @@ export type UpstashRedisCacheProps = {
    * An existing Upstash Redis client.
    */
   client?: Redis;
+  /**
+   * Time-to-live (TTL) for cached items in seconds.
+   */
+  ttl?: number;
 };
 
 /**
@@ -30,6 +34,7 @@ export type UpstashRedisCacheProps = {
  *     url: "UPSTASH_REDIS_REST_URL",
  *     token: "UPSTASH_REDIS_REST_TOKEN",
  *   },
+ *   ttl: 3600, // Optional: Cache entries will expire after 1 hour
  * });
  * // Initialize the OpenAI model with Upstash Redis cache for caching responses
  * const model = new ChatOpenAI({
@@ -42,9 +47,12 @@ export type UpstashRedisCacheProps = {
 export class UpstashRedisCache extends BaseCache {
   private redisClient: Redis;
 
+  private ttl?: number;
+
   constructor(props: UpstashRedisCacheProps) {
     super();
-    const { config, client } = props;
+    const { config, client, ttl } = props;
+    this.ttl = ttl;
 
     if (client) {
       this.redisClient = client;
@@ -84,10 +92,13 @@ export class UpstashRedisCache extends BaseCache {
   public async update(prompt: string, llmKey: string, value: Generation[]) {
     for (let i = 0; i < value.length; i += 1) {
       const key = getCacheKey(prompt, llmKey, String(i));
-      await this.redisClient.set(
-        key,
-        JSON.stringify(serializeGeneration(value[i]))
-      );
+      const serializedValue = JSON.stringify(serializeGeneration(value[i]));
+
+      if (this.ttl) {
+        await this.redisClient.set(key, serializedValue, { ex: this.ttl });
+      } else {
+        await this.redisClient.set(key, serializedValue);
+      }
     }
   }
 }