Merge pull request #106 from DefangLabs/eric/make-rebuild-run-in-background

Author: nullfunc

app/app.py

@@ -15,6 +15,7 @@
 import os
 import segment.analytics as analytics
 import uuid
+import threading
 
 import logging
 import redis
@@ -117,25 +118,15 @@ def v1_ask():
         return jsonify({"error": "Invalid or missing Ask Token"}), 401
 
 
-@app.route("/trigger-rebuild", methods=["POST"])
-@csrf.exempt
-def trigger_rebuild():
-    token = request.args.get("token")
-    if token != os.getenv("REBUILD_TOKEN"):
-        return jsonify({"error": "Unauthorized"}), 401
+def run_rebuild():
     try:
         print("Running get_knowledge_base.py script...")
         result = subprocess.run(
             ["python3", "get_knowledge_base.py"], capture_output=True, text=True
         )
         if result.returncode != 0:
             print(f"Error running get_knowledge_base.py script: {result.stderr}")
-            return jsonify(
-                {
-                    "error": "Error running get_knowledge_base.py script",
-                    "details": result.stderr,
-                }
-            ), 500
+            return
 
         print("Finished running get_knowledge_base.py script.")
 
@@ -146,12 +137,7 @@ def trigger_rebuild():
         )
         if result.returncode != 0:
             print(f"Error running get_samples_examples.py script: {result.stderr}")
-            return jsonify(
-                {
-                    "error": "Error running get_samples_examples.py script",
-                    "details": result.stderr,
-                }
-            ), 500
+            return
 
         print("Finished running get_samples_examples.py script.")
 
@@ -160,14 +146,28 @@ def trigger_rebuild():
             app.rag_system.rebuild()
         except Exception as e:
             logging.error(f"Error rebuilding embeddings: {str(e)}")
-            return jsonify({"error": "Error rebuilding embeddings"}), 500
+            return
 
         logging.info("Finished rebuilding embeddings.")
-        return jsonify({"status": "Rebuild triggered successfully"}), 200
 
     except Exception as e:
-        print(f"Error in /trigger-rebuild endpoint: {e}")
-        return jsonify({"error": "Internal Server Error"}), 500
+        logging.error(f"Error in rebuild process: {e}")
+
+
+@app.route("/trigger-rebuild", methods=["POST"])
+@csrf.exempt
+def trigger_rebuild():
+    token = request.args.get("token")
+    if token != os.getenv("REBUILD_TOKEN"):
+        return jsonify({"error": "Unauthorized"}), 401
+
+    # Start the rebuild in a background thread
+    thread = threading.Thread(target=run_rebuild)
+    thread.daemon = True  # Dies when main process dies
+    thread.start()
+
+    # Return immediately
+    return jsonify({"status": "Rebuild started successfully"}), 202
 
 
 @app.route("/data/<path:name>")

app/rag_system.py

@@ -8,12 +8,17 @@
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 import traceback
+from atomicwrites import atomic_write
 
 openai.api_base = os.getenv("OPENAI_BASE_URL")
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 
 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
 
@@ -23,29 +28,58 @@ 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()
 
         logging.info("Knowledge base embeddings created")
         self.conversation_history = []
 
+    def _atomic_save_numpy(self, file_path, data):
+        with atomic_write(file_path, mode="wb", overwrite=True) as f:
+            np.save(f, data)
+
     def rebuild_embeddings(self):
-        logging.info("No existing document embeddings found, creating new embeddings.")
-        self.doc_embeddings = self.embed_knowledge_base()
-        self.doc_about_embeddings = self.embed_knowledge_base_about()
-        # cache doc_embeddings to disk
-        np.save("./data/doc_embeddings.npy", self.doc_embeddings.cpu().numpy())
-        np.save(
-            "./data/doc_about_embeddings.npy", self.doc_about_embeddings.cpu().numpy()
+        logging.info("Rebuilding document embeddings...")
+
+        new_doc_embeddings = self.embed_knowledge_base()
+        new_about_embeddings = self.embed_knowledge_base_about()
+
+        # Atomic saves with guaranteed order
+        self._atomic_save_numpy(
+            self.DOC_EMBEDDINGS_PATH, new_doc_embeddings.cpu().numpy()
+        )
+        self._atomic_save_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):
         with open(self.knowledge_base_path, "r") as kb_file:
             return json.load(kb_file)
@@ -102,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/requirements.txt

@@ -13,6 +13,7 @@ PyYAML==6.0.2
 GitPython==3.1.44
 redis==6.2.0
 fakeredis==2.30.1
+atomicwrites==1.4.1
 
 # linter
 ruff>=0.12.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()