finish the first verion code of orm, but it still has some problems w…

…hen running scheduler
MemTensor · fridayL · Sep 24, 2025 · Sep 10, 2025 · Sep 10, 2025 · Sep 11, 2025
commit 9deb36865e94913fe83e51384f0fd9a69c1a5162
diff --git a/.gitignore b/.gitignore
@@ -48,6 +48,7 @@ share/python-wheels/
 .installed.cfg
 *.egg
 MANIFEST
+.run
 
 # PyInstaller
 #  Usually these files are written by a python script from a template

diff --git a/.vscode/settings.json b/.vscode/settings.json
diff --git a/examples/mem_scheduler/memos_w_scheduler.py b/examples/mem_scheduler/memos_w_scheduler.py
@@ -15,7 +15,7 @@
 
 
 if TYPE_CHECKING:
-    from memos.mem_scheduler.schemas import (
+    from memos.mem_scheduler.schemas.message_schemas import (
         ScheduleLogForWebItem,
     )
 

diff --git a/src/memos/configs/mem_scheduler.py b/src/memos/configs/mem_scheduler.py
@@ -64,6 +64,19 @@ class GeneralSchedulerConfig(BaseSchedulerConfig):
         default=20, description="Capacity of the activation memory monitor"
     )
 
+    # Database configuration for ORM persistence
+    db_path: str | None = Field(
+        default=None,
+        description="Path to SQLite database file for ORM persistence. If None, uses default scheduler_orm.db",
+    )
+    db_url: str | None = Field(
+        default=None,
+        description="Database URL for ORM persistence (e.g., mysql://user:pass@host/db). Takes precedence over db_path",
+    )
+    enable_orm_persistence: bool = Field(
+        default=True, description="Whether to enable ORM-based persistence for monitors"
+    )
+
 
 class SchedulerConfigFactory(BaseConfig):
     """Factory class for creating scheduler configurations."""
@@ -111,10 +124,6 @@ class RabbitMQConfig(
         le=65535,  # Port must be <= 65535
     )
 
-    @classmethod
-    def get_env_prefix(cls) -> str:
-        return "RABBITMQ_"
-
 
 class GraphDBAuthConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
     uri: str = Field(
@@ -132,10 +141,6 @@ class GraphDBAuthConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
         default=True, description="Whether to automatically create the database if it doesn't exist"
     )
 
-    @classmethod
-    def get_env_prefix(cls) -> str:
-        return "GRAPH_DB_"
-
 
 class OpenAIConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
     api_key: str = Field(default="", description="API key for OpenAI service")

diff --git a/src/memos/mem_scheduler/base_scheduler.py b/src/memos/mem_scheduler/base_scheduler.py
@@ -5,6 +5,8 @@
 from datetime import datetime
 from pathlib import Path
 
+from sqlalchemy.engine import Engine
+
 from memos.configs.mem_scheduler import AuthConfig, BaseSchedulerConfig
 from memos.llms.base import BaseLLM
 from memos.log import get_logger
@@ -62,6 +64,7 @@ def __init__(self, config: BaseSchedulerConfig):
         )
 
         self.retriever: SchedulerRetriever | None = None
+        self.db_engine: Engine | None = None
         self.monitor: SchedulerGeneralMonitor | None = None
         self.dispatcher_monitor: SchedulerDispatcherMonitor | None = None
         self.dispatcher = SchedulerDispatcher(
@@ -70,12 +73,12 @@ def __init__(self, config: BaseSchedulerConfig):
         )
 
         # internal message queue
-        self.max_internal_messae_queue_size = 100
+        self.max_internal_message_queue_size = 100
         self.memos_message_queue: Queue[ScheduleMessageItem] = Queue(
-            maxsize=self.max_internal_messae_queue_size
+            maxsize=self.max_internal_message_queue_size
         )
         self._web_log_message_queue: Queue[ScheduleLogForWebItem] = Queue(
-            maxsize=self.max_internal_messae_queue_size
+            maxsize=self.max_internal_message_queue_size
         )
         self._consumer_thread = None  # Reference to our consumer thread
         self._running = False
@@ -92,34 +95,56 @@ def __init__(self, config: BaseSchedulerConfig):
         self.auth_config = None
         self.rabbitmq_config = None
 
-    def initialize_modules(self, chat_llm: BaseLLM, process_llm: BaseLLM | None = None):
+    def initialize_modules(
+        self,
+        chat_llm: BaseLLM,
+        process_llm: BaseLLM | None = None,
+        db_engine: Engine | None = None,
+    ):
         if process_llm is None:
             process_llm = chat_llm
 
-        # initialize submodules
-        self.chat_llm = chat_llm
-        self.process_llm = process_llm
-        self.monitor = SchedulerGeneralMonitor(process_llm=self.process_llm, config=self.config)
-        self.dispatcher_monitor = SchedulerDispatcherMonitor(config=self.config)
-        self.retriever = SchedulerRetriever(process_llm=self.process_llm, config=self.config)
+        try:
+            # initialize submodules
+            self.chat_llm = chat_llm
+            self.process_llm = process_llm
+            self.db_engine = db_engine
+            self.monitor = SchedulerGeneralMonitor(
+                process_llm=self.process_llm, config=self.config, db_engine=self.db_engine
+            )
+            self.dispatcher_monitor = SchedulerDispatcherMonitor(config=self.config)
+            self.retriever = SchedulerRetriever(process_llm=self.process_llm, config=self.config)
+
+            if self.enable_parallel_dispatch:
+                self.dispatcher_monitor.initialize(dispatcher=self.dispatcher)
+                self.dispatcher_monitor.start()
+
+            # initialize with auth_config
+            if self.auth_config_path is not None and Path(self.auth_config_path).exists():
+                self.auth_config = AuthConfig.from_local_config(config_path=self.auth_config_path)
+            elif AuthConfig.default_config_exists():
+                self.auth_config = AuthConfig.from_local_config()
+            else:
+                self.auth_config = AuthConfig.from_local_env()
 
-        if self.enable_parallel_dispatch:
-            self.dispatcher_monitor.initialize(dispatcher=self.dispatcher)
-            self.dispatcher_monitor.start()
-
-        # initialize with auth_cofig
-        if self.auth_config_path is not None and Path(self.auth_config_path).exists():
-            self.auth_config = AuthConfig.from_local_config(config_path=self.auth_config_path)
-        elif AuthConfig.default_config_exists():
-            self.auth_config = AuthConfig.from_local_config()
-        else:
-            self.auth_config = None
+            if self.auth_config is not None:
+                self.rabbitmq_config = self.auth_config.rabbitmq
+                self.initialize_rabbitmq(config=self.rabbitmq_config)
 
-        if self.auth_config is not None:
-            self.rabbitmq_config = self.auth_config.rabbitmq
-            self.initialize_rabbitmq(config=self.rabbitmq_config)
+            logger.debug("GeneralScheduler has been initialized")
+        except Exception as e:
+            logger.error(f"Failed to initialize scheduler modules: {e}", exc_info=True)
+            # Clean up any partially initialized resources
+            self._cleanup_on_init_failure()
+            raise
 
-        logger.debug("GeneralScheduler has been initialized")
+    def _cleanup_on_init_failure(self):
+        """Clean up resources if initialization fails."""
+        try:
+            if hasattr(self, "dispatcher_monitor") and self.dispatcher_monitor is not None:
+                self.dispatcher_monitor.stop()
+        except Exception as e:
+            logger.warning(f"Error during cleanup: {e}")
 
     @property
     def mem_cube(self) -> GeneralMemCube:
@@ -200,8 +225,10 @@ def replace_working_memory(
             text_mem_base: TreeTextMemory = text_mem_base
 
             # process rerank memories with llm
-            query_monitor = self.monitor.query_monitors[user_id][mem_cube_id]
-            query_history = query_monitor.get_queries_with_timesort()
+            query_db_manager = self.monitor.query_monitors[user_id][mem_cube_id]
+            # Sync with database to get latest query history
+            query_db_manager.sync_with_orm()
+            query_history = query_db_manager.obj.get_queries_with_timesort()
             memories_with_new_order, rerank_success_flag = (
                 self.retriever.process_and_rerank_memories(
                     queries=query_history,
@@ -212,7 +239,7 @@ def replace_working_memory(
             )
 
             # update working memory monitors
-            query_keywords = query_monitor.get_keywords_collections()
+            query_keywords = query_db_manager.obj.get_keywords_collections()
             logger.info(
                 f"Processing {len(memories_with_new_order)} memories with {len(query_keywords)} query keywords"
             )
@@ -235,7 +262,7 @@ def replace_working_memory(
 
             mem_monitors: list[MemoryMonitorItem] = self.monitor.working_memory_monitors[user_id][
                 mem_cube_id
-            ].get_sorted_mem_monitors(reverse=True)
+            ].obj.get_sorted_mem_monitors(reverse=True)
             new_working_memories = [mem_monitor.tree_memory_item for mem_monitor in mem_monitors]
 
             text_mem_base.replace_working_memory(memories=new_working_memories)
@@ -278,6 +305,7 @@ def update_activation_memory(
             new_text_memories = new_memories
         else:
             logger.error("Not Implemented.")
+            return
 
         try:
             if isinstance(mem_cube.act_mem, VLLMKVCacheMemory):
@@ -333,7 +361,9 @@ def update_activation_memory(
             )
 
         except Exception as e:
-            logger.warning(f"MOS-based activation memory update failed: {e}", exc_info=True)
+            logger.error(f"MOS-based activation memory update failed: {e}", exc_info=True)
+            # Re-raise the exception if it's critical for the operation
+            # For now, we'll continue execution but this should be reviewed
 
     def update_activation_memory_periodically(
         self,
@@ -358,7 +388,8 @@ def update_activation_memory_periodically(
                 if (
                     user_id not in self.monitor.working_memory_monitors
                     or mem_cube_id not in self.monitor.working_memory_monitors[user_id]
-                    or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0
+                    or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].obj.memories)
+                    == 0
                 ):
                     logger.warning(
                         "No memories found in working_memory_monitors, activation memory update is skipped"
@@ -369,9 +400,13 @@ def update_activation_memory_periodically(
                     user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
                 )
 
+                # Sync with database to get latest activation memories
+                activation_db_manager = self.monitor.activation_memory_monitors[user_id][
+                    mem_cube_id
+                ]
+                activation_db_manager.sync_with_orm()
                 new_activation_memories = [
-                    m.memory_text
-                    for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id].memories
+                    m.memory_text for m in activation_db_manager.obj.memories
                 ]
 
                 logger.info(
@@ -461,25 +496,26 @@ def _message_consumer(self) -> None:
         """
         while self._running:  # Use a running flag for graceful shutdown
             try:
-                # Check if queue has messages (non-blocking)
-                if not self.memos_message_queue.empty():
-                    # Get all available messages at once
-                    messages = []
-                    while not self.memos_message_queue.empty():
-                        try:
-                            messages.append(self.memos_message_queue.get_nowait())
-                        except queue.Empty:
-                            break
-
-                    if messages:
-                        try:
-                            self.dispatcher.dispatch(messages)
-                        except Exception as e:
-                            logger.error(f"Error dispatching messages: {e!s}")
-                        finally:
-                            # Mark all messages as processed
-                            for _ in messages:
-                                self.memos_message_queue.task_done()
+                # Get all available messages at once (thread-safe approach)
+                messages = []
+                while True:
+                    try:
+                        # Use get_nowait() directly without empty() check to avoid race conditions
+                        message = self.memos_message_queue.get_nowait()
+                        messages.append(message)
+                    except queue.Empty:
+                        # No more messages available
+                        break
+
+                if messages:
+                    try:
+                        self.dispatcher.dispatch(messages)
+                    except Exception as e:
+                        logger.error(f"Error dispatching messages: {e!s}")
+                    finally:
+                        # Mark all messages as processed
+                        for _ in messages:
+                            self.memos_message_queue.task_done()
 
                 # Sleep briefly to prevent busy waiting
                 time.sleep(self._consume_interval)  # Adjust interval as needed

diff --git a/src/memos/mem_scheduler/general_modules/misc.py b/src/memos/mem_scheduler/general_modules/misc.py
@@ -20,6 +20,8 @@
 class EnvConfigMixin(Generic[T]):
     """Abstract base class for environment variable configuration."""
 
+    ENV_PREFIX = "MEMSCHEDULER_"
+
     @classmethod
     def get_env_prefix(cls) -> str:
         """Automatically generates environment variable prefix from class name.
@@ -37,7 +39,8 @@ def get_env_prefix(cls) -> str:
         if class_name.endswith("Config"):
             class_name = class_name[:-6]
         # Convert to uppercase and add trailing underscore
-        return f"{class_name.upper()}_"
+
+        return f"{cls.ENV_PREFIX}{class_name.upper()}_"
 
     @classmethod
     def from_env(cls: type[T]) -> T:
@@ -111,6 +114,26 @@ def to_dict(self) -> dict:
             dump_data["timestamp"] = self.serialize_datetime(self.timestamp, None)
         return dump_data
 
+    def to_json(self, **kwargs) -> str:
+        """
+        Convert model instance to a JSON string.
+        - Accepts the same kwargs as json.dumps (e.g., indent, ensure_ascii)
+        - Default settings make JSON human-readable and UTF-8 safe
+        """
+        return json.dumps(self.to_dict(), ensure_ascii=False, default=lambda o: str(o), **kwargs)
+
+    @classmethod
+    def from_json(cls: type[BaseModelType], json_str: str) -> BaseModelType:
+        """
+        Create model instance from a JSON string.
+        - Parses JSON into a dictionary and delegates to from_dict
+        """
+        try:
+            data = json.loads(json_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON string: {e}") from e
+        return cls.from_dict(data)
+
     @classmethod
     def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
         """
@@ -169,3 +192,11 @@ def put(self, item: T, block: bool = False, timeout: float | None = None) -> Non
     def get_queue_content_without_pop(self) -> list[T]:
         """Return a copy of the queue's contents without modifying it."""
         return list(self.queue)
+
+    def clear(self) -> None:
+        """Remove all items from the queue.
+
+        This operation is thread-safe.
+        """
+        with self.mutex:
+            self.queue.clear()