MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/general_modules/misc.py

@@ -0,0 +1,313 @@
+import json
+import os
+
+from contextlib import suppress
+from datetime import datetime
+from queue import Empty, Full, Queue
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+from dotenv import load_dotenv
+from pydantic import field_serializer
+
+
+if TYPE_CHECKING:
+    from pydantic import BaseModel
+
+T = TypeVar("T")
+
+BaseModelType = TypeVar("T", bound="BaseModel")
+
+
+class EnvConfigMixin(Generic[T]):
+    """Abstract base class for environment variable configuration."""
+
+    ENV_PREFIX = "MEMSCHEDULER_"
+
+    @classmethod
+    def get_env_prefix(cls) -> str:
+        """Automatically generates environment variable prefix from class name.
+
+        Converts the class name to uppercase and appends an underscore.
+        If the class name ends with 'Config', that suffix is removed first.
+
+        Examples:
+            RabbitMQConfig -> "RABBITMQ_"
+            OpenAIConfig -> "OPENAI_"
+            GraphDBAuthConfig -> "GRAPHDBAUTH_"
+        """
+        class_name = cls.__name__
+        # Remove 'Config' suffix if present
+        if class_name.endswith("Config"):
+            class_name = class_name[:-6]
+        # Convert to uppercase and add trailing underscore
+
+        return f"{cls.ENV_PREFIX}{class_name.upper()}_"
+
+    @classmethod
+    def from_env(cls: type[T]) -> T:
+        """Creates a config instance from environment variables.
+
+        Reads all environment variables with the class-specific prefix and maps them
+        to corresponding configuration fields (converting to the appropriate types).
+
+        Returns:
+            An instance of the config class populated from environment variables.
+
+        Raises:
+            ValueError: If required environment variables are missing.
+        """
+        load_dotenv()
+
+        prefix = cls.get_env_prefix()
+        field_values = {}
+
+        for field_name, field_info in cls.model_fields.items():
+            env_var = f"{prefix}{field_name.upper()}"
+            field_type = field_info.annotation
+
+            if field_info.is_required() and env_var not in os.environ:
+                raise ValueError(f"Required environment variable {env_var} is missing")
+
+            if env_var in os.environ:
+                raw_value = os.environ[env_var]
+                field_values[field_name] = cls._parse_env_value(raw_value, field_type)
+            elif field_info.default is not None:
+                field_values[field_name] = field_info.default
+            else:
+                raise ValueError()
+        return cls(**field_values)
+
+    @classmethod
+    def _parse_env_value(cls, value: str, target_type: type) -> Any:
+        """Converts environment variable string to appropriate type."""
+        if target_type is bool:
+            return value.lower() in ("true", "1", "t", "y", "yes")
+        if target_type is int:
+            return int(value)
+        if target_type is float:
+            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:
+    """
+    Provides conversion functionality between Pydantic models and dictionaries,
+    including datetime serialization handling.
+    """
+
+    @field_serializer("timestamp", check_fields=False)
+    def serialize_datetime(self, dt: datetime | None, _info) -> str | None:
+        """
+        Custom timestamp serialization logic.
+        - Supports timezone-aware datetime objects
+        - Compatible with models without timestamp field (via check_fields=False)
+        """
+        if dt is None:
+            return None
+        return dt.isoformat()
+
+    def to_dict(self) -> dict:
+        """
+        Convert model instance to dictionary.
+        - Uses model_dump to ensure field consistency
+        - Prioritizes custom serializer for timestamp handling
+        """
+        dump_data = self.model_dump()
+        if hasattr(self, "timestamp") and self.timestamp is not None:
+            dump_data["timestamp"] = self.serialize_datetime(self.timestamp, None)
+        return dump_data
+
+    def to_json(self, **kwargs) -> str:
+        """
+        Convert model instance to a JSON string.
+        - Accepts the same kwargs as json.dumps (e.g., indent, ensure_ascii)
+        - Default settings make JSON human-readable and UTF-8 safe
+        """
+        return json.dumps(self.to_dict(), ensure_ascii=False, default=lambda o: str(o), **kwargs)
+
+    @classmethod
+    def from_json(cls: type[BaseModelType], json_str: str) -> BaseModelType:
+        """
+        Create model instance from a JSON string.
+        - Parses JSON into a dictionary and delegates to from_dict
+        """
+        try:
+            data = json.loads(json_str)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON string: {e}") from e
+        return cls.from_dict(data)
+
+    @classmethod
+    def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
+        """
+        Create model instance from dictionary.
+        - Automatically converts timestamp strings to datetime objects
+        """
+        data_copy = data.copy()  # Avoid modifying original dictionary
+        if "timestamp" in data_copy and isinstance(data_copy["timestamp"], str):
+            try:
+                data_copy["timestamp"] = datetime.fromisoformat(data_copy["timestamp"])
+            except ValueError:
+                # Handle invalid time formats - adjust as needed (e.g., log warning or set to None)
+                data_copy["timestamp"] = None
+
+        return cls(**data_copy)
+
+    def __str__(self) -> str:
+        """
+        Convert to formatted JSON string.
+        - Used for user-friendly display in print() or str() calls
+        """
+        return json.dumps(
+            self.to_dict(),
+            indent=4,
+            ensure_ascii=False,
+            default=lambda o: str(o),  # Handle other non-serializable objects
+        )
+
+
+class AutoDroppingQueue(Queue[T]):
+    """A thread-safe queue that automatically drops the oldest item when full."""
+
+    def __init__(self, maxsize: int = 0):
+        # If maxsize <= 0, set to 0 (unlimited queue size)
+        if maxsize <= 0:
+            maxsize = 0
+        super().__init__(maxsize=maxsize)
+
+    def put(self, item: T, block: bool = False, timeout: float | None = None) -> None:
+        """Put an item into the queue.
+
+        If the queue is full, the oldest item will be automatically removed to make space.
+        IMPORTANT: When we drop an item we also call `task_done()` to keep
+        the internal `unfinished_tasks` counter consistent (the dropped task
+        will never be processed).
+
+        Args:
+            item: The item to be put into the queue
+            block: Ignored (kept for compatibility with Queue interface)
+            timeout: Ignored (kept for compatibility with Queue interface)
+        """
+        while True:
+            try:
+                # First try non-blocking put
+                super().put(item, block=block, timeout=timeout)
+                return
+            except Full:
+                # Remove the oldest item and mark it done to avoid leaking unfinished_tasks
+                with suppress(Empty):
+                    _ = self.get_nowait()
+                    # If the removed item had previously incremented unfinished_tasks,
+                    # we must decrement here since it will never be processed.
+                    with suppress(ValueError):
+                        self.task_done()
+                # Continue loop to retry putting the item
+
+    def get(
+        self, block: bool = True, timeout: float | None = None, batch_size: int | None = None
+    ) -> list[T]:
+        """Get items from the queue.
+
+        Args:
+            block: Whether to block if no items are available (default: True)
+            timeout: Timeout in seconds for blocking operations (default: None)
+            batch_size: Number of items to retrieve (default: 1)
+
+        Returns:
+            List of items (always returns a list for consistency)
+
+        Raises:
+            Empty: If no items are available and block=False or timeout expires
+        """
+
+        if batch_size is None:
+            return super().get(block=block, timeout=timeout)
+        items = []
+        for _ in range(batch_size):
+            try:
+                items.append(super().get(block=block, timeout=timeout))
+            except Empty:
+                if not items and block:
+                    # If we haven't gotten any items and we're blocking, re-raise Empty
+                    raise
+                break
+        return items
+
+    def get_nowait(self, batch_size: int | None = None) -> list[T]:
+        """Get items from the queue without blocking.
+
+        Args:
+            batch_size: Number of items to retrieve (default: 1)
+
+        Returns:
+            List of items (always returns a list for consistency)
+        """
+        if batch_size is None:
+            return super().get_nowait()
+
+        items = []
+        for _ in range(batch_size):
+            try:
+                items.append(super().get_nowait())
+            except Empty:
+                break
+        return items
+
+    def get_queue_content_without_pop(self) -> list[T]:
+        """Return a copy of the queue's contents without modifying it."""
+        # Ensure a consistent snapshot by holding the mutex
+        with self.mutex:
+            return list(self.queue)
+
+    def qsize(self) -> int:
+        """Return the approximate size of the queue.
+
+        Returns:
+            Number of items currently in the queue
+        """
+        return super().qsize()
+
+    def clear(self) -> None:
+        """Remove all items from the queue.
+
+        This operation is thread-safe.
+        IMPORTANT: We also decrement `unfinished_tasks` by the number of
+        items cleared, since those tasks will never be processed.
+        """
+        with self.mutex:
+            dropped = len(self.queue)
+            self.queue.clear()
+        # Call task_done() outside of the mutex to avoid deadlocks because
+        # Queue.task_done() acquires the same condition bound to `self.mutex`.
+        for _ in range(dropped):
+            with suppress(ValueError):
+                self.task_done()

memos/mem_scheduler/general_modules/scheduler_logger.py

@@ -0,0 +1,389 @@
+import hashlib
+
+from collections.abc import Callable
+
+from memos.log import get_logger
+from memos.mem_cube.general import GeneralMemCube
+from memos.mem_scheduler.general_modules.base import BaseSchedulerModule
+from memos.mem_scheduler.schemas.general_schemas import (
+    ACTIVATION_MEMORY_TYPE,
+    NOT_INITIALIZED,
+    PARAMETER_MEMORY_TYPE,
+    TEXT_MEMORY_TYPE,
+    WORKING_MEMORY_TYPE,
+)
+from memos.mem_scheduler.schemas.message_schemas import (
+    ScheduleLogForWebItem,
+    ScheduleMessageItem,
+)
+from memos.mem_scheduler.schemas.task_schemas import (
+    ADD_TASK_LABEL,
+    MEM_ARCHIVE_TASK_LABEL,
+    MEM_UPDATE_TASK_LABEL,
+    USER_INPUT_TYPE,
+)
+from memos.mem_scheduler.utils.filter_utils import (
+    transform_name_to_key,
+)
+from memos.mem_scheduler.utils.misc_utils import log_exceptions
+from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
+
+
+logger = get_logger(__name__)
+
+
+class SchedulerLoggerModule(BaseSchedulerModule):
+    def __init__(self):
+        """
+        Initialize RabbitMQ connection settings.
+        """
+        super().__init__()
+
+    @log_exceptions(logger=logger)
+    def create_autofilled_log_item(
+        self,
+        log_content: str,
+        label: str,
+        from_memory_type: str,
+        to_memory_type: str,
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+    ) -> ScheduleLogForWebItem:
+        if mem_cube is None:
+            logger.error(
+                "mem_cube is None — this should not happen in production!", stack_info=True
+            )
+        text_mem_base: TreeTextMemory = mem_cube.text_mem
+
+        current_memory_sizes = {}
+        if hasattr(text_mem_base, "get_current_memory_size"):
+            current_memory_sizes = text_mem_base.get_current_memory_size(user_name=mem_cube_id)
+
+        current_memory_sizes = {
+            "long_term_memory_size": current_memory_sizes.get("LongTermMemory", 0),
+            "user_memory_size": current_memory_sizes.get("UserMemory", 0),
+            "working_memory_size": current_memory_sizes.get("WorkingMemory", 0),
+            "transformed_act_memory_size": NOT_INITIALIZED,
+            "parameter_memory_size": NOT_INITIALIZED,
+        }
+
+        memory_capacities = {
+            "long_term_memory_capacity": 0,
+            "user_memory_capacity": 0,
+            "working_memory_capacity": 0,
+            "transformed_act_memory_capacity": NOT_INITIALIZED,
+            "parameter_memory_capacity": NOT_INITIALIZED,
+        }
+
+        if hasattr(text_mem_base, "memory_manager") and hasattr(
+            text_mem_base.memory_manager, "memory_size"
+        ):
+            memory_capacities.update(
+                {
+                    "long_term_memory_capacity": text_mem_base.memory_manager.memory_size.get(
+                        "LongTermMemory", 0
+                    ),
+                    "user_memory_capacity": text_mem_base.memory_manager.memory_size.get(
+                        "UserMemory", 0
+                    ),
+                    "working_memory_capacity": text_mem_base.memory_manager.memory_size.get(
+                        "WorkingMemory", 0
+                    ),
+                }
+            )
+
+        if hasattr(self, "monitor"):
+            if (
+                user_id in self.monitor.activation_memory_monitors
+                and mem_cube_id in self.monitor.activation_memory_monitors[user_id]
+            ):
+                activation_monitor = self.monitor.activation_memory_monitors[user_id][mem_cube_id]
+                transformed_act_memory_size = len(activation_monitor.obj.memories)
+                logger.info(
+                    f'activation_memory_monitors currently has "{transformed_act_memory_size}" transformed memory size'
+                )
+            else:
+                transformed_act_memory_size = 0
+                logger.info(
+                    f'activation_memory_monitors is not initialized for user "{user_id}" and mem_cube "{mem_cube_id}'
+                )
+            current_memory_sizes["transformed_act_memory_size"] = transformed_act_memory_size
+            current_memory_sizes["parameter_memory_size"] = 1
+
+            memory_capacities["transformed_act_memory_capacity"] = (
+                self.monitor.activation_mem_monitor_capacity
+            )
+            memory_capacities["parameter_memory_capacity"] = 1
+
+        log_message = ScheduleLogForWebItem(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            label=label,
+            from_memory_type=from_memory_type,
+            to_memory_type=to_memory_type,
+            log_content=log_content,
+            current_memory_sizes=current_memory_sizes,
+            memory_capacities=memory_capacities,
+        )
+        return log_message
+
+    @log_exceptions(logger=logger)
+    def create_event_log(
+        self,
+        label: str,
+        from_memory_type: str,
+        to_memory_type: str,
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+        memcube_log_content: list[dict],
+        metadata: list[dict],
+        memory_len: int,
+        memcube_name: str | None = None,
+        log_content: str | None = None,
+    ) -> ScheduleLogForWebItem:
+        item = self.create_autofilled_log_item(
+            log_content=log_content or "",
+            label=label,
+            from_memory_type=from_memory_type,
+            to_memory_type=to_memory_type,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )
+        item.memcube_log_content = memcube_log_content
+        item.metadata = metadata
+        item.memory_len = memory_len
+        item.memcube_name = memcube_name or self._map_memcube_name(mem_cube_id)
+        return item
+
+    def _map_memcube_name(self, mem_cube_id: str) -> str:
+        x = mem_cube_id or ""
+        if "public" in x.lower():
+            return "PublicMemCube"
+        return "UserMemCube"
+
+    # TODO: Log output count is incorrect
+    @log_exceptions(logger=logger)
+    def log_working_memory_replacement(
+        self,
+        original_memory: list[TextualMemoryItem],
+        new_memory: list[TextualMemoryItem],
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+        log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
+    ):
+        """Log changes when working memory is replaced."""
+        original_text_memories = [m.memory for m in original_memory]
+        new_text_memories = [m.memory for m in new_memory]
+        original_set = set(original_text_memories)
+        new_set = set(new_text_memories)
+        added_texts = []
+        for new_mem in new_set:
+            if new_mem not in original_set:
+                added_texts.append(new_mem)
+        memcube_content = []
+        meta = []
+        by_text = {m.memory: m for m in new_memory}
+        for t in added_texts:
+            itm = by_text.get(t)
+            if not itm:
+                continue
+            key_name = getattr(itm.metadata, "key", None) or itm.memory
+            k = transform_name_to_key(name=key_name)
+            memcube_content.append(
+                {
+                    "content": f"[{itm.metadata.memory_type}→{WORKING_MEMORY_TYPE}] {k}: {itm.memory}",
+                    "ref_id": itm.id,
+                }
+            )
+            meta.append(
+                {
+                    "ref_id": itm.id,
+                    "id": itm.id,
+                    "key": itm.metadata.key,
+                    "memory": itm.memory,
+                    "memory_type": itm.metadata.memory_type,
+                    "status": itm.metadata.status,
+                    "confidence": itm.metadata.confidence,
+                    "tags": itm.metadata.tags,
+                    "updated_at": getattr(itm.metadata, "updated_at", None)
+                    or getattr(itm.metadata, "update_at", None),
+                }
+            )
+        # Only create log if there are actual memory changes
+        if memcube_content:
+            ev = self.create_event_log(
+                label="scheduleMemory",
+                from_memory_type=TEXT_MEMORY_TYPE,
+                to_memory_type=WORKING_MEMORY_TYPE,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+                memcube_log_content=memcube_content,
+                metadata=meta,
+                memory_len=len(memcube_content),
+                memcube_name=self._map_memcube_name(mem_cube_id),
+            )
+            log_func_callback([ev])
+
+    @log_exceptions(logger=logger)
+    def log_activation_memory_update(
+        self,
+        original_text_memories: list[str],
+        new_text_memories: list[str],
+        label: str,
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+        log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
+    ):
+        """Log changes when activation memory is updated."""
+        original_set = set(original_text_memories)
+        new_set = set(new_text_memories)
+
+        added_memories = list(new_set - original_set)
+        memcube_content = []
+        meta = []
+        for mem in added_memories:
+            key = transform_name_to_key(mem)
+            ref_id = f"actparam-{hashlib.md5(mem.encode()).hexdigest()}"
+            memcube_content.append(
+                {
+                    "content": f"[{ACTIVATION_MEMORY_TYPE}→{PARAMETER_MEMORY_TYPE}] {key}: {mem}",
+                    "ref_id": ref_id,
+                }
+            )
+            meta.append(
+                {
+                    "ref_id": ref_id,
+                    "id": ref_id,
+                    "key": key,
+                    "memory": mem,
+                    "memory_type": ACTIVATION_MEMORY_TYPE,
+                    "status": None,
+                    "confidence": None,
+                    "tags": None,
+                    "updated_at": None,
+                }
+            )
+        # Only create log if there are actual memory changes
+        if memcube_content:
+            ev = self.create_event_log(
+                label="scheduleMemory",
+                from_memory_type=ACTIVATION_MEMORY_TYPE,
+                to_memory_type=PARAMETER_MEMORY_TYPE,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+                memcube_log_content=memcube_content,
+                metadata=meta,
+                memory_len=len(added_memories),
+                memcube_name=self._map_memcube_name(mem_cube_id),
+            )
+            log_func_callback([ev])
+
+    @log_exceptions(logger=logger)
+    def log_adding_memory(
+        self,
+        memory: str,
+        memory_type: str,
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+        log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
+    ):
+        """Deprecated: legacy text log. Use create_event_log with structured fields instead."""
+        log_message = self.create_autofilled_log_item(
+            log_content=memory,
+            label=ADD_TASK_LABEL,
+            from_memory_type=USER_INPUT_TYPE,
+            to_memory_type=memory_type,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )
+        log_func_callback([log_message])
+        logger.info(
+            f"{USER_INPUT_TYPE} memory for user {user_id} "
+            f"converted to {memory_type} memory in mem_cube {mem_cube_id}: {memory}"
+        )
+
+    @log_exceptions(logger=logger)
+    def log_updating_memory(
+        self,
+        memory: str,
+        memory_type: str,
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+        log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
+    ):
+        """Deprecated: legacy text log. Use create_event_log with structured fields instead."""
+        log_message = self.create_autofilled_log_item(
+            log_content=memory,
+            label=MEM_UPDATE_TASK_LABEL,
+            from_memory_type=memory_type,
+            to_memory_type=memory_type,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )
+        log_func_callback([log_message])
+
+    @log_exceptions(logger=logger)
+    def log_archiving_memory(
+        self,
+        memory: str,
+        memory_type: str,
+        user_id: str,
+        mem_cube_id: str,
+        mem_cube: GeneralMemCube,
+        log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
+    ):
+        """Deprecated: legacy text log. Use create_event_log with structured fields instead."""
+        log_message = self.create_autofilled_log_item(
+            log_content=memory,
+            label=MEM_ARCHIVE_TASK_LABEL,
+            from_memory_type=memory_type,
+            to_memory_type=memory_type,
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            mem_cube=mem_cube,
+        )
+        log_func_callback([log_message])
+
+    @log_exceptions(logger=logger)
+    def validate_schedule_message(self, message: ScheduleMessageItem, label: str):
+        """Validate if the message matches the expected label.
+
+        Args:
+            message: Incoming message item to validate.
+            label: Expected message label (e.g., QUERY_LABEL/ANSWER_LABEL).
+
+        Returns:
+            bool: True if validation passed, False otherwise.
+        """
+        if message.label != label:
+            logger.error(f"Handler validation failed: expected={label}, actual={message.label}")
+            return False
+        return True
+
+    @log_exceptions(logger=logger)
+    def validate_schedule_messages(self, messages: list[ScheduleMessageItem], label: str):
+        """Validate if all messages match the expected label.
+
+        Args:
+            messages: List of message items to validate.
+            label: Expected message label (e.g., QUERY_LABEL/ANSWER_LABEL).
+
+        Returns:
+            bool: True if all messages passed validation, False if any failed.
+        """
+        for message in messages:
+            if not self.validate_schedule_message(message, label):
+                logger.error("Message batch contains invalid labels, aborting processing")
+                return False
+        return True