try to avoid changing current cli logic flow

memgpt/persistence_manager.py

@@ -27,10 +27,75 @@ class PersistenceManager(ABC):
     def update_memory(self, new_memory):
         pass
 
-
 class InMemoryStateManager(PersistenceManager):
     """In-memory state manager has nothing to manage, all agents are held in-memory"""
 
+    recall_memory_cls = DummyRecallMemory
+    archival_memory_cls = DummyArchivalMemory
+
+    def __init__(self):
+        # Memory held in-state useful for debugging stateful versions
+        self.memory = None
+        self.messages = []
+        self.all_messages = []
+
+    @staticmethod
+    def load(filename):
+        with open(filename, 'rb') as f:
+            return pickle.load(f)
+
+    def save(self, filename):
+        with open(filename, 'wb') as fh:
+            pickle.dump(self, fh, protocol=pickle.HIGHEST_PROTOCOL)
+
+    def init(self, agent):
+        printd(f"Initializing InMemoryStateManager with agent object")
+        self.all_messages = [{'timestamp': get_local_time(), 'message': msg} for msg in agent.messages.copy()]
+        self.messages = [{'timestamp': get_local_time(), 'message': msg} for msg in agent.messages.copy()]
+        self.memory = agent.memory
+        printd(f"InMemoryStateManager.all_messages.len = {len(self.all_messages)}")
+        printd(f"InMemoryStateManager.messages.len = {len(self.messages)}")
+
+        # Persistence manager also handles DB-related state
+        self.recall_memory = self.recall_memory_cls(message_database=self.all_messages)
+        self.archival_memory_db = []
+        self.archival_memory = self.archival_memory_cls(archival_memory_database=self.archival_memory_db)
+
+    def trim_messages(self, num):
+        # printd(f"InMemoryStateManager.trim_messages")
+        self.messages = [self.messages[0]] + self.messages[num:]
+
+    def prepend_to_messages(self, added_messages):
+        # first tag with timestamps
+        added_messages = [{'timestamp': get_local_time(), 'message': msg} for msg in added_messages]
+
+        printd(f"InMemoryStateManager.prepend_to_message")
+        self.messages = [self.messages[0]] + added_messages + self.messages[1:]
+        self.all_messages.extend(added_messages)
+
+    def append_to_messages(self, added_messages):
+        # first tag with timestamps
+        added_messages = [{'timestamp': get_local_time(), 'message': msg} for msg in added_messages]
+
+        printd(f"InMemoryStateManager.append_to_messages")
+        self.messages = self.messages + added_messages
+        self.all_messages.extend(added_messages)
+
+    def swap_system_message(self, new_system_message):
+        # first tag with timestamps
+        new_system_message = {'timestamp': get_local_time(), 'message': new_system_message}
+
+        printd(f"InMemoryStateManager.swap_system_message")
+        self.messages[0] = new_system_message
+        self.all_messages.append(new_system_message)
+
+    def update_memory(self, new_memory):
+        printd(f"InMemoryStateManager.update_memory")
+        self.memory = new_memory
+
+class LocalStateManager(PersistenceManager):
+    """In-memory state manager has nothing to manage, all agents are held in-memory"""
+
     recall_memory_cls = DummyRecallMemory
     archival_memory_cls = LocalArchivalMemory 
 

tests/test_load_archival.py

@@ -10,7 +10,8 @@ import memgpt.constants as constants
 import memgpt.personas.personas as personas
 import memgpt.humans.humans as humans
 from memgpt.persistence_manager import (
-    InMemoryStateManager
+    InMemoryStateManager,
+    LocalStateManager
 )
 from memgpt.config import Config
 from memgpt.constants import MEMGPT_DIR, DEFAULT_MEMGPT_MODEL
@@ -39,7 +40,7 @@ def test_archival():
     )
 
     # create state manager based off loaded data
-    persistence_manager = InMemoryStateManager(archival_memory_db="tmp_hf_dataset")
+    persistence_manager = LocalStateManager(archival_memory_db="tmp_hf_dataset")
 
     # create agent
     memgpt_agent = presets.use_preset(