new feat & fix bugs: mem_scheduler update the logging feature and examples; some bugs due to refactoring have been fixed

Author: tangg555

examples/mem_scheduler/memos_w_scheduler.py

@@ -3,13 +3,16 @@
 
 from datetime import datetime
 from pathlib import Path
+from queue import Queue
+from typing import TYPE_CHECKING
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
 from memos.configs.mem_scheduler import AuthConfig, SchedulerConfigFactory
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 from memos.mem_os.main import MOS
+from memos.mem_scheduler.general_scheduler import GeneralScheduler
 from memos.mem_scheduler.scheduler_factory import SchedulerFactory
 from memos.mem_scheduler.schemas.general_schemas import (
     ANSWER_LABEL,
@@ -19,6 +22,12 @@
 from memos.mem_scheduler.utils.misc_utils import parse_yaml
 
 
+if TYPE_CHECKING:
+    from memos.mem_scheduler.schemas import (
+        ScheduleLogForWebItem,
+    )
+
+
 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
@@ -109,6 +118,8 @@ def run_with_automatic_scheduler_init():
         response = mos.chat(query, user_id=user_id)
         print(f"Query:\n {query}\n\nAnswer:\n {response}")
 
+    show_web_logs(mem_scheduler=mos.mem_scheduler)
+
     mos.mem_scheduler.stop()
 
 
@@ -184,9 +195,46 @@ def run_with_manual_scheduler_init():
         mos.mem_scheduler.submit_messages(messages=message_item)
         print(f"Query:\n {query}\n\nAnswer:\n {response}")
 
+    show_web_logs(mem_scheduler=mos.mem_scheduler)
+
     mos.mem_scheduler.stop()
 
 
+def show_web_logs(mem_scheduler: GeneralScheduler):
+    """Display all web log entries from the scheduler's log queue.
+
+    Args:
+        mem_scheduler: The scheduler instance containing web logs to display
+    """
+    if mem_scheduler._web_log_message_queue.empty():
+        print("Web log queue is currently empty.")
+        return
+
+    print("\n" + "=" * 50 + " WEB LOGS " + "=" * 50)
+
+    # Create a temporary queue to preserve the original queue contents
+    temp_queue = Queue()
+    log_count = 0
+
+    while not mem_scheduler._web_log_message_queue.empty():
+        log_item: ScheduleLogForWebItem = mem_scheduler._web_log_message_queue.get()
+        temp_queue.put(log_item)
+        log_count += 1
+
+        # Print log entry details
+        print(f"\nLog Entry #{log_count}:")
+        print(f'- "{log_item.label}" log: {log_item}')
+
+        print("-" * 50)
+
+    # Restore items back to the original queue
+    while not temp_queue.empty():
+        mem_scheduler._web_log_message_queue.put(temp_queue.get())
+
+    print(f"\nTotal {log_count} web log entries displayed.")
+    print("=" * 110 + "\n")
+
+
 if __name__ == "__main__":
     run_with_automatic_scheduler_init()
 

examples/mem_scheduler/memos_w_scheduler_for_test.py

@@ -1,25 +1,17 @@
+import json
 import shutil
 import sys
 
 from pathlib import Path
-from queue import Queue
-from typing import TYPE_CHECKING
 
 from memos.configs.mem_cube import GeneralMemCubeConfig
 from memos.configs.mem_os import MOSConfig
 from memos.configs.mem_scheduler import AuthConfig
 from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
-from memos.mem_scheduler.general_scheduler import GeneralScheduler
 from memos.mem_scheduler.mos_for_test_scheduler import MOSForTestScheduler
 
 
-if TYPE_CHECKING:
-    from memos.mem_scheduler.schemas import (
-        ScheduleLogForWebItem,
-    )
-
-
 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
@@ -90,41 +82,6 @@ def init_task():
     return conversations, questions
 
 
-def show_web_logs(mem_scheduler: GeneralScheduler):
-    """Display all web log entries from the scheduler's log queue.
-
-    Args:
-        mem_scheduler: The scheduler instance containing web logs to display
-    """
-    if mem_scheduler._web_log_message_queue.empty():
-        print("Web log queue is currently empty.")
-        return
-
-    print("\n" + "=" * 50 + " WEB LOGS " + "=" * 50)
-
-    # Create a temporary queue to preserve the original queue contents
-    temp_queue = Queue()
-    log_count = 0
-
-    while not mem_scheduler._web_log_message_queue.empty():
-        log_item: ScheduleLogForWebItem = mem_scheduler._web_log_message_queue.get()
-        temp_queue.put(log_item)
-        log_count += 1
-
-        # Print log entry details
-        print(f"\nLog Entry #{log_count}:")
-        print(f'- "{log_item.label}" log: {log_item}')
-
-        print("-" * 50)
-
-    # Restore items back to the original queue
-    while not temp_queue.empty():
-        mem_scheduler._web_log_message_queue.put(temp_queue.get())
-
-    print(f"\nTotal {log_count} web log entries displayed.")
-    print("=" * 110 + "\n")
-
-
 if __name__ == "__main__":
     # set up data
     conversations, questions = init_task()
@@ -168,13 +125,18 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
 
     mos.add(conversations, user_id=user_id, mem_cube_id=mem_cube_id)
 
+    # Add interfering conversations
+    file_path = Path(f"{BASE_DIR}/examples/data/mem_scheduler/scene_data.json")
+    scene_data = json.load(file_path.open("r", encoding="utf-8"))
+    mos.add(scene_data[0], user_id=user_id, mem_cube_id=mem_cube_id)
+    mos.add(scene_data[1], user_id=user_id, mem_cube_id=mem_cube_id)
+
     for item in questions:
         print("===== Chat Start =====")
         query = item["question"]
         print(f"Query:\n {query}\n")
         response = mos.chat(query=query, user_id=user_id)
         print(f"Answer:\n {response}")
         print("===== Chat End =====")
-    show_web_logs(mos.mem_scheduler)
 
     mos.mem_scheduler.stop()

examples/mem_scheduler/try_schedule_modules.py

@@ -1,7 +1,6 @@
 import shutil
 import sys
 
-from datetime import datetime
 from pathlib import Path
 from queue import Queue
 from typing import TYPE_CHECKING
@@ -17,9 +16,7 @@
 from memos.mem_scheduler.mos_for_test_scheduler import MOSForTestScheduler
 from memos.mem_scheduler.schemas.general_schemas import (
     NOT_APPLICABLE_TYPE,
-    QUERY_LABEL,
 )
-from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 if TYPE_CHECKING:
@@ -189,23 +186,6 @@ def show_web_logs(mem_scheduler: GeneralScheduler):
         )
         print(f"\nnew_candidates: {[one.memory for one in new_candidates]}")
 
-        # test query_consume
-        message_item = ScheduleMessageItem(
-            user_id=user_id,
-            mem_cube_id=mem_cube_id,
-            mem_cube=mem_cube,
-            label=QUERY_LABEL,
-            content=query,
-            timestamp=datetime.now(),
-        )
-        print(
-            f"\nworking memories before: {[one.memory for one in mos.mem_scheduler.mem_cubes[mem_cube_id].text_mem.get_working_memory()]}"
-        )
-        mos.mem_scheduler._query_message_consumer(messages=[message_item])
-        print(
-            f"\nworking memories after: {[one.memory for one in mos.mem_scheduler.mem_cubes[mem_cube_id].text_mem.get_working_memory()]}"
-        )
-
         # test activation memory update
         mos.mem_scheduler.update_activation_memory_periodically(
             interval_seconds=0,

src/memos/configs/mem_scheduler.py

@@ -106,7 +106,20 @@ class RabbitMQConfig(
 
 
 class GraphDBAuthConfig(BaseConfig):
-    uri: str = Field(default="localhost", description="URI for graph database access")
+    uri: str = Field(
+        default="bolt://localhost:7687",
+        description="URI for graph database access (e.g., bolt://host:port)",
+    )
+    user: str = Field(default="neo4j", description="Username for graph database authentication")
+    password: str = Field(
+        default="",
+        description="Password for graph database authentication",
+        min_length=8,  # 建议密码最小长度
+    )
+    db_name: str = Field(default="neo4j", description="Database name to connect to")
+    auto_create: bool = Field(
+        default=True, description="Whether to automatically create the database if it doesn't exist"
+    )
 
 
 class OpenAIConfig(BaseConfig):

src/memos/mem_os/core.py

@@ -586,6 +586,7 @@ def add(
             user_id (str, optional): The identifier of the user to add the memories to.
                 If None, the default user is used.
         """
+        # user input messages
         assert (messages is not None) or (memory_content is not None) or (doc_path is not None), (
             "messages_or_doc_path or memory_content or doc_path must be provided."
         )
@@ -625,23 +626,26 @@ def add(
                     type="chat",
                     info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
                 )
+
+                mem_ids = []
                 for mem in memories:
-                    self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_ids.extend(mem_id_list)
 
                 # submit messages for scheduler
-                mem_cube = self.mem_cubes[mem_cube_id]
                 if self.enable_mem_scheduler and self.mem_scheduler is not None:
-                    text_messages = [message["content"] for message in messages]
+                    mem_cube = self.mem_cubes[mem_cube_id]
                     message_item = ScheduleMessageItem(
                         user_id=target_user_id,
                         mem_cube_id=mem_cube_id,
                         mem_cube=mem_cube,
                         label=ADD_LABEL,
-                        content=json.dumps(text_messages),
+                        content=json.dumps(mem_ids),
                         timestamp=datetime.now(),
                     )
                     self.mem_scheduler.submit_messages(messages=[message_item])
 
+        # user profile
         if (
             (memory_content is not None)
             and self.config.enable_textual_memory
@@ -669,21 +673,56 @@ def add(
                     type="chat",
                     info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
                 )
+
+                mem_ids = []
                 for mem in memories:
-                    self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                    mem_ids.extend(mem_id_list)
+
+                # submit messages for scheduler
+                if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                    mem_cube = self.mem_cubes[mem_cube_id]
+                    message_item = ScheduleMessageItem(
+                        user_id=target_user_id,
+                        mem_cube_id=mem_cube_id,
+                        mem_cube=mem_cube,
+                        label=ADD_LABEL,
+                        content=json.dumps(mem_ids),
+                        timestamp=datetime.now(),
+                    )
+                    self.mem_scheduler.submit_messages(messages=[message_item])
+
+        # user doc input
         if (
             (doc_path is not None)
             and self.config.enable_textual_memory
             and self.mem_cubes[mem_cube_id].text_mem
         ):
             documents = self._get_all_documents(doc_path)
-            doc_memory = self.mem_reader.get_memory(
+            doc_memories = self.mem_reader.get_memory(
                 documents,
                 type="doc",
                 info={"user_id": target_user_id, "session_id": str(uuid.uuid4())},
             )
-            for mem in doc_memory:
-                self.mem_cubes[mem_cube_id].text_mem.add(mem)
+
+            mem_ids = []
+            for mem in doc_memories:
+                mem_id_list: list[str] = self.mem_cubes[mem_cube_id].text_mem.add(mem)
+                mem_ids.extend(mem_id_list)
+
+            # submit messages for scheduler
+            if self.enable_mem_scheduler and self.mem_scheduler is not None:
+                mem_cube = self.mem_cubes[mem_cube_id]
+                message_item = ScheduleMessageItem(
+                    user_id=target_user_id,
+                    mem_cube_id=mem_cube_id,
+                    mem_cube=mem_cube,
+                    label=ADD_LABEL,
+                    content=json.dumps(mem_ids),
+                    timestamp=datetime.now(),
+                )
+                self.mem_scheduler.submit_messages(messages=[message_item])
+
         logger.info(f"Add memory to {mem_cube_id} successfully")
 
     def get(

src/memos/mem_scheduler/base_scheduler.py

@@ -351,7 +351,7 @@ def update_activation_memory_periodically(
         new_activation_memories = []
 
         if self.monitor.timed_trigger(
-            last_time=self.monitor._last_activation_mem_update_time,
+            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}")
@@ -387,15 +387,15 @@ def update_activation_memory_periodically(
                 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}"
+                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"Last update time is {self.monitor.last_activation_mem_update_time} and now is"
                 f"{datetime.now()}"
             )
 
@@ -421,10 +421,11 @@ def _submit_web_logs(
 
         for message in messages:
             self._web_log_message_queue.put(message)
-            logger.info(f"Submitted Scheduling log for web: {message.log_content}")
+            message_info = message.debug_info()
+            logger.debug(f"Submitted Scheduling log for web: {message_info}")
 
             if self.is_rabbitmq_connected():
-                logger.info("Submitted Scheduling log to rabbitmq")
+                logger.info(f"Submitted Scheduling log to rabbitmq: {message_info}")
                 self.rabbitmq_publish_message(message=message.to_dict())
         logger.debug(f"{len(messages)} submitted. {self._web_log_message_queue.qsize()} in queue.")