Merge branch 'dev' into feat/file_parser_debug

Author: CaralHsi

evaluation/scripts/locomo/locomo_eval.py

@@ -311,7 +311,7 @@ async def main(frame, version="default", options=None, num_runs=1, max_workers=4
     with open(response_path) as file:
         locomo_responses = json.load(file)
 
-    num_users = 2
+    num_users = 10
     all_grades = {}
 
     total_responses_count = sum(

evaluation/scripts/utils/client.py

@@ -189,7 +189,9 @@ def search(self, query, user_id, top_k):
         )
         response = requests.request("POST", url, data=payload, headers=self.headers)
         assert response.status_code == 200, response.text
-        assert json.loads(response.text)["message"] == "Memory searched successfully", response.text
+        assert json.loads(response.text)["message"] == "Search completed successfully", (
+            response.text
+        )
         return json.loads(response.text)["data"]
 
 

examples/data/config/mem_scheduler/general_scheduler_config.yaml

@@ -4,7 +4,7 @@ config:
   act_mem_update_interval: 30
   context_window_size: 10
   thread_pool_max_workers: 5
-  consume_interval_seconds: 1
+  consume_interval_seconds: 0.01
   working_mem_monitor_capacity: 20
   activation_mem_monitor_capacity: 5
   enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_optimized_scheduler.yaml

@@ -38,7 +38,7 @@ mem_scheduler:
     act_mem_update_interval: 30
     context_window_size: 10
     thread_pool_max_workers: 10
-    consume_interval_seconds: 1
+    consume_interval_seconds: 0.01
     working_mem_monitor_capacity: 20
     activation_mem_monitor_capacity: 5
     enable_parallel_dispatch: true

examples/data/config/mem_scheduler/memos_config_w_scheduler.yaml

@@ -38,7 +38,7 @@ mem_scheduler:
     act_mem_update_interval: 30
     context_window_size: 10
     thread_pool_max_workers: 10
-    consume_interval_seconds: 1
+    consume_interval_seconds: 0.01
     working_mem_monitor_capacity: 20
     activation_mem_monitor_capacity: 5
     enable_parallel_dispatch: true

examples/mem_chunk/markdown_chunk.py

@@ -0,0 +1,33 @@
+from memos.chunkers import ChunkerFactory
+from memos.configs.chunker import ChunkerConfigFactory
+
+
+config = ChunkerConfigFactory.model_validate(
+    {
+        "backend": "markdown",
+        "config": {
+            "chunk_size": 1000,
+            "chunk_overlap": 100,
+            "recursive": True,
+        },
+    }
+)
+
+chunker = ChunkerFactory.from_config(config)
+
+text = """
+# Header 1
+This is the first sentence. This is the second sentence.
+And here's a third one with some additional context.
+
+# Header 2
+This is the fourth sentence. This is the fifth sentence.
+And here's a sixth one with some additional context.
+
+# Header 3
+This is the seventh sentence. This is the eighth sentence.
+And here's a ninth one with some additional context.
+"""
+chunks = chunker.chunk(text)
+for chunk in chunks:
+    print("doc:", chunk)

examples/mem_scheduler/api_w_scheduler.py

@@ -17,7 +17,6 @@
 print(f"Queue maxsize: {getattr(mem_scheduler.memos_message_queue, 'maxsize', 'N/A')}")
 print("=====================================\n")
 
-mem_scheduler.memos_message_queue.debug_mode_on()
 queue = mem_scheduler.memos_message_queue
 queue.clear()
 

examples/mem_scheduler/task_fair_schedule.py

@@ -0,0 +1,87 @@
+import sys
+
+from collections import defaultdict
+from pathlib import Path
+
+from memos.api.routers.server_router import mem_scheduler
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent
+sys.path.insert(0, str(BASE_DIR))
+
+
+def make_message(user_id: str, mem_cube_id: str, label: str, idx: int | str) -> ScheduleMessageItem:
+    return ScheduleMessageItem(
+        item_id=f"{user_id}:{mem_cube_id}:{label}:{idx}",
+        user_id=user_id,
+        mem_cube_id=mem_cube_id,
+        label=label,
+        content=f"msg-{idx} for {user_id}/{mem_cube_id}/{label}",
+    )
+
+
+def seed_messages_for_test_fairness(queue, combos, per_stream):
+    # send overwhelm message by one user
+    (u, c, label) = combos[0]
+    task_target = 100
+    print(f"{u}:{c}:{label} submit {task_target} messages")
+    for i in range(task_target):
+        msg = make_message(u, c, label, f"overwhelm_{i}")
+        queue.submit_messages(msg)
+
+    for u, c, label in combos:
+        print(f"{u}:{c}:{label} submit {per_stream} messages")
+        for i in range(per_stream):
+            msg = make_message(u, c, label, i)
+            queue.submit_messages(msg)
+    print("======= seed_messages Done ===========")
+
+
+def count_by_stream(messages):
+    counts = defaultdict(int)
+    for m in messages:
+        key = f"{m.user_id}:{m.mem_cube_id}:{m.label}"
+        counts[key] += 1
+    return counts
+
+
+def run_fair_redis_schedule(batch_size: int = 3):
+    print("=== Redis Fairness Demo ===")
+    print(f"use_redis_queue: {mem_scheduler.use_redis_queue}")
+    mem_scheduler.consume_batch = batch_size
+    queue = mem_scheduler.memos_message_queue
+
+    # Isolate and clear queue
+    queue.clear()
+
+    # Define multiple streams: (user_id, mem_cube_id, task_label)
+    combos = [
+        ("u1", "u1", "labelX"),
+        ("u1", "u1", "labelY"),
+        ("u2", "u2", "labelX"),
+        ("u2", "u2", "labelY"),
+    ]
+    per_stream = 5
+
+    # Seed messages evenly across streams
+    seed_messages_for_test_fairness(queue, combos, per_stream)
+
+    # Compute target batch size (fair split across streams)
+    print(f"Request batch_size={batch_size} for {len(combos)} streams")
+
+    for _ in range(len(combos)):
+        # Fetch one brokered pack
+        msgs = queue.get_messages(batch_size=batch_size)
+        print(f"Fetched {len(msgs)} messages in first pack")
+
+        # Check fairness: counts per stream
+        counts = count_by_stream(msgs)
+        for k in sorted(counts):
+            print(f"{k}: {counts[k]}")
+
+
+if __name__ == "__main__":
+    # task 1 fair redis schedule
+    run_fair_redis_schedule()

examples/mem_scheduler/task_stop_rerun.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+from time import sleep
+
+# Note: we skip API handler status/wait utilities in this demo
+from memos.api.routers.server_router import mem_scheduler
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+
+
+# Debug: Print scheduler configuration
+print("=== Scheduler Configuration Debug ===")
+print(f"Scheduler type: {type(mem_scheduler).__name__}")
+print(f"Config: {mem_scheduler.config}")
+print(f"use_redis_queue: {mem_scheduler.use_redis_queue}")
+print(f"Queue type: {type(mem_scheduler.memos_message_queue).__name__}")
+print(f"Queue maxsize: {getattr(mem_scheduler.memos_message_queue, 'maxsize', 'N/A')}")
+print("=====================================\n")
+
+queue = mem_scheduler.memos_message_queue
+
+
+# Define a handler function
+def my_test_handler(messages: list[ScheduleMessageItem]):
+    print(f"My test handler received {len(messages)} messages: {[one.item_id for one in messages]}")
+    for msg in messages:
+        # Create a file named by task_id (use item_id as numeric id 0..99)
+        task_id = str(msg.item_id)
+        file_path = tmp_dir / f"{task_id}.txt"
+        try:
+            print(f"writing {file_path}...")
+            file_path.write_text(f"Task {task_id} processed.\n")
+        except Exception as e:
+            print(f"Failed to write {file_path}: {e}")
+
+
+def submit_tasks():
+    mem_scheduler.memos_message_queue.clear()
+
+    # Create 100 messages (task_id 0..99)
+    users = ["user_A", "user_B"]
+    messages_to_send = [
+        ScheduleMessageItem(
+            item_id=str(i),
+            user_id=users[i % 2],
+            mem_cube_id="test_mem_cube",
+            label=TEST_HANDLER_LABEL,
+            content=f"Create file for task {i}",
+        )
+        for i in range(100)
+    ]
+    # Submit messages in batch and print completion
+    print(f"Submitting {len(messages_to_send)} messages to the scheduler...")
+    mem_scheduler.memos_message_queue.submit_messages(messages_to_send)
+    print(f"Task submission done! tasks in queue: {mem_scheduler.get_tasks_status()}")
+
+
+# Register the handler
+TEST_HANDLER_LABEL = "test_handler"
+mem_scheduler.register_handlers({TEST_HANDLER_LABEL: my_test_handler})
+
+
+tmp_dir = Path("./tmp")
+tmp_dir.mkdir(exist_ok=True)
+
+# Test stop-and-restart: if tmp already has >1 files, skip submission and print info
+existing_count = len(list(Path("tmp").glob("*.txt"))) if Path("tmp").exists() else 0
+if existing_count > 1:
+    print(f"Skip submission: found {existing_count} files in tmp (>1), continue processing")
+else:
+    submit_tasks()
+
+# 6. Wait until tmp has 100 files or timeout
+poll_interval = 0.01
+expected = 100
+tmp_dir = Path("tmp")
+while mem_scheduler.get_tasks_status()["remaining"] != 0:
+    count = len(list(tmp_dir.glob("*.txt"))) if tmp_dir.exists() else 0
+    tasks_status = mem_scheduler.get_tasks_status()
+    mem_scheduler.print_tasks_status(tasks_status=tasks_status)
+    print(f"[Monitor] Files in tmp: {count}/{expected}")
+    sleep(poll_interval)
+print(f"[Result] Final files in tmp: {len(list(tmp_dir.glob('*.txt')))})")
+
+# 7. Stop the scheduler
+print("Stopping the scheduler...")
+mem_scheduler.stop()