feat & fix bugs & refactor: mem scheudler add more functions of monitors and web logs, and refactor the code structure

Author: tangg555

examples/data/config/mem_scheduler/general_scheduler_config.yaml

@@ -1,9 +1,9 @@
 backend: general_scheduler
 config:
   top_k: 10
-  top_n: 5
+  top_n: 10
   act_mem_update_interval: 30
-  context_window_size: 5
+  context_window_size: 10
   thread_pool_max_workers: 5
   consume_interval_seconds: 1
   enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml

@@ -33,9 +33,9 @@ mem_scheduler:
   backend: "general_scheduler"
   config:
     top_k: 10
-    top_n: 5
+    top_n: 10
     act_mem_update_interval: 30
-    context_window_size: 5
+    context_window_size: 10
     thread_pool_max_workers: 10
     consume_interval_seconds: 1
     enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml

@@ -34,10 +34,10 @@ mem_reader:
 mem_scheduler:
   backend: "general_scheduler"
   config:
-    top_k: 2
-    top_n: 5
+    top_k: 10
+    top_n: 10
     act_mem_update_interval: 30
-    context_window_size: 5
+    context_window_size: 10
     thread_pool_max_workers: 10
     consume_interval_seconds: 1
     enable_parallel_dispatch: true

examples/mem_os/chat_w_scheduler.py

@@ -7,7 +7,7 @@
 from memos.configs.mem_os import MOSConfig
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_os.main import MOS
-from memos.mem_scheduler.utils import parse_yaml
+from memos.mem_scheduler.utils.misc_utils import parse_yaml
 
 
 # init MOS

examples/mem_scheduler/memos_w_scheduler.py

@@ -10,13 +10,13 @@
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_os.main import MOS
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.scheduler_factory import SchedulerFactory
+from memos.mem_scheduler.schemas.general_schemas import (
     ANSWER_LABEL,
     QUERY_LABEL,
-    ScheduleMessageItem,
 )
-from memos.mem_scheduler.scheduler_factory import SchedulerFactory
-from memos.mem_scheduler.utils import parse_yaml
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.utils.misc_utils import parse_yaml
 
 
 FILE_PATH = Path(__file__).absolute()

…s/mem_scheduler/schedule_chat_and_web.py …_scheduler/memos_w_scheduler_for_test.pyexamples/mem_scheduler/schedule_chat_and_web.py renamed to examples/mem_scheduler/memos_w_scheduler_for_test.py examples/mem_scheduler/schedule_chat_and_web.py renamed to examples/mem_scheduler/memos_w_scheduler_for_test.py

@@ -15,7 +15,7 @@
 
 
 if TYPE_CHECKING:
-    from memos.mem_scheduler.modules.schemas import (
+    from memos.mem_scheduler.schemas import (
         ScheduleLogForWebItem,
     )
 
@@ -169,11 +169,12 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
     mos.add(conversations, user_id=user_id, mem_cube_id=mem_cube_id)
 
     for item in questions:
+        print("===== Chat Start =====")
         query = item["question"]
-
+        print(f"Query:\n {query}\n")
         response = mos.chat(query=query, user_id=user_id)
-        print(f"Query:\n {query}\n\nAnswer:\n {response}")
-
+        print(f"Answer:\n {response}")
+        print("===== Chat End =====")
     show_web_logs(mos.mem_scheduler)
 
     mos.mem_scheduler.stop()

examples/mem_scheduler/rabbitmq_example.py

@@ -8,7 +8,7 @@
 def publish_message(rabbitmq_module, message):
     """Function to publish a message."""
     rabbitmq_module.rabbitmq_publish_message(message)
-    print(f"Published message: {message}")
+    print(f"Published message: {message}\n")
 
 
 def main():
@@ -24,8 +24,7 @@ def main():
         rabbitmq_module.initialize_rabbitmq(config=AuthConfig.from_local_yaml().rabbitmq)
 
     try:
-        # Start consumer
-        rabbitmq_module.rabbitmq_start_consuming()
+        rabbitmq_module.wait_for_connection_ready()
 
         # === Publish some test messages ===
         # List to hold thread references
@@ -38,6 +37,9 @@ def main():
             thread.start()
             threads.append(thread)
 
+        # Start consumer
+        rabbitmq_module.rabbitmq_start_consuming()
+
         # Join threads to ensure all messages are published before proceeding
         for thread in threads:
             thread.join()
@@ -47,7 +49,7 @@ def main():
 
     finally:
         # Give some time for cleanup
-        time.sleep(5)
+        time.sleep(3)
 
         # Close connections
         rabbitmq_module.rabbitmq_close()

examples/mem_scheduler/redis_example.py

@@ -8,8 +8,9 @@
 
 from memos.configs.mem_scheduler import SchedulerConfigFactory
 from memos.mem_cube.general import GeneralMemCube
-from memos.mem_scheduler.modules.schemas import QUERY_LABEL, ScheduleMessageItem
 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
 
 
 if TYPE_CHECKING:

examples/mem_scheduler/try_schedule_modules.py

@@ -14,16 +14,16 @@
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_scheduler.general_scheduler import GeneralScheduler
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.mos_for_test_scheduler import MOSForTestScheduler
+from memos.mem_scheduler.schemas.general_schemas import (
     NOT_APPLICABLE_TYPE,
     QUERY_LABEL,
-    ScheduleMessageItem,
 )
-from memos.mem_scheduler.mos_for_test_scheduler import MOSForTestScheduler
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 if TYPE_CHECKING:
-    from memos.mem_scheduler.modules.schemas import (
+    from memos.mem_scheduler.schemas import (
         ScheduleLogForWebItem,
     )
 
@@ -180,14 +180,14 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
         query = item["question"]
 
         # test process_session_turn
-        mos.mem_scheduler.process_session_turn(
+        working_memory, new_candidates = mos.mem_scheduler.process_session_turn(
             queries=[query],
             user_id=user_id,
             mem_cube_id=mem_cube_id,
             mem_cube=mem_cube,
             top_k=10,
-            query_history=None,
         )
+        print(f"\nnew_candidates: {[one.memory for one in new_candidates]}")
 
         # test query_consume
         message_item = ScheduleMessageItem(
@@ -198,7 +198,13 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
             content=query,
             timestamp=datetime.now(),
         )
+        print(
+            f"\nworking memories before: {[one.memory for one in mos.mem_scheduler.mem_cubes[mem_cube_id].text_mem.get_working_memory()]}"
+        )
         mos.mem_scheduler._query_message_consumer(messages=[message_item])
+        print(
+            f"\nworking memories after: {[one.memory for one in mos.mem_scheduler.mem_cubes[mem_cube_id].text_mem.get_working_memory()]}"
+        )
 
         # test activation memory update
         mos.mem_scheduler.update_activation_memory_periodically(

src/memos/configs/mem_scheduler.py

@@ -6,12 +6,12 @@
 from pydantic import ConfigDict, Field, field_validator, model_validator
 
 from memos.configs.base import BaseConfig
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.modules.misc import DictConversionMixin
+from memos.mem_scheduler.schemas.general_schemas import (
     BASE_DIR,
     DEFAULT_ACT_MEM_DUMP_PATH,
     DEFAULT_CONSUME_INTERVAL_SECONDS,
     DEFAULT_THREAD__POOL_MAX_WORKERS,
-    DictConversionMixin,
 )
 
 
@@ -21,6 +21,7 @@ class BaseSchedulerConfig(BaseConfig):
     top_k: int = Field(
         default=10, description="Number of top candidates to consider in initial retrieval"
     )
+    # TODO: The 'top_n' field is deprecated and will be removed in future versions.
     top_n: int = Field(default=5, description="Number of final results to return after processing")
     enable_parallel_dispatch: bool = Field(
         default=True, description="Whether to enable parallel message processing using thread pool"
@@ -48,7 +49,7 @@ class GeneralSchedulerConfig(BaseSchedulerConfig):
         default=300, description="Interval in seconds for updating activation memory"
     )
     context_window_size: int | None = Field(
-        default=5, description="Size of the context window for conversation history"
+        default=10, description="Size of the context window for conversation history"
     )
     act_mem_dump_path: str | None = Field(
         default=DEFAULT_ACT_MEM_DUMP_PATH,  # Replace with DEFAULT_ACT_MEM_DUMP_PATH