Implement Cloudflare KV as cache backend (#5)

Author: dlackty

README.md

@@ -21,10 +21,12 @@ Clone the repo and install dependencies.
 You will need to sign up for two services (which both have free tiers):
 
 - [Cloudflare](https://www.cloudflare.com): Where our worker will be hosted.
-- [Upstash](https://upstash.com): We use Upstash's redis-over-HTTP service for storing cached OpenAI responses. Depending on your usage, you may try replacing Redis with [Cloudflare KV](https://developers.cloudflare.com/workers/runtime-apis/kv/) instead which is eventually consistent but will likely provide better read latency.
+- [Upstash](https://upstash.com): We use Upstash's redis-over-HTTP service for storing cached OpenAI responses.
 
 Finally, set up your redis secrets based on instructions in `wrangler.toml`.
 
+Depending on your usage, you may try replacing Redis with [Cloudflare KV](https://developers.cloudflare.com/workers/runtime-apis/kv/) instead which is eventually consistent but will likely provide better read latency. Check `wrangler.toml` for setup instructions.
+
 ### Usage
 
 Start the proxy server at http://localhost:8787 with:

src/cache.ts

@@ -80,17 +80,29 @@ interface CachedResponse {
 }
 
 export class ResponseCache {
-  redis: Redis;
+  redis: Redis | undefined;
+  kv: KVNamespace | undefined;
 
   constructor({ env }: { env: Env }) {
-    this.redis = new Redis({
-      url: env.UPSTASH_REDIS_REST_URL,
-      token: env.UPSTASH_REDIS_REST_TOKEN,
-    });
+    if (env.UPSTASH_REDIS_REST_URL && env.UPSTASH_REDIS_REST_TOKEN) {
+      this.redis = new Redis({
+        url: env.UPSTASH_REDIS_REST_URL,
+        token: env.UPSTASH_REDIS_REST_TOKEN,
+      });
+    } else {
+      this.kv = env.OPENAI_CACHE;
+    }
   }
 
   read = async ({ cacheKey }: { cacheKey: string }): Promise<Response | null> => {
-    const result = await this.redis.get(cacheKey);
+    let result;
+    if (this.redis) {
+      result = await this.redis.get(cacheKey);
+    } else if (this.kv) {
+      result = await this.kv.get(cacheKey);
+    } else {
+      return null;
+    }
     if (result) {
       // Note: Upstash seems to automatically parse it to JSON:
       const cachedResponse =
@@ -118,7 +130,12 @@ export class ResponseCache {
       headers: getHeadersAsObject(response.headers),
       body: body ? body : null,
     };
-    const options = ttl != null ? { ex: ttl } : {};
-    await this.redis.set(cacheKey, JSON.stringify(responseObject), options);
+    if (this.redis) {
+      const options = ttl != null ? { ex: ttl } : {};
+      await this.redis.set(cacheKey, JSON.stringify(responseObject), options);
+    } else if (this.kv) {
+      const options = ttl != null ? { expirationTtl: ttl } : {};
+      await this.kv.put(cacheKey, JSON.stringify(responseObject), options);
+    }
   };
 }

src/env.ts

@@ -10,4 +10,5 @@ export interface Env {
 
   UPSTASH_REDIS_REST_URL: string;
   UPSTASH_REDIS_REST_TOKEN: string;
+  OPENAI_CACHE: KVNamespace;
 }

wrangler.toml

@@ -2,16 +2,22 @@ name = "openai-caching-proxy-worker"
 main = "src/index.ts"
 compatibility_date = "2022-12-11"
 
-# https://developers.cloudflare.com/workers/platform/environment-variables/#comparing-secrets-and-environment-variables
-[vars]
+# This cloudflare worker uses either Upstash or workers KV to store cache.
 
-# This cloudflare worker uses Upstash (hosted redis-over-REST):
+# To use Upstash (hosted redis-over-REST),
 # https://docs.upstash.com/redis/quickstarts/cloudflareworkers
 # 
-# After signing up for a free account at upstash.com, set your secrets below based on instructions
+# 1. signing up for a free account at upstash.com, set your secrets below based on instructions
 # here for `wrangler secret put` (or alternatively use the Cloudflare Workers web UI):
 # https://developers.cloudflare.com/workers/platform/environment-variables/#adding-secrets-via-wrangler
-
-# The necessary secrets are:
+# 2. Uncomment following environment variables config:
+# [vars]
 # - UPSTASH_REDIS_REST_TOKEN
 # - UPSTASH_REDIS_REST_URL
+
+# To use workers KV,
+# 1. Run `wrangler kv:namespace create <YOUR_NAMESPACE>` in your terminal
+# 2. Uncomment following config and fill in <YOUR_ID>:
+# kv_namespaces = [
+#    { binding = "OPENAI_CACHE", id = "<YOUR_ID>" }
+# ]