Refactor: Conflict resolution, bug fixes mainly for the mem scheduler

evaluation/scripts/run_locomo_eval.sh

@@ -6,12 +6,12 @@ VERSION="072001"
 WORKERS=10
 TOPK=20
 
-# echo "Running locomo_ingestion.py..."
-# CUDA_VISIBLE_DEVICES=0 python scripts/locomo/locomo_ingestion.py --lib $LIB --version $VERSION --workers $WORKERS
-# if [ $? -ne 0 ]; then
-#     echo "Error running locomo_ingestion.py"
-#     exit 1
-# fi
+ echo "Running locomo_ingestion.py..."
+ CUDA_VISIBLE_DEVICES=0 python scripts/locomo/locomo_ingestion.py --lib $LIB --version $VERSION --workers $WORKERS
+ if [ $? -ne 0 ]; then
+     echo "Error running locomo_ingestion.py"
+     exit 1
+ fi
 
 echo "Running locomo_search.py..."
 CUDA_VISIBLE_DEVICES=0 python scripts/locomo/locomo_search.py --lib $LIB --version $VERSION --top_k $TOPK --workers $WORKERS

examples/data/config/mem_scheduler/general_scheduler_config.yaml

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

examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml

@@ -33,12 +33,14 @@ mem_scheduler:
   backend: "general_scheduler"
   config:
     top_k: 10
-    top_n: 10
     act_mem_update_interval: 30
     context_window_size: 10
     thread_pool_max_workers: 10
     consume_interval_seconds: 1
+    working_mem_monitor_capacity: 20
+    activation_mem_monitor_capacity: 5
     enable_parallel_dispatch: true
+    enable_activation_memory: true
 max_turns_window: 20
 top_k: 5
 enable_textual_memory: true

examples/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml

@@ -35,13 +35,14 @@ mem_scheduler:
   backend: "general_scheduler"
   config:
     top_k: 10
-    top_n: 10
     act_mem_update_interval: 30
     context_window_size: 10
     thread_pool_max_workers: 10
     consume_interval_seconds: 1
+    working_mem_monitor_capacity: 20
+    activation_mem_monitor_capacity: 5
     enable_parallel_dispatch: true
-    enable_act_memory_update: false
+    enable_activation_memory: true
 max_turns_window: 20
 top_k: 5
 enable_textual_memory: true

examples/mem_os/chat_w_scheduler.py

@@ -1,40 +1,61 @@
 import shutil
-import uuid
+import sys
 
 from pathlib import Path
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
+from memos.configs.mem_scheduler import AuthConfig
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_os.main import MOS
-from memos.mem_scheduler.utils.misc_utils import parse_yaml
 
 
-# init MOS
-config = parse_yaml("./examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml")
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent
+sys.path.insert(0, str(BASE_DIR))  # Enable execution from any working directory
 
-mos_config = MOSConfig(**config)
-mos = MOS(mos_config)
 
-# create user
-user_id = str(uuid.uuid4())
-mos.create_user(user_id=user_id)
+# set configs
+mos_config = MOSConfig.from_yaml_file(
+    f"{BASE_DIR}/examples/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml"
+)
 
-config = GeneralMemCubeConfig.from_yaml_file(
-    "./examples/data/config/mem_scheduler/mem_cube_config.yaml"
+mem_cube_config = GeneralMemCubeConfig.from_yaml_file(
+    f"{BASE_DIR}/examples/data/config/mem_scheduler/mem_cube_config.yaml"
 )
+
+# default local graphdb uri
+if AuthConfig.default_config_exists():
+    auth_config = AuthConfig.from_local_config()
+
+    mos_config.mem_reader.config.llm.config.api_key = auth_config.openai.api_key
+    mos_config.mem_reader.config.llm.config.api_base = auth_config.openai.base_url
+
+    mem_cube_config.text_mem.config.graph_db.config.uri = auth_config.graph_db.uri
+    mem_cube_config.text_mem.config.graph_db.config.user = auth_config.graph_db.user
+    mem_cube_config.text_mem.config.graph_db.config.password = auth_config.graph_db.password
+    mem_cube_config.text_mem.config.graph_db.config.db_name = auth_config.graph_db.db_name
+    mem_cube_config.text_mem.config.graph_db.config.auto_create = auth_config.graph_db.auto_create
+
+# Initialization
+mos = MOS(mos_config)
+
+user_id = "user_1"
+mos.create_user(user_id)
+
 mem_cube_id = "mem_cube_5"
-mem_cube_name_or_path = f"./outputs/mem_scheduler/{user_id}/{mem_cube_id}"
+mem_cube_name_or_path = f"{BASE_DIR}/outputs/mem_scheduler/{user_id}/{mem_cube_id}"
+
 if Path(mem_cube_name_or_path).exists():
     shutil.rmtree(mem_cube_name_or_path)
     print(f"{mem_cube_name_or_path} is not empty, and has been removed.")
 
-mem_cube = GeneralMemCube(config)
+mem_cube = GeneralMemCube(mem_cube_config)
 mem_cube.dump(mem_cube_name_or_path)
-
 mos.register_mem_cube(
     mem_cube_name_or_path=mem_cube_name_or_path, mem_cube_id=mem_cube_id, user_id=user_id
 )
+
 messages = [
     {"role": "user", "content": "I like playing football."},
     {"role": "assistant", "content": "I like playing football too."},
@@ -51,8 +72,3 @@
     for node in retrieved_memories["text_mem"][0]["memories"]["nodes"]:
         if node["metadata"]["memory_type"] == "WorkingMemory":
             print(f"🤖 [Assistant]working mem : {node['memory']}\n")
-    if retrieved_memories["act_mem"][0]["memories"]:
-        for act_mem in retrieved_memories["act_mem"][0]["memories"]:
-            print(f"🤖 [Assistant]act_mem: {act_mem['memory']}\n")
-    else:
-        print("🤖 [Assistant]act_mem: None\n")

examples/mem_scheduler/memos_w_scheduler.py

@@ -85,7 +85,7 @@ def run_with_scheduler_init():
 
     # default local graphdb uri
     if AuthConfig.default_config_exists():
-        auth_config = AuthConfig.from_local_yaml()
+        auth_config = AuthConfig.from_local_config()
 
         mos_config.mem_reader.config.llm.config.api_key = auth_config.openai.api_key
         mos_config.mem_reader.config.llm.config.api_base = auth_config.openai.base_url

examples/mem_scheduler/memos_w_scheduler_for_test.py

@@ -97,7 +97,7 @@ def init_task():
 
     # default local graphdb uri
     if AuthConfig.default_config_exists():
-        auth_config = AuthConfig.from_local_yaml()
+        auth_config = AuthConfig.from_local_config()
 
         mos_config.mem_reader.config.llm.config.api_key = auth_config.openai.api_key
         mos_config.mem_reader.config.llm.config.api_base = auth_config.openai.base_url

examples/mem_scheduler/rabbitmq_example.py

@@ -21,7 +21,7 @@ def main():
         print("Please set configs for rabbitmq.")
         return
     else:
-        rabbitmq_module.initialize_rabbitmq(config=AuthConfig.from_local_yaml().rabbitmq)
+        rabbitmq_module.initialize_rabbitmq(config=AuthConfig.from_local_config().rabbitmq)
 
     try:
         rabbitmq_module.wait_for_connection_ready()

examples/mem_scheduler/try_schedule_modules.py

@@ -145,7 +145,7 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
 
     # default local graphdb uri
     if AuthConfig.default_config_exists():
-        auth_config = AuthConfig.from_local_yaml()
+        auth_config = AuthConfig.from_local_config()
 
         mos_config.mem_reader.config.llm.config.api_key = auth_config.openai.api_key
         mos_config.mem_reader.config.llm.config.api_base = auth_config.openai.base_url

src/memos/api/config.py

@@ -266,7 +266,7 @@ def get_scheduler_config() -> dict[str, Any]:
                     "MOS_SCHEDULER_ENABLE_PARALLEL_DISPATCH", "true"
                 ).lower()
                 == "true",
-                "enable_act_memory_update": True,
+                "enable_activation_memory": True,
             },
         }
 

src/memos/configs/mem_scheduler.py

@@ -21,8 +21,6 @@ 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"
     )
@@ -55,9 +53,15 @@ class GeneralSchedulerConfig(BaseSchedulerConfig):
         default=DEFAULT_ACT_MEM_DUMP_PATH,  # Replace with DEFAULT_ACT_MEM_DUMP_PATH
         description="File path for dumping activation memory",
     )
-    enable_act_memory_update: bool = Field(
+    enable_activation_memory: bool = Field(
         default=False, description="Whether to enable automatic activation memory updates"
     )
+    working_mem_monitor_capacity: int = Field(
+        default=30, description="Capacity of the working memory monitor"
+    )
+    activation_mem_monitor_capacity: int = Field(
+        default=20, description="Capacity of the activation memory monitor"
+    )
 
 
 class SchedulerConfigFactory(BaseConfig):
@@ -137,29 +141,46 @@ class AuthConfig(BaseConfig, DictConversionMixin):
     )
 
     @classmethod
-    def from_local_yaml(cls, config_path: str | None = None) -> "AuthConfig":
+    def from_local_config(cls, config_path: str | Path | None = None) -> "AuthConfig":
         """
-        Load configuration from YAML file
+        Load configuration from either a YAML or JSON file based on file extension.
+
+        Automatically detects file type (YAML or JSON) from the file extension
+        and uses the appropriate parser. If no path is provided, uses the default
+        configuration path (YAML) or its JSON counterpart.
 
         Args:
-            config_path: Path to YAML configuration file
+            config_path: Optional path to configuration file.
+                         If not provided, uses default configuration path.
 
         Returns:
-            AuthConfig instance
+            AuthConfig instance populated with data from the configuration file.
 
         Raises:
-            FileNotFoundError: If config file doesn't exist
-            ValueError: If YAML parsing or validation fails
+            FileNotFoundError: If the specified or default configuration file does not exist.
+            ValueError: If file extension is not .yaml/.yml or .json, or if parsing fails.
         """
-
+        # Determine config path
         if config_path is None:
             config_path = cls.default_config_path
 
-        # Check file exists
-        if not Path(config_path).exists():
-            raise FileNotFoundError(f"Config file not found: {config_path}")
-
-        return cls.from_yaml_file(yaml_path=config_path)
+        # Validate file existence
+        config_path_obj = Path(config_path)
+        if not config_path_obj.exists():
+            raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+        # Get file extension and determine parser
+        file_ext = config_path_obj.suffix.lower()
+
+        if file_ext in (".yaml", ".yml"):
+            return cls.from_yaml_file(yaml_path=str(config_path_obj))
+        elif file_ext == ".json":
+            return cls.from_json_file(json_path=str(config_path_obj))
+        else:
+            raise ValueError(
+                f"Unsupported file format: {file_ext}. "
+                "Please use YAML (.yaml, .yml) or JSON (.json) files."
+            )
 
     def set_openai_config_to_environment(self):
         # Set environment variables

src/memos/mem_os/main.py

@@ -219,7 +219,7 @@ def _chat_with_cot_enhancement(
                         mem_cube=mem_cube,
                         label=ANSWER_LABEL,
                         content=enhanced_response,
-                        timestamp=datetime.now(),
+                        timestamp=datetime.now().isoformat(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 

src/memos/mem_os/product.py

@@ -491,7 +491,7 @@ def _send_message_to_scheduler(
                 mem_cube=self.mem_cubes[mem_cube_id],
                 label=label,
                 content=query,
-                timestamp=datetime.now(),
+                timestamp=datetime.utcnow(),
             )
             self.mem_scheduler.submit_messages(messages=[message_item])
 

src/memos/mem_os/utils/default_config.py

@@ -112,7 +112,7 @@ def get_default_config(
                 "thread_pool_max_workers": kwargs.get("scheduler_thread_pool_max_workers", 10),
                 "consume_interval_seconds": kwargs.get("scheduler_consume_interval_seconds", 3),
                 "enable_parallel_dispatch": kwargs.get("scheduler_enable_parallel_dispatch", True),
-                "enable_act_memory_update": True,
+                "enable_activation_memory": True,
             },
         }