Scheduler: new feat about orchestrator task schedule

examples/mem_scheduler/memos_w_scheduler.py

@@ -13,15 +13,15 @@
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_os.main import MOS
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
-from memos.mem_scheduler.schemas.general_schemas import (
-    ADD_LABEL,
-    ANSWER_LABEL,
-    MEM_ARCHIVE_LABEL,
-    MEM_ORGANIZE_LABEL,
-    MEM_UPDATE_LABEL,
-    QUERY_LABEL,
-)
 from memos.mem_scheduler.schemas.message_schemas import ScheduleLogForWebItem
+from memos.mem_scheduler.schemas.task_schemas import (
+    ADD_TASK_LABEL,
+    ANSWER_TASK_LABEL,
+    MEM_ARCHIVE_TASK_LABEL,
+    MEM_ORGANIZE_TASK_LABEL,
+    MEM_UPDATE_TASK_LABEL,
+    QUERY_TASK_LABEL,
+)
 from memos.mem_scheduler.utils.filter_utils import transform_name_to_key
 
 
@@ -118,24 +118,24 @@ def _first_content() -> str:
             return memcube_content[0].get("content", "") or content
         return content
 
-    if label in ("addMessage", QUERY_LABEL, ANSWER_LABEL):
+    if label in ("addMessage", QUERY_TASK_LABEL, ANSWER_TASK_LABEL):
         target_cube = cube_display.replace("MemCube", "")
         title = _format_title(item.timestamp, f"addMessages to {target_cube} MemCube")
         return title, _truncate_with_rules(_first_content())
 
-    if label in ("addMemory", ADD_LABEL):
+    if label in ("addMemory", ADD_TASK_LABEL):
         title = _format_title(item.timestamp, f"{cube_display} added {memory_len} memories")
         return title, _truncate_with_rules(_first_content())
 
-    if label in ("updateMemory", MEM_UPDATE_LABEL):
+    if label in ("updateMemory", MEM_UPDATE_TASK_LABEL):
         title = _format_title(item.timestamp, f"{cube_display} updated {memory_len} memories")
         return title, _truncate_with_rules(_first_content())
 
-    if label in ("archiveMemory", MEM_ARCHIVE_LABEL):
+    if label in ("archiveMemory", MEM_ARCHIVE_TASK_LABEL):
         title = _format_title(item.timestamp, f"{cube_display} archived {memory_len} memories")
         return title, _truncate_with_rules(_first_content())
 
-    if label in ("mergeMemory", MEM_ORGANIZE_LABEL):
+    if label in ("mergeMemory", MEM_ORGANIZE_TASK_LABEL):
         title = _format_title(item.timestamp, f"{cube_display} merged {memory_len} memories")
         merged = [c for c in memcube_content if c.get("type") == "merged"]
         post = [c for c in memcube_content if c.get("type") == "postMerge"]

examples/mem_scheduler/redis_example.py

@@ -9,8 +9,8 @@
 from memos.configs.mem_scheduler import SchedulerConfigFactory
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_scheduler.scheduler_factory import SchedulerFactory
-from memos.mem_scheduler.schemas.general_schemas import QUERY_LABEL
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import QUERY_TASK_LABEL
 
 
 if TYPE_CHECKING:
@@ -55,7 +55,7 @@ def service_run():
         message_item = ScheduleMessageItem(
             user_id=user_id,
             mem_cube_id="mem_cube_2",
-            label=QUERY_LABEL,
+            label=QUERY_TASK_LABEL,
             mem_cube=mem_cube,
             content=query,
             timestamp=datetime.now(),

examples/mem_scheduler/try_schedule_modules.py

@@ -14,7 +14,7 @@
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_scheduler.analyzer.mos_for_test_scheduler import MOSForTestScheduler
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
-from memos.mem_scheduler.schemas.general_schemas import (
+from memos.mem_scheduler.schemas.task_schemas import (
     NOT_APPLICABLE_TYPE,
 )
 

src/memos/api/handlers/chat_handler.py

@@ -30,11 +30,11 @@
     prepare_reference_data,
     process_streaming_references_complete,
 )
-from memos.mem_scheduler.schemas.general_schemas import (
-    ANSWER_LABEL,
-    QUERY_LABEL,
-)
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import (
+    ANSWER_TASK_LABEL,
+    QUERY_TASK_LABEL,
+)
 from memos.templates.mos_prompts import (
     FURTHER_SUGGESTION_PROMPT,
     get_memos_prompt,
@@ -244,7 +244,7 @@ def generate_chat_response() -> Generator[str, None, None]:
                         user_id=chat_req.user_id,
                         mem_cube_id=scheduler_cube_id,
                         query=chat_req.query,
-                        label=QUERY_LABEL,
+                        label=QUERY_TASK_LABEL,
                     )
                     # Extract memories from search results
                     memories_list = []
@@ -423,7 +423,7 @@ def generate_chat_response() -> Generator[str, None, None]:
                         user_id=chat_req.user_id,
                         mem_cube_id=scheduler_cube_id,
                         query=chat_req.query,
-                        label=QUERY_LABEL,
+                        label=QUERY_TASK_LABEL,
                     )
                     # Extract memories from search results
                     memories_list = []
@@ -1034,7 +1034,7 @@ async def _post_chat_processing(
 
             # Send answer to scheduler
             self._send_message_to_scheduler(
-                user_id=user_id, mem_cube_id=cube_id, query=clean_response, label=ANSWER_LABEL
+                user_id=user_id, mem_cube_id=cube_id, query=clean_response, label=ANSWER_TASK_LABEL
             )
 
             self.logger.info(f"Post-chat processing completed for user {user_id}")

src/memos/configs/mem_reader.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import Any, ClassVar
 
-from pydantic import Field, field_validator, model_validator
+from pydantic import ConfigDict, Field, field_validator, model_validator
 
 from memos.configs.base import BaseConfig
 from memos.configs.chunker import ChunkerConfigFactory
@@ -44,6 +44,9 @@ def parse_datetime(cls, value):
 class SimpleStructMemReaderConfig(BaseMemReaderConfig):
     """SimpleStruct MemReader configuration class."""
 
+    # Allow passing additional fields without raising validation errors
+    model_config = ConfigDict(extra="allow", strict=True)
+
 
 class MultiModalStructMemReaderConfig(BaseMemReaderConfig):
     """MultiModalStruct MemReader configuration class."""

src/memos/mem_os/core.py

@@ -15,14 +15,14 @@
 from memos.mem_reader.factory import MemReaderFactory
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
 from memos.mem_scheduler.scheduler_factory import SchedulerFactory
-from memos.mem_scheduler.schemas.general_schemas import (
-    ADD_LABEL,
-    ANSWER_LABEL,
-    MEM_READ_LABEL,
-    PREF_ADD_LABEL,
-    QUERY_LABEL,
-)
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import (
+    ADD_TASK_LABEL,
+    ANSWER_TASK_LABEL,
+    MEM_READ_TASK_LABEL,
+    PREF_ADD_TASK_LABEL,
+    QUERY_TASK_LABEL,
+)
 from memos.mem_user.user_manager import UserManager, UserRole
 from memos.memories.activation.item import ActivationMemoryItem
 from memos.memories.parametric.item import ParametricMemoryItem
@@ -283,7 +283,7 @@ def chat(self, query: str, user_id: str | None = None, base_prompt: str | None =
                     message_item = ScheduleMessageItem(
                         user_id=target_user_id,
                         mem_cube_id=mem_cube_id,
-                        label=QUERY_LABEL,
+                        label=QUERY_TASK_LABEL,
                         content=query,
                         timestamp=datetime.utcnow(),
                     )
@@ -343,7 +343,7 @@ def chat(self, query: str, user_id: str | None = None, base_prompt: str | None =
                 message_item = ScheduleMessageItem(
                     user_id=target_user_id,
                     mem_cube_id=mem_cube_id,
-                    label=ANSWER_LABEL,
+                    label=ANSWER_TASK_LABEL,
                     content=response,
                     timestamp=datetime.utcnow(),
                 )
@@ -771,7 +771,7 @@ def process_textual_memory():
                             message_item = ScheduleMessageItem(
                                 user_id=target_user_id,
                                 mem_cube_id=mem_cube_id,
-                                label=MEM_READ_LABEL,
+                                label=MEM_READ_TASK_LABEL,
                                 content=json.dumps(mem_ids),
                                 timestamp=datetime.utcnow(),
                                 task_id=task_id,
@@ -783,7 +783,7 @@ def process_textual_memory():
                             message_item = ScheduleMessageItem(
                                 user_id=target_user_id,
                                 mem_cube_id=mem_cube_id,
-                                label=ADD_LABEL,
+                                label=ADD_TASK_LABEL,
                                 content=json.dumps(mem_ids),
                                 timestamp=datetime.utcnow(),
                                 task_id=task_id,
@@ -824,7 +824,7 @@ def process_preference_memory():
                         user_id=target_user_id,
                         session_id=target_session_id,
                         mem_cube_id=mem_cube_id,
-                        label=PREF_ADD_LABEL,
+                        label=PREF_ADD_TASK_LABEL,
                         content=json.dumps(messages_list),
                         timestamp=datetime.utcnow(),
                     )
@@ -878,7 +878,7 @@ def process_preference_memory():
                         message_item = ScheduleMessageItem(
                             user_id=target_user_id,
                             mem_cube_id=mem_cube_id,
-                            label=MEM_READ_LABEL,
+                            label=MEM_READ_TASK_LABEL,
                             content=json.dumps(mem_ids),
                             timestamp=datetime.utcnow(),
                         )
@@ -889,7 +889,7 @@ def process_preference_memory():
                         message_item = ScheduleMessageItem(
                             user_id=target_user_id,
                             mem_cube_id=mem_cube_id,
-                            label=ADD_LABEL,
+                            label=ADD_TASK_LABEL,
                             content=json.dumps(mem_ids),
                             timestamp=datetime.utcnow(),
                         )
@@ -920,7 +920,7 @@ def process_preference_memory():
                 message_item = ScheduleMessageItem(
                     user_id=target_user_id,
                     mem_cube_id=mem_cube_id,
-                    label=ADD_LABEL,
+                    label=ADD_TASK_LABEL,
                     content=json.dumps(mem_ids),
                     timestamp=datetime.utcnow(),
                 )

src/memos/mem_os/product.py

@@ -29,11 +29,11 @@
     prepare_reference_data,
     process_streaming_references_complete,
 )
-from memos.mem_scheduler.schemas.general_schemas import (
-    ANSWER_LABEL,
-    QUERY_LABEL,
-)
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import (
+    ANSWER_TASK_LABEL,
+    QUERY_TASK_LABEL,
+)
 from memos.mem_user.persistent_factory import PersistentUserManagerFactory
 from memos.mem_user.user_manager import UserRole
 from memos.memories.textual.item import (
@@ -710,7 +710,7 @@ async def _post_chat_processing(
                     logger.warning(f"Failed to send chat notification (async): {e}")
 
             self._send_message_to_scheduler(
-                user_id=user_id, mem_cube_id=cube_id, query=clean_response, label=ANSWER_LABEL
+                user_id=user_id, mem_cube_id=cube_id, query=clean_response, label=ANSWER_TASK_LABEL
             )
 
             self.add(
@@ -1151,7 +1151,7 @@ def chat_with_references(
             f"time chat: search text_mem time user_id: {user_id} time is: {search_time_end - time_start}"
         )
         self._send_message_to_scheduler(
-            user_id=user_id, mem_cube_id=cube_id, query=query, label=QUERY_LABEL
+            user_id=user_id, mem_cube_id=cube_id, query=query, label=QUERY_TASK_LABEL
         )
         if memories_result:
             memories_list = memories_result[0]["memories"]

src/memos/mem_scheduler/analyzer/mos_for_test_scheduler.py

@@ -4,11 +4,13 @@
 from memos.log import get_logger
 from memos.mem_os.main import MOS
 from memos.mem_scheduler.schemas.general_schemas import (
-    ANSWER_LABEL,
     MONITOR_WORKING_MEMORY_TYPE,
-    QUERY_LABEL,
 )
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.schemas.task_schemas import (
+    ANSWER_TASK_LABEL,
+    QUERY_TASK_LABEL,
+)
 
 
 logger = get_logger(__name__)
@@ -427,7 +429,7 @@ def chat(self, query: str, user_id: str | None = None) -> str:
                 message_item = ScheduleMessageItem(
                     user_id=target_user_id,
                     mem_cube_id=mem_cube_id,
-                    label=QUERY_LABEL,
+                    label=QUERY_TASK_LABEL,
                     content=query,
                     timestamp=datetime.now(),
                 )
@@ -517,7 +519,7 @@ def chat(self, query: str, user_id: str | None = None) -> str:
                 message_item = ScheduleMessageItem(
                     user_id=target_user_id,
                     mem_cube_id=mem_cube_id,
-                    label=ANSWER_LABEL,
+                    label=ANSWER_TASK_LABEL,
                     content=response,
                     timestamp=datetime.now(),
                 )

src/memos/mem_scheduler/analyzer/scheduler_for_eval.py

@@ -7,10 +7,10 @@
 
 from memos.log import get_logger
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
-from memos.mem_scheduler.schemas.general_schemas import (
+from memos.mem_scheduler.schemas.monitor_schemas import QueryMonitorItem
+from memos.mem_scheduler.schemas.task_schemas import (
     DEFAULT_MAX_QUERY_KEY_WORDS,
 )
-from memos.mem_scheduler.schemas.monitor_schemas import QueryMonitorItem
 
 
 if TYPE_CHECKING:

src/memos/mem_scheduler/base_scheduler.py

@@ -43,6 +43,7 @@
 )
 from memos.mem_scheduler.schemas.monitor_schemas import MemoryMonitorItem
 from memos.mem_scheduler.task_schedule_modules.dispatcher import SchedulerDispatcher
+from memos.mem_scheduler.task_schedule_modules.orchestrator import SchedulerOrchestrator
 from memos.mem_scheduler.task_schedule_modules.task_queue import ScheduleTaskQueue
 from memos.mem_scheduler.utils import metrics
 from memos.mem_scheduler.utils.db_utils import get_utc_now
@@ -121,10 +122,12 @@ def __init__(self, config: BaseSchedulerConfig):
         self.max_internal_message_queue_size = self.config.get(
             "max_internal_message_queue_size", DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE
         )
+        self.orchestrator = SchedulerOrchestrator()
         self.memos_message_queue = ScheduleTaskQueue(
             use_redis_queue=self.use_redis_queue,
             maxsize=self.max_internal_message_queue_size,
             disabled_handlers=self.disabled_handlers,
+            orchestrator=self.orchestrator,
         )
         self.searcher: Searcher | None = None
         self.retriever: SchedulerRetriever | None = None
@@ -143,6 +146,7 @@ def __init__(self, config: BaseSchedulerConfig):
             status_tracker=self.status_tracker,
             metrics=self.metrics,
             submit_web_logs=self._submit_web_logs,
+            orchestrator=self.orchestrator,
         )
         # Task schedule monitor: initialize with underlying queue implementation
         self.get_status_parallel = self.config.get("get_status_parallel", True)
@@ -697,22 +701,22 @@ def get_web_log_messages(self) -> list[dict]:
                 break
 
         def _map_label(label: str) -> str:
-            from memos.mem_scheduler.schemas.general_schemas import (
-                ADD_LABEL,
-                ANSWER_LABEL,
-                MEM_ARCHIVE_LABEL,
-                MEM_ORGANIZE_LABEL,
-                MEM_UPDATE_LABEL,
-                QUERY_LABEL,
+            from memos.mem_scheduler.schemas.task_schemas import (
+                ADD_TASK_LABEL,
+                ANSWER_TASK_LABEL,
+                MEM_ARCHIVE_TASK_LABEL,
+                MEM_ORGANIZE_TASK_LABEL,
+                MEM_UPDATE_TASK_LABEL,
+                QUERY_TASK_LABEL,
             )
 
             mapping = {
-                QUERY_LABEL: "addMessage",
-                ANSWER_LABEL: "addMessage",
-                ADD_LABEL: "addMemory",
-                MEM_UPDATE_LABEL: "updateMemory",
-                MEM_ORGANIZE_LABEL: "mergeMemory",
-                MEM_ARCHIVE_LABEL: "archiveMemory",
+                QUERY_TASK_LABEL: "addMessage",
+                ANSWER_TASK_LABEL: "addMessage",
+                ADD_TASK_LABEL: "addMemory",
+                MEM_UPDATE_TASK_LABEL: "updateMemory",
+                MEM_ORGANIZE_TASK_LABEL: "mergeMemory",
+                MEM_ARCHIVE_TASK_LABEL: "archiveMemory",
             }
             return mapping.get(label, label)
 
@@ -785,7 +789,7 @@ def _message_consumer(self) -> None:
                         if enqueue_epoch is not None:
                             queue_wait_ms = max(0.0, now - enqueue_epoch) * 1000
 
-                        # Avoid pydantic attribute enforcement
+                        # Avoid pydantic field enforcement by using object.__setattr__
                         object.__setattr__(msg, "_dequeue_ts", now)
                         emit_monitor_event(
                             "dequeue",