Merge pull request #55 from OmniJax/fix/thread-safety-issue

Fix: Add thread locks to memory storage classes to prevent race conditions

Author: kkkjz

memoryos-pypi/long_term.py

@@ -2,6 +2,7 @@
 import numpy as np
 import faiss
 from collections import deque
+import threading
 try:
     from .utils import get_timestamp, get_embedding, normalize_vector, ensure_directory_exists
 except ImportError:
@@ -19,6 +20,7 @@ def __init__(self, file_path, knowledge_capacity=100, embedding_model_name: str
 
         self.embedding_model_name = embedding_model_name
         self.embedding_model_kwargs = embedding_model_kwargs if embedding_model_kwargs is not None else {}
+        self.lock = threading.Lock()
         self.load()
 
     def update_user_profile(self, user_id, new_data, merge=True):
@@ -144,8 +146,9 @@ def save(self):
             "assistant_knowledge": list(self.assistant_knowledge)
         }
         try:
-            with open(self.file_path, "w", encoding="utf-8") as f:
-                json.dump(data, f, ensure_ascii=False, indent=2)
+            with self.lock:
+                with open(self.file_path, "w", encoding="utf-8") as f:
+                    json.dump(data, f, ensure_ascii=False, indent=2)
         except IOError as e:
             print(f"Error saving LongTermMemory to {self.file_path}: {e}")
 

memoryos-pypi/mid_term.py

@@ -3,6 +3,7 @@
 from collections import defaultdict
 import faiss
 import heapq
+import threading
 from datetime import datetime
 
 try:
@@ -46,6 +47,7 @@ def __init__(self, file_path: str, client: OpenAIClient, max_capacity=2000, embe
 
         self.embedding_model_name = embedding_model_name
         self.embedding_model_kwargs = embedding_model_kwargs if embedding_model_kwargs is not None else {}
+        self.lock = threading.Lock()
         self.load()
 
     def get_page_by_id(self, page_id):
@@ -370,8 +372,9 @@ def save(self):
             # "heap_snapshot": self.heap 
         }
         try:
-            with open(self.file_path, "w", encoding="utf-8") as f:
-                json.dump(data_to_save, f, ensure_ascii=False, indent=2)
+            with self.lock:
+                with open(self.file_path, "w", encoding="utf-8") as f:
+                    json.dump(data_to_save, f, ensure_ascii=False, indent=2)
         except IOError as e:
             print(f"Error saving MidTermMemory to {self.file_path}: {e}")
 

memoryos-pypi/short_term.py

@@ -1,5 +1,6 @@
 import json
 from collections import deque
+import threading
 try:
     from .utils import get_timestamp, ensure_directory_exists
 except ImportError:
@@ -11,6 +12,7 @@ def __init__(self, file_path, max_capacity=10):
         self.file_path = file_path
         ensure_directory_exists(self.file_path)
         self.memory = deque(maxlen=max_capacity)
+        self.lock = threading.Lock()
         self.load()
 
     def add_qa_pair(self, qa_pair):
@@ -38,8 +40,9 @@ def pop_oldest(self):
 
     def save(self):
         try:
-            with open(self.file_path, "w", encoding="utf-8") as f:
-                json.dump(list(self.memory), f, ensure_ascii=False, indent=2)
+            with self.lock:
+                with open(self.file_path, "w", encoding="utf-8") as f:
+                    json.dump(list(self.memory), f, ensure_ascii=False, indent=2)
         except IOError as e:
             print(f"Error saving ShortTermMemory to {self.file_path}: {e}")