refresh if cache has changed.

Author: Eric Liu

app/rag_system.py

@@ -15,6 +15,10 @@
 
 
 class RAGSystem:
+    # Cache file paths
+    DOC_EMBEDDINGS_PATH = "./data/doc_embeddings.npy"
+    DOC_ABOUT_EMBEDDINGS_PATH = "./data/doc_about_embeddings.npy"
+
     def __init__(self, knowledge_base_path="./data/knowledge_base.json"):
         self.knowledge_base_path = knowledge_base_path
 
@@ -24,13 +28,22 @@ def __init__(self, knowledge_base_path="./data/knowledge_base.json"):
         # load existing embeddings if available
         logging.info("Embedding knowledge base...")
 
-        if os.path.exists("./data/doc_about_embeddings.npy") and os.path.exists(
-            "./data/doc_embeddings.npy"
+        if os.path.exists(self.DOC_ABOUT_EMBEDDINGS_PATH) and os.path.exists(
+            self.DOC_EMBEDDINGS_PATH
         ):
-            self.doc_about_embeddings = np.load("./data/doc_about_embeddings.npy")
+            self.doc_about_embeddings = np.load(self.DOC_ABOUT_EMBEDDINGS_PATH)
             logging.info("Loaded existing about document about embeddings from disk.")
-            self.doc_embeddings = np.load("./data/doc_embeddings.npy")
+            self.doc_embeddings = np.load(self.DOC_EMBEDDINGS_PATH)
             logging.info("Loaded existing document embeddings from disk.")
+
+            # Save file timestamps when loading cache
+            self.doc_embeddings_timestamp = os.path.getmtime(self.DOC_EMBEDDINGS_PATH)
+            self.doc_about_embeddings_timestamp = os.path.getmtime(
+                self.DOC_ABOUT_EMBEDDINGS_PATH
+            )
+            logging.info(
+                f"Cache loaded - doc_embeddings timestamp: {self.doc_embeddings_timestamp}, doc_about_embeddings timestamp: {self.doc_about_embeddings_timestamp}"
+            )
         else:
             self.rebuild_embeddings()
 
@@ -49,16 +62,22 @@ def rebuild_embeddings(self):
 
         # Atomic saves with guaranteed order
         self._atomic_save_numpy(
-            "./data/doc_embeddings.npy", new_doc_embeddings.cpu().numpy()
+            self.DOC_EMBEDDINGS_PATH, new_doc_embeddings.cpu().numpy()
         )
         self._atomic_save_numpy(
-            "./data/doc_about_embeddings.npy", new_about_embeddings.cpu().numpy()
+            self.DOC_ABOUT_EMBEDDINGS_PATH, new_about_embeddings.cpu().numpy()
         )
 
         # Update in-memory embeddings only after successful saves
         self.doc_embeddings = new_doc_embeddings
         self.doc_about_embeddings = new_about_embeddings
 
+        # Update file timestamps after successful saves
+        self.doc_embeddings_timestamp = os.path.getmtime(self.DOC_EMBEDDINGS_PATH)
+        self.doc_about_embeddings_timestamp = os.path.getmtime(
+            self.DOC_ABOUT_EMBEDDINGS_PATH
+        )
+
         logging.info("Embeddings rebuilt successfully.")
 
     def load_knowledge_base(self):
@@ -117,6 +136,43 @@ def compute_document_scores(
 
         return result
 
+    def cache_check(func):
+        """Decorator to automatically check cache consistency"""
+
+        def wrapper(self, *args, **kwargs):
+            try:
+                current_times = [
+                    os.path.getmtime(self.DOC_EMBEDDINGS_PATH),
+                    os.path.getmtime(self.DOC_ABOUT_EMBEDDINGS_PATH),
+                ]
+                stored_times = [
+                    self.doc_embeddings_timestamp,
+                    self.doc_about_embeddings_timestamp,
+                ]
+
+                # update cache if timestamps are different from out last load
+                if current_times != stored_times:
+                    self._reload_cache()
+
+            except (OSError, FileNotFoundError, PermissionError):
+                logging.warning("Cache files inaccessible, rebuilding...")
+                self.rebuild_embeddings()
+
+            return func(self, *args, **kwargs)
+
+        return wrapper
+
+    def _reload_cache(self):
+        self.doc_embeddings = np.load(self.DOC_EMBEDDINGS_PATH)
+        self.doc_about_embeddings = np.load(self.DOC_ABOUT_EMBEDDINGS_PATH)
+
+        # update our timestamps of the cached files
+        self.doc_embeddings_timestamp = os.path.getmtime(self.DOC_EMBEDDINGS_PATH)
+        self.doc_about_embeddings_timestamp = os.path.getmtime(
+            self.DOC_ABOUT_EMBEDDINGS_PATH
+        )
+
+    @cache_check
     def retrieve(
         self, query, similarity_threshold=0.4, high_match_threshold=0.8, max_docs=5
     ):

app/test_rag_system.py

@@ -1,6 +1,7 @@
 import unittest
 
 from rag_system import RAGSystem
+import os
 
 
 class TestRAGSystem(unittest.TestCase):
@@ -123,6 +124,79 @@ def test_compute_document_scores(self):
 
         print("Test for compute_document_scores passed successfully!")
 
+    def test_cache_check_reload_cache(self):
+        # Simulate cache file timestamp change to trigger _reload_cache
+        original_doc_embeddings_timestamp = self.rag_system.doc_embeddings_timestamp
+        original_doc_about_embeddings_timestamp = (
+            self.rag_system.doc_about_embeddings_timestamp
+        )
+
+        # Patch os.path.getmtime to return different timestamps
+        def fake_getmtime(path):
+            if path == self.rag_system.DOC_EMBEDDINGS_PATH:
+                return original_doc_embeddings_timestamp + 1
+            if path == self.rag_system.DOC_ABOUT_EMBEDDINGS_PATH:
+                return original_doc_about_embeddings_timestamp + 1
+            return 0
+
+        self.rag_system._reload_cache_called = False
+
+        def fake_reload_cache():
+            self.rag_system._reload_cache_called = True
+            real_reload_cache()
+
+        real_getmtime = os.path.getmtime
+        os.path.getmtime = fake_getmtime
+
+        # Patch _reload_cache to set a flag
+        real_reload_cache = self.rag_system._reload_cache
+        self.rag_system._reload_cache = fake_reload_cache
+
+        # Call a cache_check-decorated method
+        self.rag_system.retrieve("test query")
+
+        self.assertTrue(
+            self.rag_system._reload_cache_called,
+            "Cache reload was not triggered when timestamps changed.",
+        )
+
+        # Restore patched methods
+        os.path.getmtime = real_getmtime
+        self.rag_system._reload_cache = real_reload_cache
+        print("Test for cache_check reload_cache passed successfully!")
+
+    def test_cache_check_rebuild_embeddings_on_error(self):
+        # Patch os.path.getmtime to raise OSError
+        real_getmtime = os.path.getmtime
+
+        def raise_oserror(path):
+            raise OSError("Simulated error")
+
+        os.path.getmtime = raise_oserror
+
+        self.rag_system._rebuild_embeddings_called = False
+
+        def fake_rebuild_embeddings():
+            self.rag_system._rebuild_embeddings_called = True
+            return real_rebuild_embeddings()
+
+        self.rag_system.rebuild_embeddings = fake_rebuild_embeddings
+        # Patch rebuild_embeddings to set a flag
+        real_rebuild_embeddings = self.rag_system.rebuild_embeddings
+
+        # Call a cache_check-decorated method
+        self.rag_system.retrieve("test query")
+
+        self.assertTrue(
+            self.rag_system._rebuild_embeddings_called,
+            "rebuild_embeddings was not triggered on cache access error.",
+        )
+
+        # Restore patched methods
+        os.path.getmtime = real_getmtime
+        self.rag_system.rebuild_embeddings = real_rebuild_embeddings
+        print("Test for cache_check rebuild_embeddings on error passed successfully!")
+
 
 if __name__ == "__main__":
     unittest.main()