MemoryOS 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

memos/mem_scheduler/modules/scheduler_logger.py

@@ -0,0 +1,255 @@
+from collections.abc import Callable
+
+from memos.log import get_logger
+from memos.mem_cube.general import GeneralMemCube
+from memos.mem_scheduler.modules.base import BaseSchedulerModule
+from memos.mem_scheduler.schemas.general_schemas import (
+    ACTIVATION_MEMORY_TYPE,
+    ADD_LABEL,
+    LONG_TERM_MEMORY_TYPE,
+    NOT_INITIALIZED,
+    PARAMETER_MEMORY_TYPE,
+    QUERY_LABEL,
+    TEXT_MEMORY_TYPE,
+    USER_INPUT_TYPE,
+    WORKING_MEMORY_TYPE,
+)
+from memos.mem_scheduler.schemas.message_schemas import (
+    ScheduleLogForWebItem,
+    ScheduleMessageItem,
+)
+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:
+        text_mem_base: TreeTextMemory = mem_cube.text_mem
+        current_memory_sizes = text_mem_base.get_current_memory_size()
+        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": text_mem_base.memory_manager.memory_size["LongTermMemory"],
+            "user_memory_capacity": text_mem_base.memory_manager.memory_size["UserMemory"],
+            "working_memory_capacity": text_mem_base.memory_manager.memory_size["WorkingMemory"],
+            "transformed_act_memory_capacity": NOT_INITIALIZED,
+            "parameter_memory_capacity": NOT_INITIALIZED,
+        }
+
+        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.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 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."""
+        memory_type_map = {
+            transform_name_to_key(name=m.memory): m.metadata.memory_type
+            for m in original_memory + new_memory
+        }
+
+        original_text_memories = [m.memory for m in original_memory]
+        new_text_memories = [m.memory for m in new_memory]
+
+        # Convert to sets for efficient difference operations
+        original_set = set(original_text_memories)
+        new_set = set(new_text_memories)
+
+        # Identify changes
+        added_memories = list(new_set - original_set)  # Present in new but not original
+
+        # recording messages
+        for memory in added_memories:
+            normalized_mem = transform_name_to_key(name=memory)
+            if normalized_mem not in memory_type_map:
+                logger.error(f"Memory text not found in type mapping: {memory[:50]}...")
+            # Get the memory type from the map, default to LONG_TERM_MEMORY_TYPE if not found
+            mem_type = memory_type_map.get(normalized_mem, LONG_TERM_MEMORY_TYPE)
+
+            if mem_type == WORKING_MEMORY_TYPE:
+                logger.warning(f"Memory already in working memory: {memory[:50]}...")
+                continue
+
+            log_message = self.create_autofilled_log_item(
+                log_content=memory,
+                label=QUERY_LABEL,
+                from_memory_type=mem_type,
+                to_memory_type=WORKING_MEMORY_TYPE,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+            )
+            log_func_callback([log_message])
+            logger.info(
+                f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
+                f"transformed to {WORKING_MEMORY_TYPE} memories."
+            )
+
+    @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)
+
+        # Identify changes
+        added_memories = list(new_set - original_set)  # Present in new but not original
+
+        # recording messages
+        for mem in added_memories:
+            log_message_a = self.create_autofilled_log_item(
+                log_content=mem,
+                label=label,
+                from_memory_type=TEXT_MEMORY_TYPE,
+                to_memory_type=ACTIVATION_MEMORY_TYPE,
+                user_id=user_id,
+                mem_cube_id=mem_cube_id,
+                mem_cube=mem_cube,
+            )
+            log_message_b = self.create_autofilled_log_item(
+                log_content=mem,
+                label=label,
+                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,
+            )
+            log_func_callback([log_message_a, log_message_b])
+            logger.info(
+                f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
+                f"transformed to {WORKING_MEMORY_TYPE} memories."
+            )
+
+    @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],
+    ):
+        """Log changes when working memory is replaced."""
+        log_message = self.create_autofilled_log_item(
+            log_content=memory,
+            label=ADD_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 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

memos/mem_scheduler/mos_for_test_scheduler.py

@@ -3,12 +3,12 @@ from datetime import datetime
 from memos.configs.mem_os import MOSConfig
 from memos.log import get_logger
 from memos.mem_os.main import MOS
-from memos.mem_scheduler.modules.schemas import (
+from memos.mem_scheduler.schemas.general_schemas import (
     ANSWER_LABEL,
     MONITOR_WORKING_MEMORY_TYPE,
     QUERY_LABEL,
-    ScheduleMessageItem,
 )
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
 
 
 logger = get_logger(__name__)
@@ -54,24 +54,20 @@ class MOSForTestScheduler(MOS):
                 if not mem_cube.text_mem:
                     continue
 
-                # submit message to scheduler
-                if self.enable_mem_scheduler and self.mem_scheduler is not None:
-                    message_item = ScheduleMessageItem(
-                        user_id=target_user_id,
-                        mem_cube_id=mem_cube_id,
-                        mem_cube=mem_cube,
-                        label=QUERY_LABEL,
-                        content=query,
-                        timestamp=datetime.now(),
-                    )
-                    self.mem_scheduler.submit_messages(messages=[message_item])
-
-                self.mem_scheduler.monitor.register_memory_manager_if_not_exists(
-                    user_id=user_id,
+                message_item = ScheduleMessageItem(
+                    user_id=target_user_id,
                     mem_cube_id=mem_cube_id,
-                    memory_monitors=self.mem_scheduler.monitor.working_memory_monitors,
-                    max_capacity=self.mem_scheduler.monitor.working_mem_monitor_capacity,
+                    mem_cube=mem_cube,
+                    label=QUERY_LABEL,
+                    content=query,
+                    timestamp=datetime.now(),
                 )
+                cur_working_memories = [m.memory for m in mem_cube.text_mem.get_working_memory()]
+                print(f"Working memories before schedule: {cur_working_memories}")
+
+                # --- force to run mem_scheduler ---
+                self.mem_scheduler.monitor.query_trigger_interval = 0
+                self.mem_scheduler._query_message_consumer(messages=[message_item])
 
                 # from scheduler
                 scheduler_memories = self.mem_scheduler.monitor.get_monitor_memories(
@@ -80,6 +76,7 @@ class MOSForTestScheduler(MOS):
                     memory_type=MONITOR_WORKING_MEMORY_TYPE,
                     top_k=topk_for_scheduler,
                 )
+                print(f"Working memories after schedule: {scheduler_memories}")
                 memories_all.extend(scheduler_memories)
 
                 # from mem_cube
@@ -87,6 +84,7 @@ class MOSForTestScheduler(MOS):
                     query, top_k=self.config.top_k - topk_for_scheduler
                 )
                 text_memories = [m.memory for m in memories]
+                print(f"Search results with new working memories: {text_memories}")
                 memories_all.extend(text_memories)
 
                 memories_all = list(set(memories_all))
@@ -139,5 +137,4 @@ class MOSForTestScheduler(MOS):
                     timestamp=datetime.now(),
                 )
                 self.mem_scheduler.submit_messages(messages=[message_item])
-
         return response

memos/mem_scheduler/schemas/general_schemas.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+from typing import NewType
+
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+QUERY_LABEL = "query"
+ANSWER_LABEL = "answer"
+ADD_LABEL = "add"
+
+TreeTextMemory_SEARCH_METHOD = "tree_text_memory_search"
+TextMemory_SEARCH_METHOD = "text_memory_search"
+DIRECT_EXCHANGE_TYPE = "direct"
+FANOUT_EXCHANGE_TYPE = "fanout"
+DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT = 20
+DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT = 5
+DEFAULT_ACT_MEM_DUMP_PATH = f"{BASE_DIR}/outputs/mem_scheduler/mem_cube_scheduler_test.kv_cache"
+DEFAULT_THREAD__POOL_MAX_WORKERS = 5
+DEFAULT_CONSUME_INTERVAL_SECONDS = 3
+NOT_INITIALIZED = -1
+
+
+# web log
+LONG_TERM_MEMORY_TYPE = "LongTermMemory"
+USER_MEMORY_TYPE = "UserMemory"
+WORKING_MEMORY_TYPE = "WorkingMemory"
+TEXT_MEMORY_TYPE = "TextMemory"
+ACTIVATION_MEMORY_TYPE = "ActivationMemory"
+PARAMETER_MEMORY_TYPE = "ParameterMemory"
+USER_INPUT_TYPE = "UserInput"
+NOT_APPLICABLE_TYPE = "NotApplicable"
+
+# monitors
+MONITOR_WORKING_MEMORY_TYPE = "MonitorWorkingMemoryType"
+MONITOR_ACTIVATION_MEMORY_TYPE = "MonitorActivationMemoryType"
+DEFAULT_MAX_QUERY_KEY_WORDS = 1000
+DEFAULT_WEIGHT_VECTOR_FOR_RANKING = [0.9, 0.05, 0.05]
+
+
+# new types
+UserID = NewType("UserID", str)
+MemCubeID = NewType("CubeID", str)

memos/mem_scheduler/schemas/message_schemas.py

@@ -0,0 +1,148 @@
+from datetime import datetime
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer
+from typing_extensions import TypedDict
+
+from memos.log import get_logger
+from memos.mem_cube.general import GeneralMemCube
+from memos.mem_scheduler.modules.misc import DictConversionMixin
+
+from .general_schemas import NOT_INITIALIZED
+
+
+logger = get_logger(__name__)
+
+DEFAULT_MEMORY_SIZES = {
+    "long_term_memory_size": NOT_INITIALIZED,
+    "user_memory_size": NOT_INITIALIZED,
+    "working_memory_size": NOT_INITIALIZED,
+    "transformed_act_memory_size": NOT_INITIALIZED,
+    "parameter_memory_size": NOT_INITIALIZED,
+}
+
+DEFAULT_MEMORY_CAPACITIES = {
+    "long_term_memory_capacity": 10000,
+    "user_memory_capacity": 10000,
+    "working_memory_capacity": 20,
+    "transformed_act_memory_capacity": NOT_INITIALIZED,
+    "parameter_memory_capacity": NOT_INITIALIZED,
+}
+
+
+class ScheduleMessageItem(BaseModel, DictConversionMixin):
+    item_id: str = Field(description="uuid", default_factory=lambda: str(uuid4()))
+    user_id: str = Field(..., description="user id")
+    mem_cube_id: str = Field(..., description="memcube id")
+    label: str = Field(..., description="Label of the schedule message")
+    mem_cube: GeneralMemCube | str = Field(..., description="memcube for schedule")
+    content: str = Field(..., description="Content of the schedule message")
+    timestamp: datetime = Field(
+        default_factory=lambda: datetime.utcnow(), description="submit time for schedule_messages"
+    )
+
+    # Pydantic V2 model configuration
+    model_config = ConfigDict(
+        # Allows arbitrary Python types as model fields without validation
+        # Required when using custom types like GeneralMemCube that aren't Pydantic models
+        arbitrary_types_allowed=True,
+        # Additional metadata for JSON Schema generation
+        json_schema_extra={
+            # Example payload demonstrating the expected structure and sample values
+            # Used for API documentation, testing, and developer reference
+            "example": {
+                "item_id": "123e4567-e89b-12d3-a456-426614174000",  # Sample UUID
+                "user_id": "user123",  # Example user identifier
+                "mem_cube_id": "cube456",  # Sample memory cube ID
+                "label": "sample_label",  # Demonstration label value
+                "mem_cube": "obj of GeneralMemCube",  # Added mem_cube example
+                "content": "sample content",  # Example message content
+                "timestamp": "2024-07-22T12:00:00Z",  # Added timestamp example
+            }
+        },
+    )
+
+    @field_serializer("mem_cube")
+    def serialize_mem_cube(self, cube: GeneralMemCube | str, _info) -> str:
+        """Custom serializer for GeneralMemCube objects to string representation"""
+        if isinstance(cube, str):
+            return cube
+        return f"<GeneralMemCube:{id(cube)}>"
+
+    def to_dict(self) -> dict:
+        """Convert model to dictionary suitable for Redis Stream"""
+        return {
+            "item_id": self.item_id,
+            "user_id": self.user_id,
+            "cube_id": self.mem_cube_id,
+            "label": self.label,
+            "cube": "Not Applicable",  # Custom cube serialization
+            "content": self.content,
+            "timestamp": self.timestamp.isoformat(),
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ScheduleMessageItem":
+        """Create model from Redis Stream dictionary"""
+        return cls(
+            item_id=data.get("item_id", str(uuid4())),
+            user_id=data["user_id"],
+            cube_id=data["cube_id"],
+            label=data["label"],
+            cube="Not Applicable",  # Custom cube deserialization
+            content=data["content"],
+            timestamp=datetime.fromisoformat(data["timestamp"]),
+        )
+
+
+class MemorySizes(TypedDict):
+    long_term_memory_size: int
+    user_memory_size: int
+    working_memory_size: int
+    transformed_act_memory_size: int
+
+
+class MemoryCapacities(TypedDict):
+    long_term_memory_capacity: int
+    user_memory_capacity: int
+    working_memory_capacity: int
+    transformed_act_memory_capacity: int
+
+
+class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
+    item_id: str = Field(
+        description="Unique identifier for the log entry", default_factory=lambda: str(uuid4())
+    )
+    user_id: str = Field(..., description="Identifier for the user associated with the log")
+    mem_cube_id: str = Field(
+        ..., description="Identifier for the memcube associated with this log entry"
+    )
+    label: str = Field(..., description="Label categorizing the type of log")
+    from_memory_type: str = Field(..., description="Source memory type")
+    to_memory_type: str = Field(..., description="Destination memory type")
+    log_content: str = Field(..., description="Detailed content of the log entry")
+    current_memory_sizes: MemorySizes = Field(
+        default_factory=lambda: dict(DEFAULT_MEMORY_SIZES),
+        description="Current utilization of memory partitions",
+    )
+    memory_capacities: MemoryCapacities = Field(
+        default_factory=lambda: dict(DEFAULT_MEMORY_CAPACITIES),
+        description="Maximum capacities of memory partitions",
+    )
+    timestamp: datetime = Field(
+        default_factory=lambda: datetime.utcnow(),
+        description="Timestamp indicating when the log entry was created",
+    )
+
+    def debug_info(self) -> dict[str, Any]:
+        """Return structured debug information for logging purposes."""
+        return {
+            "log_id": self.item_id,
+            "user_id": self.user_id,
+            "mem_cube_id": self.mem_cube_id,
+            "operation": f"{self.from_memory_type} → {self.to_memory_type}",
+            "label": self.label,
+            "content_length": len(self.log_content),
+            "timestamp": self.timestamp.isoformat(),
+        }