fix indexing race condition where knowledge_base is updated by the embeddings have not yet fully calculated, any queries may result in out of bounds issue.

Author: Eric Liu

app/data/doc_about_embeddings.npy

@@ -22,7 +22,7 @@ class RAGSystem:
     def __init__(self, knowledge_base_path="./data/knowledge_base.json"):
         self.knowledge_base_path = knowledge_base_path
 
-        self.knowledge_base = self.load_knowledge_base()
+        knowledge_base = self.load_knowledge_base()
         self.model = SentenceTransformer("all-MiniLM-L6-v2")
 
         # load existing embeddings if available
@@ -35,6 +35,7 @@ def __init__(self, knowledge_base_path="./data/knowledge_base.json"):
             logging.info("Loaded existing about document about embeddings from disk.")
             self.doc_embeddings = np.load(self.DOC_EMBEDDINGS_PATH)
             logging.info("Loaded existing document embeddings from disk.")
+            self.knowledge_base = knowledge_base
 
             # Save file timestamps when loading cache
             self.doc_embeddings_timestamp = os.path.getmtime(self.DOC_EMBEDDINGS_PATH)
@@ -45,7 +46,7 @@ def __init__(self, knowledge_base_path="./data/knowledge_base.json"):
                 f"Cache loaded - doc_embeddings timestamp: {self.doc_embeddings_timestamp}, doc_about_embeddings timestamp: {self.doc_about_embeddings_timestamp}"
             )
         else:
-            self.rebuild_embeddings()
+            self.rebuild_embeddings(knowledge_base)
 
         logging.info("Knowledge base embeddings created")
         self.conversation_history = []
@@ -54,11 +55,11 @@ 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):
+    def rebuild_embeddings(self, knowledge_base):
         logging.info("Rebuilding document embeddings...")
 
-        new_doc_embeddings = self.embed_knowledge_base()
-        new_about_embeddings = self.embed_knowledge_base_about()
+        new_doc_embeddings = self.embed_knowledge_base(knowledge_base)
+        new_about_embeddings = self.embed_knowledge_base_about(knowledge_base)
 
         # Atomic saves with guaranteed order
         self._atomic_save_numpy(
@@ -69,6 +70,7 @@ def rebuild_embeddings(self):
         )
 
         # Update in-memory embeddings only after successful saves
+        self.knowledge_base = knowledge_base
         self.doc_embeddings = new_doc_embeddings
         self.doc_about_embeddings = new_about_embeddings
 
@@ -84,13 +86,13 @@ def load_knowledge_base(self):
         with open(self.knowledge_base_path, "r") as kb_file:
             return json.load(kb_file)
 
-    def embed_knowledge_base(self):
-        docs = [f"{doc['about']}. {doc['text']}" for doc in self.knowledge_base]
+    def embed_knowledge_base(self, knowledge_base):
+        docs = [f"{doc['about']}. {doc['text']}" for doc in knowledge_base]
         return self.model.encode(docs, convert_to_tensor=True)
 
-    def embed_knowledge_base_about(self):
+    def embed_knowledge_base_about(self, knowledge_base):
         return self.model.encode(
-            [doc["about"] for doc in self.knowledge_base], convert_to_tensor=True
+            [doc["about"] for doc in knowledge_base], convert_to_tensor=True
         )
 
     def normalize_query(self, query):
@@ -193,7 +195,20 @@ def compute_relevance_scores(
         self, text_similarities, about_similarities, high_match_threshold
     ):
         relevance_scores = []
-        for i, _ in enumerate(self.knowledge_base):
+
+        # Defensive check for size mismatches
+        sizes = [
+            len(text_similarities),
+            len(about_similarities),
+            len(self.knowledge_base),
+        ]
+        if len(set(sizes)) > 1:  # Not all sizes are equal
+            logging.warning(
+                f"Array size mismatch detected: text_similarities={sizes[0]}, about_similarities={sizes[1]}, knowledge_base={sizes[2]}"
+            )
+
+        max_index = min(sizes)
+        for i in range(max_index):
             about_similarity = about_similarities[i]
             text_similarity = text_similarities[i]
             # If either about or text similarity is above the high match threshold, prioritize it
@@ -321,8 +336,10 @@ def rebuild(self):
         Rebuild the embeddings for the knowledge base. This should be called whenever the knowledge base is updated.
         """
         print("Rebuilding embeddings for the knowledge base...")
-        self.knowledge_base = self.load_knowledge_base()  # Reload the knowledge base
-        self.doc_embeddings = self.rebuild_embeddings()  # Rebuild the embeddings
+        knowledge_base = self.load_knowledge_base()  # Reload the knowledge base
+        self.doc_embeddings = self.rebuild_embeddings(
+            knowledge_base
+        )  # Rebuild the embeddings
         print("Embeddings have been rebuilt.")
 
     def get_citations(self, retrieved_docs):