feat: introduce async memory add for TreeTextMemory using MemScheduler (#373)

* feat: define mem-read schedular message&consumer; add async mem-reader mode in core;

* feat: add fast/fine mode in mem-reader;

* feat: add mem-reader in scheduler

* feat: change async remove

* feat: modify async-add in core.py

* feat: add 'remove and refresh memory in schedular'

* feat: add naive fast mode in mem-reader

* feat: finish fast mode in mem-reader

* feat: add token-based window splitting and concurrency improvements

* feat: add split chunker into  mode in simple struct mem reader

* feat: update async-mode add

* chore: update gitignore

* feat: improve database note write performance

* feat: fix mem-read scheduler

* fix: nebula group-by bug

* fix: bug in adding mem scheduler

* fix: nebula index; mem-reader chat-time;

* format: searcher

* fix: some bug in shceduler and mem-reader

* feat: add mem-organize in scheduler

* feat: add tree.mode to config; modify scheduler config

* fix: test bug

Author: CaralHsi

.gitignore

@@ -15,6 +15,7 @@ evaluation/.env
 !evaluation/configs-example/*.json
 evaluation/configs/*
 **tree_textual_memory_locomo**
+**script.py**
 .env
 evaluation/scripts/personamem
 

examples/mem_reader/reader.py

@@ -28,7 +28,7 @@ def __init__(self, config: SentenceChunkerConfig):
         )
         logger.info(f"Initialized SentenceChunker with config: {config}")
 
-    def chunk(self, text: str) -> list[Chunk]:
+    def chunk(self, text: str) -> list[str] | list[Chunk]:
         """Chunk the given text into smaller chunks based on sentences."""
         chonkie_chunks = self.chunker.chunk(text)
 

src/memos/chunkers/sentence_chunker.py

@@ -28,13 +28,11 @@ class BaseSchedulerConfig(BaseConfig):
     thread_pool_max_workers: int = Field(
         default=DEFAULT_THREAD_POOL_MAX_WORKERS,
         gt=1,
-        lt=20,
         description=f"Maximum worker threads in pool (default: {DEFAULT_THREAD_POOL_MAX_WORKERS})",
     )
     consume_interval_seconds: float = Field(
         default=DEFAULT_CONSUME_INTERVAL_SECONDS,
         gt=0,
-        le=60,
         description=f"Interval for consuming messages from queue in seconds (default: {DEFAULT_CONSUME_INTERVAL_SECONDS})",
     )
     auth_config_path: str | None = Field(

src/memos/configs/mem_scheduler.py

@@ -179,6 +179,11 @@ class TreeTextMemoryConfig(BaseTextMemoryConfig):
         ),
     )
 
+    mode: str | None = Field(
+        default="sync",
+        description=("whether use asynchronous mode in memory add"),
+    )
+
 
 class SimpleTreeTextMemoryConfig(TreeTextMemoryConfig):
     """Simple tree text memory configuration class."""

src/memos/configs/memory.py

@@ -440,20 +440,22 @@ def remove_oldest_memory(
             memory_type (str): Memory type (e.g., 'WorkingMemory', 'LongTermMemory').
             keep_latest (int): Number of latest WorkingMemory entries to keep.
         """
-        optional_condition = ""
-
-        user_name = user_name if user_name else self.config.user_name
-
-        optional_condition = f"AND n.user_name = '{user_name}'"
-        query = f"""
-            MATCH (n@Memory /*+ INDEX(idx_memory_user_name) */)
-            WHERE n.memory_type = '{memory_type}'
-            {optional_condition}
-            ORDER BY n.updated_at DESC
-            OFFSET {int(keep_latest)}
-            DETACH DELETE n
-        """
-        self.execute_query(query)
+        try:
+            user_name = user_name if user_name else self.config.user_name
+            optional_condition = f"AND n.user_name = '{user_name}'"
+            count = self.count_nodes(memory_type, user_name)
+            if count > keep_latest:
+                delete_query = f"""
+                    MATCH (n@Memory /*+ INDEX(idx_memory_user_name) */)
+                    WHERE n.memory_type = '{memory_type}'
+                    {optional_condition}
+                    ORDER BY n.updated_at DESC
+                    OFFSET {int(keep_latest)}
+                    DETACH DELETE n
+                """
+                self.execute_query(delete_query)
+        except Exception as e:
+            logger.warning(f"Delete old mem error: {e}")
 
     @timed
     def add_node(
@@ -1175,7 +1177,6 @@ def get_grouped_counts(
             MATCH (n /*+ INDEX(idx_memory_user_name) */)
             {where_clause}
             RETURN {", ".join(return_fields)}, COUNT(n) AS count
-            GROUP BY {", ".join(group_by_fields)}
             """
         result = self.execute_query(gql)  # Pure GQL string execution
 
@@ -1620,7 +1621,13 @@ def _create_basic_property_indexes(self) -> None:
         Create standard B-tree indexes on user_name when use Shared Database
         Multi-Tenant Mode.
         """
-        fields = ["status", "memory_type", "created_at", "updated_at", "user_name"]
+        fields = [
+            "status",
+            "memory_type",
+            "created_at",
+            "updated_at",
+            "user_name",
+        ]
 
         for field in fields:
             index_name = f"idx_memory_{field}"

src/memos/graph_dbs/nebular.py

@@ -669,7 +669,7 @@ def search_by_embedding(
             vector (list[float]): The embedding vector representing query semantics.
             top_k (int): Number of top similar nodes to retrieve.
             scope (str, optional): Memory type filter (e.g., 'WorkingMemory', 'LongTermMemory').
-            status (str, optional): Node status filter (e.g., 'active', 'archived').
+            status (str, optional): Node status filter (e.g., 'activated', 'archived').
                             If provided, restricts results to nodes with matching status.
             threshold (float, optional): Minimum similarity score threshold (0 ~ 1).
             search_filter (dict, optional): Additional metadata filters for search results.

src/memos/graph_dbs/neo4j.py

@@ -17,6 +17,7 @@
 from memos.mem_scheduler.schemas.general_schemas import (
     ADD_LABEL,
     ANSWER_LABEL,
+    MEM_READ_LABEL,
     QUERY_LABEL,
 )
 from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
@@ -70,6 +71,7 @@ def __init__(self, config: MOSConfig, user_manager: UserManager | None = None):
         if self.enable_mem_scheduler:
             self._mem_scheduler = self._initialize_mem_scheduler()
             self._mem_scheduler.mem_cubes = self.mem_cubes
+            self._mem_scheduler.mem_reader = self.mem_reader
         else:
             self._mem_scheduler: GeneralScheduler = None
 
@@ -681,6 +683,12 @@ def add(
         logger.info(
             f"time add: get mem_cube_id check in mem_cubes time user_id: {target_user_id} time is: {time.time() - time_start_0}"
         )
+        sync_mode = self.mem_cubes[mem_cube_id].text_mem.mode
+        if sync_mode == "async":
+            assert self.mem_scheduler is not None, (
+                "Mem-Scheduler must be working when use asynchronous memory adding."
+            )
+        logger.debug(f"Mem-reader mode is: {sync_mode}")
         time_start_1 = time.time()
         if (
             (messages is not None)
@@ -690,6 +698,7 @@ def add(
             logger.info(
                 f"time add: messages is not None and enable_textual_memory and text_mem is not None time user_id: {target_user_id} time is: {time.time() - time_start_1}"
             )
+
             if self.mem_cubes[mem_cube_id].config.text_mem.backend != "tree_text":
                 add_memory = []
                 metadata = TextualMemoryMetadata(
@@ -707,21 +716,30 @@ def add(
                     messages_list,
                     type="chat",
                     info={"user_id": target_user_id, "session_id": target_session_id},
+                    mode="fast" if sync_mode == "async" else "fine",
                 )
                 logger.info(
                     f"time add: get mem_reader time user_id: {target_user_id} time is: {time.time() - time_start_2}"
                 )
-                mem_ids = []
-                for mem in memories:
-                    mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
-                    mem_ids.extend(mem_id_list)
-                    logger.info(
-                        f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_id_list}"
-                    )
-
+                memories_flatten = [m for m_list in memories for m in m_list]
+                mem_ids: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(memories_flatten)
+                logger.info(
+                    f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_ids}"
+                )
                 # submit messages for scheduler
                 if self.enable_mem_scheduler and self.mem_scheduler is not None:
                     mem_cube = self.mem_cubes[mem_cube_id]
+                    if sync_mode == "async":
+                        message_item = ScheduleMessageItem(
+                            user_id=target_user_id,
+                            mem_cube_id=mem_cube_id,
+                            mem_cube=mem_cube,
+                            label=MEM_READ_LABEL,
+                            content=json.dumps(mem_ids),
+                            timestamp=datetime.utcnow(),
+                        )
+                        self.mem_scheduler.submit_messages(messages=[message_item])
+
                     message_item = ScheduleMessageItem(
                         user_id=target_user_id,
                         mem_cube_id=mem_cube_id,
@@ -749,10 +767,12 @@ def add(
                 messages_list = [
                     [{"role": "user", "content": memory_content}]
                 ]  # for only user-str input and convert message
+
                 memories = self.mem_reader.get_memory(
                     messages_list,
                     type="chat",
                     info={"user_id": target_user_id, "session_id": target_session_id},
+                    mode="fast" if sync_mode == "async" else "fine",
                 )
 
                 mem_ids = []
@@ -766,6 +786,16 @@ def add(
                 # submit messages for scheduler
                 if self.enable_mem_scheduler and self.mem_scheduler is not None:
                     mem_cube = self.mem_cubes[mem_cube_id]
+                    if sync_mode == "async":
+                        message_item = ScheduleMessageItem(
+                            user_id=target_user_id,
+                            mem_cube_id=mem_cube_id,
+                            mem_cube=mem_cube,
+                            label=MEM_READ_LABEL,
+                            content=json.dumps(mem_ids),
+                            timestamp=datetime.utcnow(),
+                        )
+                        self.mem_scheduler.submit_messages(messages=[message_item])
                     message_item = ScheduleMessageItem(
                         user_id=target_user_id,
                         mem_cube_id=mem_cube_id,

src/memos/mem_os/core.py

@@ -18,10 +18,17 @@ def get_scene_data_info(self, scene_data: list, type: str) -> list[str]:
 
     @abstractmethod
     def get_memory(
-        self, scene_data: list, type: str, info: dict[str, Any]
+        self, scene_data: list, type: str, info: dict[str, Any], mode: str = "fast"
     ) -> list[list[TextualMemoryItem]]:
         """Various types of memories extracted from scene_data"""
 
     @abstractmethod
     def transform_memreader(self, data: dict) -> list[TextualMemoryItem]:
         """Transform the memory data into a list of TextualMemoryItem objects."""
+
+    @abstractmethod
+    def fine_transfer_simple_mem(
+        self, input_memories: list[list[TextualMemoryItem]], type: str
+    ) -> list[list[TextualMemoryItem]]:
+        """Fine Transform TextualMemoryItem List into another list of
+        TextualMemoryItem objects via calling llm to better understand users."""