feat: add try catch for mem scheduler (#144)

* feat: update config

* fix:dim

* change dim

* fix:change default db

* fix：delay

* fix:len

* fix:change recently mem size

* fix:dup node error

* fix: remove mock_data

* fix: change config

* feat: reorganize code

* add: add json parse for en

* fix:change user_id

* fix: logger info

* fix: remove unsed change

* feat: add topk for api

* feat: add logger

* fix:fix  scheduler logs and

---------

Co-authored-by: Travis Tang <[email protected]>

Author: fridayL

src/memos/api/product_models.py

@@ -150,6 +150,7 @@ class SearchRequest(BaseRequest):
     user_id: str = Field(..., description="User ID")
     query: str = Field(..., description="Search query")
     mem_cube_id: str | None = Field(None, description="Cube ID to search in")
+    top_k: int = Field(10, description="Number of results to return")
 
 
 class SuggestionRequest(BaseRequest):

src/memos/api/routers/product_router.py

@@ -194,6 +194,7 @@ async def search_memories(search_req: SearchRequest):
             query=search_req.query,
             user_id=search_req.user_id,
             install_cube_ids=[search_req.mem_cube_id] if search_req.mem_cube_id else None,
+            top_k=search_req.top_k,
         )
         return SearchResponse(message="Search completed successfully", data=result)
 

src/memos/mem_os/core.py

@@ -631,6 +631,9 @@ def add(
                 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}"
+                    )
 
                 # submit messages for scheduler
                 if self.enable_mem_scheduler and self.mem_scheduler is not None:
@@ -671,6 +674,9 @@ def add(
                 mem_ids = []
                 for mem in memories:
                     mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    logger.info(
+                        f"Added memory user {target_user_id} to memcube {mem_cube_id}: {mem_id_list}"
+                    )
                     mem_ids.extend(mem_id_list)
 
                 # submit messages for scheduler

src/memos/mem_os/product.py

@@ -23,7 +23,11 @@
     remove_embedding_recursive,
     sort_children_by_memory_type,
 )
-from memos.mem_scheduler.schemas import ANSWER_LABEL, QUERY_LABEL, ScheduleMessageItem
+from memos.mem_scheduler.schemas.general_schemas import (
+    ANSWER_LABEL,
+    QUERY_LABEL,
+)
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 from memos.mem_user.persistent_user_manager import PersistentUserManager
 from memos.mem_user.user_manager import UserRole
 from memos.memories.textual.item import (
@@ -601,7 +605,7 @@ def user_register(
                 try:
                     default_mem_cube.dump(mem_cube_name_or_path)
                 except Exception as e:
-                    print(e)
+                    logger.error(f"Failed to dump default cube: {e}")
 
             # Register the default cube with MOS
             self.register_mem_cube(
@@ -679,57 +683,6 @@ def get_suggestion_query(self, user_id: str, language: str = "zh") -> list[str]:
         response_json = json.loads(clean_response)
         return response_json["query"]
 
-    def chat(
-        self,
-        query: str,
-        user_id: str,
-        cube_id: str | None = None,
-        history: MessageList | None = None,
-    ) -> Generator[str, None, None]:
-        """Chat with LLM SSE Type.
-        Args:
-            query (str): Query string.
-            user_id (str): User ID.
-            cube_id (str, optional): Custom cube ID for user.
-            history (list[dict], optional): Chat history.
-
-        Returns:
-            Generator[str, None, None]: The response string generator.
-        """
-        # Use MOSCore's built-in validation
-        if cube_id:
-            self._validate_cube_access(user_id, cube_id)
-        else:
-            self._validate_user_exists(user_id)
-
-        # Load user cubes if not already loaded
-        self._load_user_cubes(user_id, self.default_cube_config)
-        time_start = time.time()
-        memories_list = super().search(query, user_id)["text_mem"]
-        # Get response from parent MOSCore (returns string, not generator)
-        response = super().chat(query, user_id)
-        time_end = time.time()
-
-        # Use tiktoken for proper token-based chunking
-        for chunk in self._chunk_response_with_tiktoken(response, chunk_size=5):
-            chunk_data = f"data: {json.dumps({'type': 'text', 'content': chunk})}\n\n"
-            yield chunk_data
-
-        # Prepare reference data
-        reference = []
-        for memories in memories_list:
-            memories_json = memories.model_dump()
-            memories_json["metadata"]["ref_id"] = f"[{memories.id.split('-')[0]}]"
-            memories_json["metadata"]["embedding"] = []
-            memories_json["metadata"]["sources"] = []
-            reference.append(memories_json)
-
-        yield f"data: {json.dumps({'type': 'reference', 'content': reference})}\n\n"
-        total_time = round(float(time_end - time_start), 1)
-
-        yield f"data: {json.dumps({'type': 'time', 'content': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n"
-        yield f"data: {json.dumps({'type': 'end'})}\n\n"
-
     def chat_with_references(
         self,
         query: str,
@@ -768,6 +721,8 @@ def chat_with_references(
             self._register_chat_history(user_id)
 
         chat_history = self.chat_history_manager[user_id]
+        if history:
+            chat_history.chat_history = history[-10:]
         current_messages = [
             {"role": "system", "content": system_prompt},
             *chat_history.chat_history,
@@ -853,15 +808,12 @@ def chat_with_references(
         yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n"
         total_time = round(float(time_end - time_start), 1)
         yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n"
-        chat_history.chat_history.append({"role": "user", "content": query})
-        chat_history.chat_history.append({"role": "assistant", "content": full_response})
         self._send_message_to_scheduler(
             user_id=user_id, mem_cube_id=cube_id, query=query, label=QUERY_LABEL
         )
         self._send_message_to_scheduler(
             user_id=user_id, mem_cube_id=cube_id, query=full_response, label=ANSWER_LABEL
         )
-        self.chat_history_manager[user_id] = chat_history
 
         yield f"data: {json.dumps({'type': 'end'})}\n\n"
         self.add(
@@ -880,12 +832,6 @@ def chat_with_references(
             ],
             mem_cube_id=cube_id,
         )
-        # Keep chat history under 30 messages by removing oldest conversation pair
-        if len(self.chat_history_manager[user_id].chat_history) > 10:
-            self.chat_history_manager[user_id].chat_history.pop(0)  # Remove oldest user message
-            self.chat_history_manager[user_id].chat_history.pop(
-                0
-            )  # Remove oldest assistant response
 
     def get_all(
         self,
@@ -1030,11 +976,9 @@ def get_subgraph(
         return reformat_memory_list
 
     def search(
-        self, query: str, user_id: str, install_cube_ids: list[str] | None = None, top_k: int = 20
+        self, query: str, user_id: str, install_cube_ids: list[str] | None = None, top_k: int = 10
     ):
         """Search memories for a specific user."""
-        # Validate user access
-        self._validate_user_access(user_id)
 
         # Load user cubes if not already loaded
         self._load_user_cubes(user_id, self.default_cube_config)

src/memos/mem_scheduler/base_scheduler.py

@@ -194,10 +194,10 @@ def replace_working_memory(
             text_mem_base: TreeTextMemory = text_mem_base
 
             # process rerank memories with llm
-            quey_history = self.monitor.query_monitors.get_queries_with_timesort()
+            query_history = self.monitor.query_monitors.get_queries_with_timesort()
             memories_with_new_order, rerank_success_flag = (
                 self.retriever.process_and_rerank_memories(
-                    queries=quey_history,
+                    queries=query_history,
                     original_memory=original_memory,
                     new_memory=new_memory,
                     top_k=self.top_k,
@@ -350,54 +350,63 @@ def update_activation_memory_periodically(
     ):
         new_activation_memories = []
 
-        if self.monitor.timed_trigger(
-            last_time=self.monitor.last_activation_mem_update_time,
-            interval_seconds=interval_seconds,
-        ):
-            logger.info(f"Updating activation memory for user {user_id} and mem_cube {mem_cube_id}")
+        try:
+            if self.monitor.timed_trigger(
+                last_time=self.monitor.last_activation_mem_update_time,
+                interval_seconds=interval_seconds,
+            ):
+                logger.info(
+                    f"Updating activation memory for user {user_id} and mem_cube {mem_cube_id}"
+                )
+
+                if (
+                    user_id not in self.monitor.working_memory_monitors
+                    or mem_cube_id not in self.monitor.working_memory_monitors[user_id]
+                    or len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0
+                ):
+                    logger.warning(
+                        "No memories found in working_memory_monitors, initializing from current working_memories"
+                    )
+                    self.initialize_working_memory_monitors(
+                        user_id=user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=mem_cube,
+                    )
 
-            if len(self.monitor.working_memory_monitors[user_id][mem_cube_id].memories) == 0:
-                logger.warning(
-                    "No memories found in working_memory_monitors, initializing from current working_memories"
+                self.monitor.update_activation_memory_monitors(
+                    user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
                 )
-                self.initialize_working_memory_monitors(
+
+                new_activation_memories = [
+                    m.memory_text
+                    for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id].memories
+                ]
+
+                logger.info(
+                    f"Collected {len(new_activation_memories)} new memory entries for processing"
+                )
+
+                self.update_activation_memory(
+                    new_memories=new_activation_memories,
+                    label=label,
                     user_id=user_id,
                     mem_cube_id=mem_cube_id,
                     mem_cube=mem_cube,
                 )
 
-            self.monitor.update_activation_memory_monitors(
-                user_id=user_id, mem_cube_id=mem_cube_id, mem_cube=mem_cube
-            )
-
-            new_activation_memories = [
-                m.memory_text
-                for m in self.monitor.activation_memory_monitors[user_id][mem_cube_id].memories
-            ]
-
-            logger.info(
-                f"Collected {len(new_activation_memories)} new memory entries for processing"
-            )
-
-            self.update_activation_memory(
-                new_memories=new_activation_memories,
-                label=label,
-                user_id=user_id,
-                mem_cube_id=mem_cube_id,
-                mem_cube=mem_cube,
-            )
-
-            self.monitor.last_activation_mem_update_time = datetime.now()
+                self.monitor.last_activation_mem_update_time = datetime.now()
 
-            logger.debug(
-                f"Activation memory update completed at {self.monitor.last_activation_mem_update_time}"
-            )
-        else:
-            logger.info(
-                f"Skipping update - {interval_seconds} second interval not yet reached. "
-                f"Last update time is {self.monitor.last_activation_mem_update_time} and now is"
-                f"{datetime.now()}"
-            )
+                logger.debug(
+                    f"Activation memory update completed at {self.monitor.last_activation_mem_update_time}"
+                )
+            else:
+                logger.info(
+                    f"Skipping update - {interval_seconds} second interval not yet reached. "
+                    f"Last update time is {self.monitor.last_activation_mem_update_time} and now is"
+                    f"{datetime.now()}"
+                )
+        except Exception as e:
+            logger.error(f"Error: {e}", exc_info=True)
 
     def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
         """Submit multiple messages to the message queue."""

src/memos/mem_scheduler/general_scheduler.py

@@ -145,43 +145,45 @@ def _add_message_consumer(self, messages: list[ScheduleMessageItem]) -> None:
         grouped_messages = self.dispatcher.group_messages_by_user_and_cube(messages=messages)
 
         self.validate_schedule_messages(messages=messages, label=ADD_LABEL)
-
-        for user_id in grouped_messages:
-            for mem_cube_id in grouped_messages[user_id]:
-                messages = grouped_messages[user_id][mem_cube_id]
-                if len(messages) == 0:
-                    return
-
-                # for status update
-                self._set_current_context_from_message(msg=messages[0])
-
-                # submit logs
-                for msg in messages:
-                    userinput_memory_ids = json.loads(msg.content)
-                    mem_cube = msg.mem_cube
-                    for memory_id in userinput_memory_ids:
-                        mem_item: TextualMemoryItem = mem_cube.text_mem.get(memory_id=memory_id)
-                        mem_type = mem_item.meta_data.memory_type
-                        mem_content = mem_item.memory
-
-                        self.log_adding_memory(
-                            memory=mem_content,
-                            memory_type=mem_type,
-                            user_id=msg.user_id,
-                            mem_cube_id=msg.mem_cube_id,
-                            mem_cube=msg.mem_cube,
-                            log_func_callback=self._submit_web_logs,
+        try:
+            for user_id in grouped_messages:
+                for mem_cube_id in grouped_messages[user_id]:
+                    messages = grouped_messages[user_id][mem_cube_id]
+                    if len(messages) == 0:
+                        return
+
+                    # for status update
+                    self._set_current_context_from_message(msg=messages[0])
+
+                    # submit logs
+                    for msg in messages:
+                        userinput_memory_ids = json.loads(msg.content)
+                        mem_cube = msg.mem_cube
+                        for memory_id in userinput_memory_ids:
+                            mem_item: TextualMemoryItem = mem_cube.text_mem.get(memory_id=memory_id)
+                            mem_type = mem_item.metadata.memory_type
+                            mem_content = mem_item.memory
+
+                            self.log_adding_memory(
+                                memory=mem_content,
+                                memory_type=mem_type,
+                                user_id=msg.user_id,
+                                mem_cube_id=msg.mem_cube_id,
+                                mem_cube=msg.mem_cube,
+                                log_func_callback=self._submit_web_logs,
+                            )
+
+                    # update activation memories
+                    if self.enable_act_memory_update:
+                        self.update_activation_memory_periodically(
+                            interval_seconds=self.monitor.act_mem_update_interval,
+                            label=ADD_LABEL,
+                            user_id=user_id,
+                            mem_cube_id=mem_cube_id,
+                            mem_cube=messages[0].mem_cube,
                         )
-
-                # update activation memories
-                if self.enable_act_memory_update:
-                    self.update_activation_memory_periodically(
-                        interval_seconds=self.monitor.act_mem_update_interval,
-                        label=ADD_LABEL,
-                        user_id=user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=messages[0].mem_cube,
-                    )
+        except Exception as e:
+            logger.error(f"Error: {e}", exc_info=True)
 
     def process_session_turn(
         self,