Refactor: Conflict resolution, bug fixes mainly for the mem scheduler (#221)

* rebase to address conflicts

* fix bugs: fix a bug in retriever, and add new auth info for neo4j db

* fix bugs & new feat: fix bugs in mem_scheduler examples, and remove initialize working memories (logically uneccessary). change the function parameters of search as the function input info as an addition

* fix bugs: modify configs, examples, schedule handlers of mem_scheduler, and rename some variables

* modify datetime to utc date time

---------

Co-authored-by: CaralHsi <[email protected]>

Author: tangg555

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/mem_chat_config.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.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/data/config/mem_scheduler/memos_config_w_scheduler_and_openai.yaml

@@ -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/data/config/mem_scheduler/memos_config_wo_scheduler.yaml

@@ -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_os/chat_w_scheduler.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/memos_w_scheduler.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()