Merge remote-tracking branch 'immorez/feat/prompt-lru-cache'

Author: yousshim

src/prompts/promptCache.ts

@@ -1,28 +1,38 @@
+import createLRUCache from "../utils/lru-cache";
 
 // Remove all newlines, double spaces, etc
 function normalizeText(src: string) {
-    src = src.split('\n').join(' ');
-    src = src.replace(/\s+/gm, ' ');
-    return src;
+  src = src.split("\n").join(" ");
+  src = src.replace(/\s+/gm, " ");
+  return src;
 }
 
-function extractPromptCacheKey(args: { prefix: string, suffix: string | null }) {
-    if (args.suffix) {
-        return normalizeText(args.prefix + ' ##CURSOR## ' + args.suffix);
-    } else {
-        return normalizeText(args.prefix);
-    }
+function extractPromptCacheKey(args: {
+  prefix: string;
+  suffix: string | null;
+}) {
+  if (args.suffix) {
+    return normalizeText(args.prefix + " ##CURSOR## " + args.suffix);
+  } else {
+    return normalizeText(args.prefix);
+  }
 }
 
-// TODO: make it LRU
-let cache: { [key: string]: string | null } = {};
+const promptCache = createLRUCache<string, string | null>({ maxSize: 1000 });
 
-export function getFromPromptCache(args: { prefix: string, suffix: string | null }): string | undefined | null {
-    const key = extractPromptCacheKey(args);
-    return cache[key];
+export function getFromPromptCache(args: {
+  prefix: string;
+  suffix: string | null;
+}): string | undefined | null {
+  const key = extractPromptCacheKey(args);
+  return promptCache.get(key);
 }
 
-export function setPromptToCache(args: { prefix: string, suffix: string | null, value: string | null }) {
-    const key = extractPromptCacheKey(args);
-    cache[key] = args.value;
-}
+export function setPromptToCache(args: {
+  prefix: string;
+  suffix: string | null;
+  value: string | null;
+}) {
+  const key = extractPromptCacheKey(args);
+  promptCache.set(key, args.value);
+}

src/utils/lru-cache.ts

@@ -0,0 +1,58 @@
+interface LRUCacheConfig<K, V> {
+  maxSize: number;
+}
+
+class LRUCache<K, V> {
+  private cache: Map<K, V>;
+  private maxSize: number;
+
+  constructor(config: LRUCacheConfig<K, V>) {
+    this.cache = new Map<K, V>();
+    this.maxSize = config.maxSize;
+  }
+
+  get(key: K): V | undefined {
+    if (!this.cache.has(key)) {
+      return undefined;
+    }
+
+    const value = this.cache.get(key)!;
+    this.cache.delete(key);
+    this.cache.set(key, value);
+    return value;
+  }
+
+  set(key: K, value: V): void {
+    if (this.cache.has(key)) {
+      this.cache.delete(key);
+    }
+
+    this.cache.set(key, value);
+
+    if (this.cache.size > this.maxSize) {
+      const oldestKey = this.cache.keys().next().value;
+
+      if (oldestKey) {
+        this.cache.delete(oldestKey);
+      }
+    }
+  }
+
+  clear(): void {
+    this.cache.clear();
+  }
+
+  has(key: K): boolean {
+    return this.cache.has(key);
+  }
+
+  get size(): number {
+    return this.cache.size;
+  }
+}
+
+function createLRUCache<K, V>(config: LRUCacheConfig<K, V>): LRUCache<K, V> {
+  return new LRUCache<K, V>(config);
+}
+
+export default createLRUCache;