feat & refactor: enable mem scheduler to load auth config from environment variables, refactor AuthConfig and EnvConfigMixin for improved robustness and smarter configuration handling, and allow the mem scheduler to initialize modules with RabbitMQ support.

Author: tangg555

evaluation/.env-example

@@ -9,3 +9,24 @@ ZEP_API_KEY="z_***REDACTED***"
 CHAT_MODEL="gpt-4o-mini"
 CHAT_MODEL_BASE_URL="http://***.***.***.***:3000/v1"
 CHAT_MODEL_API_KEY="sk-***REDACTED***"
+
+# Configuration Only For Scheduler
+# RabbitMQ Configuration
+MEMSCHEDULER_RABBITMQ_HOST_NAME=rabbitmq-cn-***.cn-***.amqp-32.net.mq.amqp.aliyuncs.com
+MEMSCHEDULER_RABBITMQ_USER_NAME=***
+MEMSCHEDULER_RABBITMQ_PASSWORD=***
+MEMSCHEDULER_RABBITMQ_VIRTUAL_HOST=memos
+MEMSCHEDULER_RABBITMQ_ERASE_ON_CONNECT=true
+MEMSCHEDULER_RABBITMQ_PORT=5672
+
+# OpenAI Configuration
+MEMSCHEDULER_OPENAI_API_KEY=sk-***
+MEMSCHEDULER_OPENAI_BASE_URL=http://***.***.***.***:3000/v1
+MEMSCHEDULER_OPENAI_DEFAULT_MODEL=gpt-4o-mini
+
+# Graph DB Configuration
+MEMSCHEDULER_GRAPHDBAUTH_URI=bolt://localhost:7687
+MEMSCHEDULER_GRAPHDBAUTH_USER=neo4j
+MEMSCHEDULER_GRAPHDBAUTH_PASSWORD=***
+MEMSCHEDULER_GRAPHDBAUTH_DB_NAME=neo4j
+MEMSCHEDULER_GRAPHDBAUTH_AUTO_CREATE=true

evaluation/scripts/temporal_locomo/modules/base_eval_module.py

@@ -91,51 +91,53 @@ def __init__(self, args):
         self.ingestion_storage_dir = self.result_dir / "storages"
         self.mos_config_path = Path(f"{BASE_DIR}/configs-example/mos_w_scheduler_config.json")
         self.mem_cube_config_path = Path(f"{BASE_DIR}/configs-example/mem_cube_config.json")
+
         self.openai_api_key = os.getenv("CHAT_MODEL_API_KEY")
         self.openai_base_url = os.getenv("CHAT_MODEL_BASE_URL")
         self.openai_chat_model = os.getenv("CHAT_MODEL")
 
         auth_config_path = Path(f"{BASE_DIR}/scripts/temporal_locomo/eval_auth.json")
         if auth_config_path.exists():
             auth_config = AuthConfig.from_local_config(config_path=auth_config_path)
-
-            self.openai_api_key = auth_config.openai.api_key
-            self.openai_base_url = auth_config.openai.base_url
-            self.openai_chat_model = auth_config.openai.default_model
-
-            self.mos_config_data = json.load(self.mos_config_path.open("r", encoding="utf-8"))
-            self.mem_cube_config_data = json.load(
-                self.mem_cube_config_path.open("r", encoding="utf-8")
-            )
-
-            # Update LLM authentication information in MOS configuration using dictionary assignment
-            self.mos_config_data["mem_reader"]["config"]["llm"]["config"]["api_key"] = (
-                auth_config.openai.api_key
-            )
-            self.mos_config_data["mem_reader"]["config"]["llm"]["config"]["api_base"] = (
-                auth_config.openai.base_url
-            )
-
-            # Update graph database authentication information in memory cube configuration using dictionary assignment
-            self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["uri"] = (
-                auth_config.graph_db.uri
-            )
-            self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["user"] = (
-                auth_config.graph_db.user
-            )
-            self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["password"] = (
-                auth_config.graph_db.password
-            )
-            self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["db_name"] = (
-                auth_config.graph_db.db_name
+            print(
+                f"✅ Configuration loaded successfully: from local config file {auth_config_path}"
             )
-            self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["auto_create"] = (
-                auth_config.graph_db.auto_create
-            )
-
         else:
-            print("Please referring to configs-example to provide valid configs.")
-            exit()
+            # Load .env file first before reading environment variables
+            load_dotenv()
+            auth_config = AuthConfig.from_local_env()
+            print("✅ Configuration loaded successfully: from environment variables")
+        self.openai_api_key = auth_config.openai.api_key
+        self.openai_base_url = auth_config.openai.base_url
+        self.openai_chat_model = auth_config.openai.default_model
+
+        self.mos_config_data = json.load(self.mos_config_path.open("r", encoding="utf-8"))
+        self.mem_cube_config_data = json.load(self.mem_cube_config_path.open("r", encoding="utf-8"))
+
+        # Update LLM authentication information in MOS configuration using dictionary assignment
+        self.mos_config_data["mem_reader"]["config"]["llm"]["config"]["api_key"] = (
+            auth_config.openai.api_key
+        )
+        self.mos_config_data["mem_reader"]["config"]["llm"]["config"]["api_base"] = (
+            auth_config.openai.base_url
+        )
+
+        # Update graph database authentication information in memory cube configuration using dictionary assignment
+        self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["uri"] = (
+            auth_config.graph_db.uri
+        )
+        self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["user"] = (
+            auth_config.graph_db.user
+        )
+        self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["password"] = (
+            auth_config.graph_db.password
+        )
+        self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["db_name"] = (
+            auth_config.graph_db.db_name
+        )
+        self.mem_cube_config_data["text_mem"]["config"]["graph_db"]["config"]["auto_create"] = (
+            auth_config.graph_db.auto_create
+        )
 
         # Logger initialization
         self.logger = logger
@@ -158,7 +160,6 @@ def __init__(self, args):
 
         self.can_answer_cases: list[RecordingCase] = []
         self.cannot_answer_cases: list[RecordingCase] = []
-        load_dotenv()
 
     def print_eval_info(self):
         """

src/memos/configs/mem_scheduler.py

@@ -1,3 +1,4 @@
+import logging
 import os
 
 from pathlib import Path
@@ -135,7 +136,7 @@ class GraphDBAuthConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
     password: str = Field(
         default="",
         description="Password for graph database authentication",
-        min_length=8,  # 建议密码最小长度
+        min_length=8,  # Recommended minimum password length
     )
     db_name: str = Field(default="neo4j", description="Database name to connect to")
     auto_create: bool = Field(
@@ -150,13 +151,51 @@ class OpenAIConfig(BaseConfig, DictConversionMixin, EnvConfigMixin):
 
 
 class AuthConfig(BaseConfig, DictConversionMixin):
-    rabbitmq: RabbitMQConfig
-    openai: OpenAIConfig
-    graph_db: GraphDBAuthConfig
+    rabbitmq: RabbitMQConfig | None = None
+    openai: OpenAIConfig | None = None
+    graph_db: GraphDBAuthConfig | None = None
     default_config_path: ClassVar[str] = (
         f"{BASE_DIR}/examples/data/config/mem_scheduler/scheduler_auth.yaml"
     )
 
+    @model_validator(mode="after")
+    def validate_partial_initialization(self) -> "AuthConfig":
+        """
+        Validate that at least one configuration component is successfully initialized.
+        Log warnings for any failed initializations but allow partial success.
+        """
+        logger = logging.getLogger(__name__)
+
+        initialized_components = []
+        failed_components = []
+
+        if self.rabbitmq is not None:
+            initialized_components.append("rabbitmq")
+        else:
+            failed_components.append("rabbitmq")
+
+        if self.openai is not None:
+            initialized_components.append("openai")
+        else:
+            failed_components.append("openai")
+
+        if self.graph_db is not None:
+            initialized_components.append("graph_db")
+        else:
+            failed_components.append("graph_db")
+
+        # Allow all components to be None for flexibility, but log a warning
+        if not initialized_components:
+            logger.warning(
+                "All configuration components are None. This may indicate missing environment variables or configuration files."
+            )
+        elif failed_components:
+            logger.warning(
+                f"Failed to initialize components: {', '.join(failed_components)}. Successfully initialized: {', '.join(initialized_components)}"
+            )
+
+        return self
+
     @classmethod
     def from_local_config(cls, config_path: str | Path | None = None) -> "AuthConfig":
         """
@@ -205,24 +244,75 @@ def from_local_env(cls) -> "AuthConfig":
 
         This method loads configuration for all nested components (RabbitMQ, OpenAI, GraphDB)
         from their respective environment variables using each component's specific prefix.
+        If any component fails to initialize, it will be set to None and a warning will be logged.
 
         Returns:
             AuthConfig: Configured instance with values from environment variables
 
         Raises:
-            ValueError: If any required environment variables are missing
+            ValueError: If all components fail to initialize
         """
+        logger = logging.getLogger(__name__)
+
+        rabbitmq_config = None
+        openai_config = None
+        graph_db_config = None
+
+        # Try to initialize RabbitMQ config - check if any RabbitMQ env vars exist
+        try:
+            rabbitmq_prefix = RabbitMQConfig.get_env_prefix()
+            has_rabbitmq_env = any(key.startswith(rabbitmq_prefix) for key in os.environ)
+            if has_rabbitmq_env:
+                rabbitmq_config = RabbitMQConfig.from_env()
+                logger.info("Successfully initialized RabbitMQ configuration")
+            else:
+                logger.info(
+                    "No RabbitMQ environment variables found, skipping RabbitMQ initialization"
+                )
+        except (ValueError, Exception) as e:
+            logger.warning(f"Failed to initialize RabbitMQ config from environment: {e}")
+
+        # Try to initialize OpenAI config - check if any OpenAI env vars exist
+        try:
+            openai_prefix = OpenAIConfig.get_env_prefix()
+            has_openai_env = any(key.startswith(openai_prefix) for key in os.environ)
+            if has_openai_env:
+                openai_config = OpenAIConfig.from_env()
+                logger.info("Successfully initialized OpenAI configuration")
+            else:
+                logger.info("No OpenAI environment variables found, skipping OpenAI initialization")
+        except (ValueError, Exception) as e:
+            logger.warning(f"Failed to initialize OpenAI config from environment: {e}")
+
+        # Try to initialize GraphDB config - check if any GraphDB env vars exist
+        try:
+            graphdb_prefix = GraphDBAuthConfig.get_env_prefix()
+            has_graphdb_env = any(key.startswith(graphdb_prefix) for key in os.environ)
+            if has_graphdb_env:
+                graph_db_config = GraphDBAuthConfig.from_env()
+                logger.info("Successfully initialized GraphDB configuration")
+            else:
+                logger.info(
+                    "No GraphDB environment variables found, skipping GraphDB initialization"
+                )
+        except (ValueError, Exception) as e:
+            logger.warning(f"Failed to initialize GraphDB config from environment: {e}")
+
         return cls(
-            rabbitmq=RabbitMQConfig.from_env(),
-            openai=OpenAIConfig.from_env(),
-            graph_db=GraphDBAuthConfig.from_env(),
+            rabbitmq=rabbitmq_config,
+            openai=openai_config,
+            graph_db=graph_db_config,
         )
 
     def set_openai_config_to_environment(self):
-        # Set environment variables
-        os.environ["OPENAI_API_KEY"] = self.openai.api_key
-        os.environ["OPENAI_BASE_URL"] = self.openai.base_url
-        os.environ["MODEL"] = self.openai.default_model
+        # Set environment variables only if openai config is available
+        if self.openai is not None:
+            os.environ["OPENAI_API_KEY"] = self.openai.api_key
+            os.environ["OPENAI_BASE_URL"] = self.openai.base_url
+            os.environ["MODEL"] = self.openai.default_model
+        else:
+            logger = logging.getLogger(__name__)
+            logger.warning("OpenAI config is not available, skipping environment variable setup")
 
     @classmethod
     def default_config_exists(cls) -> bool:

src/memos/mem_os/core.py

@@ -124,11 +124,6 @@ def _initialize_mem_scheduler(self) -> GeneralScheduler:
                     f"Memory reader of type {type(self.mem_reader).__name__} "
                     "missing required 'llm' attribute"
                 )
-                self._mem_scheduler.initialize_modules(
-                    chat_llm=self.chat_llm,
-                    process_llm=self.chat_llm,
-                    db_engine=self.user_manager.engine,
-                )
             else:
                 # Configure scheduler general_modules
                 self._mem_scheduler.initialize_modules(

src/memos/mem_scheduler/base_scheduler.py

@@ -138,7 +138,8 @@ def initialize_modules(
 
             if self.auth_config is not None:
                 self.rabbitmq_config = self.auth_config.rabbitmq
-                self.initialize_rabbitmq(config=self.rabbitmq_config)
+                if self.rabbitmq_config is not None:
+                    self.initialize_rabbitmq(config=self.rabbitmq_config)
 
             logger.debug("GeneralScheduler has been initialized")
         except Exception as e:
@@ -497,6 +498,9 @@ def _submit_web_logs(
         Args:
             messages: Single log message or list of log messages
         """
+        if self.rabbitmq_config is None:
+            return
+
         if isinstance(messages, ScheduleLogForWebItem):
             messages = [messages]  # transform single message to list
 
@@ -526,7 +530,7 @@ def get_web_log_messages(self) -> list[dict]:
         messages = []
         while True:
             try:
-                item = self._web_log_message_queue.get_nowait()  # 线程安全的 get
+                item = self._web_log_message_queue.get_nowait()  # Thread-safe get
                 messages.append(item.to_dict())
             except queue.Empty:
                 break

src/memos/mem_scheduler/general_modules/dispatcher.py

@@ -206,7 +206,7 @@ def join(self, timeout: float | None = None) -> bool:
             bool: True if all tasks completed, False if timeout occurred.
         """
         if not self.enable_parallel_dispatch or self.dispatcher_executor is None:
-            return True  # 串行模式无需等待
+            return True  # Serial mode requires no waiting
 
         done, not_done = concurrent.futures.wait(
             self._futures, timeout=timeout, return_when=concurrent.futures.ALL_COMPLETED

src/memos/mem_scheduler/general_modules/misc.py

@@ -6,6 +6,7 @@
 from queue import Empty, Full, Queue
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
+from dotenv import load_dotenv
 from pydantic import field_serializer
 
 
@@ -32,7 +33,7 @@ def get_env_prefix(cls) -> str:
         Examples:
             RabbitMQConfig -> "RABBITMQ_"
             OpenAIConfig -> "OPENAI_"
-            GraphDBAuthConfig -> "GRAPH_DB_AUTH_"
+            GraphDBAuthConfig -> "GRAPHDBAUTH_"
         """
         class_name = cls.__name__
         # Remove 'Config' suffix if present
@@ -55,6 +56,8 @@ def from_env(cls: type[T]) -> T:
         Raises:
             ValueError: If required environment variables are missing.
         """
+        load_dotenv()
+
         prefix = cls.get_env_prefix()
         field_values = {}
 
@@ -85,6 +88,35 @@ def _parse_env_value(cls, value: str, target_type: type) -> Any:
             return float(value)
         return value
 
+    @classmethod
+    def print_env_mapping(cls) -> None:
+        """Print the mapping between class fields and their corresponding environment variable names.
+
+        Displays each field's name, type, whether it's required, default value, and corresponding environment variable name.
+        """
+        prefix = cls.get_env_prefix()
+        print(f"\n=== {cls.__name__} Environment Variable Mapping ===")
+        print(f"Environment Variable Prefix: {prefix}")
+        print("-" * 60)
+
+        if not hasattr(cls, "model_fields"):
+            print("This class does not define model_fields, may not be a Pydantic model")
+            return
+
+        for field_name, field_info in cls.model_fields.items():
+            env_var = f"{prefix}{field_name.upper()}"
+            field_type = field_info.annotation
+            is_required = field_info.is_required()
+            default_value = field_info.default if field_info.default is not None else "None"
+
+            print(f"Field Name: {field_name}")
+            print(f"  Environment Variable: {env_var}")
+            print(f"  Type: {field_type}")
+            print(f"  Required: {'Yes' if is_required else 'No'}")
+            print(f"  Default Value: {default_value}")
+            print(f"  Current Environment Value: {os.environ.get(env_var, 'Not Set')}")
+            print("-" * 40)
+
 
 class DictConversionMixin:
     """

src/memos/mem_scheduler/general_modules/scheduler_logger.py

@@ -98,7 +98,7 @@ def create_autofilled_log_item(
         )
         return log_message
 
-    # TODO: 日志打出来数量不对
+    # TODO: Log output count is incorrect
     @log_exceptions(logger=logger)
     def log_working_memory_replacement(
         self,