feat : server response cache

server/sources/cache.py

@@ -0,0 +1,36 @@
+import os
+import json
+from pathlib import Path
+
+class Cache:
+    def __init__(self, cache_dir='.cache', cache_file='messages.json'):
+        self.cache_dir = Path(cache_dir)
+        self.cache_file = self.cache_dir / cache_file
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        if not self.cache_file.exists():
+            with open(self.cache_file, 'w') as f:
+                json.dump([], f)
+
+        with open(self.cache_file, 'r') as f:
+            self.cache = set(json.load(f))
+
+    def add_message_pair(self, user_message: str, assistant_message: str):
+        """Add a user/assistant pair to the cache if not present."""
+        if not any(entry["user"] == user_message for entry in self.cache):
+            self.cache.append({"user": user_message, "assistant": assistant_message})
+            self._save()
+
+    def is_cached(self, user_message: str) -> bool:
+        """Check if a user msg is cached."""
+        return any(entry["user"] == user_message for entry in self.cache)
+
+    def get_cached_response(self, user_message: str) -> str | None:
+        """Return the assistant response to a user message if cached."""
+        for entry in self.cache:
+            if entry["user"] == user_message:
+                return entry["assistant"]
+        return None
+
+    def _save(self):
+        with open(self.cache_file, 'w') as f:
+            json.dump(self.cache, f, indent=2)

server/sources/generator.py

@@ -2,6 +2,7 @@
 import threading
 import logging
 from abc import abstractmethod
+from .cache import Cache
 
 class GenerationState:
     def __init__(self):
@@ -29,6 +30,7 @@ class GeneratorLLM():
         handler.setFormatter(formatter)
         self.logger.addHandler(handler)
         self.logger.setLevel(logging.INFO)
+        cache = Cache()
     
     def set_model(self, model: str) -> None:
         self.logger.info(f"Model set to {model}")

server/sources/ollama_handler.py

@@ -13,6 +13,10 @@ class OllamaLLM(GeneratorLLM):
 
     def generate(self, history):
         self.logger.info(f"Using {self.model} for generation with Ollama")
+        if cache.is_cached(history[-1]['content']):
+            self.state.current_buffer = cache.get_cached_response(history[-1]['content'])
+            self.state.is_generating = False
+            return
         try:
             with self.state.lock:
                 self.state.is_generating = True
@@ -43,6 +47,7 @@ class OllamaLLM(GeneratorLLM):
             self.logger.info("Generation complete")
             with self.state.lock:
                 self.state.is_generating = False
+                self.cache.add_message_pair(history[-1]['content'], self.state.current_buffer)
 
 if __name__ == "__main__":
     generator = OllamaLLM()