MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/schemas/monitor_schemas.py

@@ -0,0 +1,406 @@
+import json
+import threading
+
+from collections import Counter
+from datetime import datetime
+from pathlib import Path
+from typing import ClassVar
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, computed_field, field_validator
+
+from memos.log import get_logger
+from memos.mem_scheduler.general_modules.misc import AutoDroppingQueue, DictConversionMixin
+from memos.mem_scheduler.schemas.general_schemas import (
+    DEFAULT_WEIGHT_VECTOR_FOR_RANKING,
+    NOT_INITIALIZED,
+)
+from memos.mem_scheduler.schemas.task_schemas import (
+    DEFAULT_MAX_QUERY_KEY_WORDS,
+)
+from memos.mem_scheduler.utils.filter_utils import transform_name_to_key
+from memos.memories.textual.tree import TextualMemoryItem
+
+
+logger = get_logger(__name__)
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+
+# ============== Queries ==============
+class QueryMonitorItem(BaseModel, DictConversionMixin):
+    item_id: str = Field(
+        description="Unique identifier for the query item", default_factory=lambda: str(uuid4())
+    )
+    user_id: str = Field(..., description="Required user identifier", min_length=1)
+    mem_cube_id: str = Field(..., description="Required memory cube identifier", min_length=1)
+    query_text: str = Field(
+        ...,
+        description="The actual user query text content",
+        min_length=1,
+    )
+    keywords: list[str] | None = Field(
+        default=None,
+        min_length=1,  # If provided, shouldn't be empty
+        description="Semantic keywords extracted from the query text",
+    )
+    max_keywords: ClassVar[int] = DEFAULT_MAX_QUERY_KEY_WORDS
+
+    timestamp: datetime = Field(
+        default_factory=datetime.now, description="Timestamp indicating when query was submitted"
+    )
+
+    @field_validator("keywords", mode="before")
+    @classmethod
+    def validate_keywords(cls, v, values):
+        if v is None:
+            return None
+
+        if not isinstance(v, list):
+            raise ValueError("Keywords must be a list")
+
+        if len(v) > cls.max_keywords:
+            logger.warning(
+                f"Keywords list truncated from {len(v)} to {cls.max_keywords} items. "
+                f"Configure max_keywords class attribute to adjust this limit."
+            )
+            return v[: cls.max_keywords]
+        return v
+
+    @classmethod
+    def with_max_keywords(cls, limit: int):
+        """Create a new class with custom keywords limit."""
+        if not isinstance(limit, int) or limit <= 0:
+            raise ValueError("Max keywords limit must be positive integer")
+
+        return type(f"{cls.__name__}_MaxKeywords{limit}", (cls,), {"max_keywords": limit})
+
+
+class QueryMonitorQueue(AutoDroppingQueue[QueryMonitorItem]):
+    """
+    A thread-safe queue for monitoring queries with timestamp and keyword tracking.
+    Each item is expected to be a dictionary containing:
+    """
+
+    def put(self, item: QueryMonitorItem, block: bool = True, timeout: float | None = 5.0) -> None:
+        """
+        Add a query item to the queue. Ensures the item is of correct type.
+
+        Args:
+            item: A QueryMonitorItem instance
+        """
+        if not isinstance(item, QueryMonitorItem):
+            raise ValueError("Item must be an instance of QueryMonitorItem")
+        logger.debug(
+            f"Thread {threading.get_ident()} acquired mutex. Timeout is set to {timeout} seconds"
+        )
+        super().put(item, block, timeout)
+
+    def get_queries_by_timestamp(
+        self, start_time: datetime, end_time: datetime
+    ) -> list[QueryMonitorItem]:
+        """
+        Retrieve queries added between the specified time range.
+        """
+        with self.mutex:
+            logger.debug(f"Thread {threading.get_ident()} acquired mutex.")
+            return [item for item in self.queue if start_time <= item.timestamp <= end_time]
+
+    def get_keywords_collections(self) -> Counter:
+        """
+        Generate a Counter containing keyword frequencies across all queries.
+
+        Returns:
+            Counter object with keyword counts
+        """
+        with self.mutex:
+            logger.debug(f"Thread {threading.get_ident()} acquired mutex.")
+            # Fix: Handle None keywords safely
+            all_keywords = [kw for item in self.queue if item.keywords for kw in item.keywords]
+            return Counter(all_keywords)
+
+    def get_queries_with_timesort(self, reverse: bool = True) -> list[str]:
+        """
+        Retrieve all queries sorted by timestamp.
+
+        Args:
+            reverse: If True, sort in descending order (newest first),
+                     otherwise sort in ascending order (oldest first)
+
+        Returns:
+            List of query items sorted by timestamp
+        """
+        with self.mutex:
+            logger.debug(f"Thread {threading.get_ident()} acquired mutex.")
+            return [
+                monitor.query_text
+                for monitor in sorted(self.queue, key=lambda x: x.timestamp, reverse=reverse)
+            ]
+
+    def to_json(self) -> str:
+        """Serialize the queue to a JSON string.
+
+        Args:
+            item_serializer: Optional function to serialize individual items.
+                             If not provided, items must be JSON-serializable.
+
+        Returns:
+            A JSON string representing the queue's content and maxsize.
+        """
+        with self.mutex:
+            serialized_items = [item.to_json() for item in self.queue]
+
+        data = {"maxsize": self.maxsize, "items": serialized_items}
+        return json.dumps(data, ensure_ascii=False, indent=2)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "QueryMonitorQueue":
+        """Create a new AutoDroppingQueue from a JSON string.
+
+        Args:
+            json_str: JSON string created by to_json()
+            item_deserializer: Optional function to reconstruct items from dicts.
+                               If not provided, items are used as-is.
+
+        Returns:
+            A new AutoDroppingQueue instance with deserialized data.
+        """
+        data = json.loads(json_str)
+        maxsize = data.get("maxsize", 0)
+        item_strs = data.get("items", [])
+
+        queue = cls(maxsize=maxsize)
+
+        items = [QueryMonitorItem.from_json(json_str=item_str) for item_str in item_strs]
+
+        # Fix: Add error handling for put operations
+        for item in items:
+            try:
+                queue.put(item)  # Use put() to respect maxsize and auto-drop behavior
+            except Exception as e:
+                logger.error(f"Failed to add item to queue: {e}")
+                # Continue with other items instead of failing completely
+
+        return queue
+
+
+# ============== Memories ==============
+class MemoryMonitorItem(BaseModel, DictConversionMixin):
+    """
+    Represents a memory item in the monitoring system.
+
+    Note: This class does NOT have a timestamp field, unlike QueryMonitorItem.
+    For sorting by recency, use sorting_score or importance_score instead.
+    """
+
+    item_id: str = Field(
+        description="Unique identifier for the memory item", default_factory=lambda: str(uuid4())
+    )
+    memory_text: str = Field(
+        ...,
+        description="The actual content of the memory",
+        min_length=1,
+    )
+    tree_memory_item: TextualMemoryItem | None = Field(
+        default=None, description="Optional textual memory item"
+    )
+    tree_memory_item_mapping_key: str = Field(
+        description="Key generated from memory_text using transform_name_to_key",
+    )
+    keywords_score: float = Field(
+        default=NOT_INITIALIZED,
+        description="The score generate by counting keywords in queries",
+        ge=NOT_INITIALIZED,  # Minimum value of 0
+    )
+    sorting_score: float = Field(
+        default=NOT_INITIALIZED,
+        description="The score generate from rerank process",
+        ge=NOT_INITIALIZED,  # Minimum value of 0
+    )
+    importance_score: float = Field(
+        default=NOT_INITIALIZED,
+        description="Numerical score representing the memory's importance",
+        ge=NOT_INITIALIZED,  # Minimum value of 0
+    )
+    recording_count: int = Field(
+        default=1,
+        description="How many times this memory has been recorded",
+        ge=1,
+    )
+
+    @field_validator("tree_memory_item_mapping_key", mode="before")
+    def generate_mapping_key(cls, v, values):  # noqa: N805
+        if v is None and "memory_text" in values:
+            return transform_name_to_key(values["memory_text"])
+        return v
+
+    def get_importance_score(self, weight_vector: list[float] | None = None) -> float:
+        return self._get_complex_importance_score(weight_vector=weight_vector)
+
+    def _get_complex_importance_score(self, weight_vector: list[float] | None = None) -> float:
+        """Calculate traditional importance score using existing logic"""
+        if weight_vector is None:
+            logger.warning("weight_vector of get_complex_score is None.")
+            weight_vector = DEFAULT_WEIGHT_VECTOR_FOR_RANKING
+
+        # Fix: Add proper validation for weight_vector
+        if not weight_vector or len(weight_vector) != 3 or abs(sum(weight_vector) - 1.0) > 1e-6:
+            raise ValueError("weight_vector must be provided, have length 3, and sum to 1.0")
+
+        # Fix: Handle uninitialized scores safely
+        sorting_score = self.sorting_score if self.sorting_score != NOT_INITIALIZED else 0.0
+        keywords_score = self.keywords_score if self.keywords_score != NOT_INITIALIZED else 0.0
+
+        normalized_keywords_score = min(keywords_score * weight_vector[1], 5)
+        normalized_recording_count_score = min(self.recording_count * weight_vector[2], 2)
+        self.importance_score = (
+            sorting_score * weight_vector[0]
+            + normalized_keywords_score * weight_vector[1]
+            + normalized_recording_count_score * weight_vector[2]
+        )
+        return self.importance_score
+
+
+class MemoryMonitorManager(BaseModel, DictConversionMixin):
+    user_id: str = Field(..., description="Required user identifier", min_length=1)
+    mem_cube_id: str = Field(..., description="Required memory cube identifier", min_length=1)
+    memories: list[MemoryMonitorItem] = Field(
+        default_factory=list, description="Collection of memory items"
+    )
+    max_capacity: int | None = Field(
+        default=None, description="Maximum number of memories allowed (None for unlimited)", ge=1
+    )
+
+    @computed_field
+    @property
+    def memory_size(self) -> int:
+        """Automatically calculated count of memory items."""
+        return len(self.memories)
+
+    @property
+    def memories_mapping_dict(self) -> dict[str, MemoryMonitorItem]:
+        """
+        Generate a mapping dictionary for the memories in MemoryMonitorManager,
+        using tree_memory_item_mapping_key as the key and MemoryMonitorItem as the value.
+
+        Returns:
+            Dict[str, MemoryMonitorItem]: A dictionary where keys are
+            tree_memory_item_mapping_key values from MemoryMonitorItem,
+            and values are the corresponding MemoryMonitorItem objects.
+        """
+        mapping_dict = {
+            mem_item.tree_memory_item_mapping_key: mem_item for mem_item in self.memories
+        }
+
+        logger.debug(
+            f"Generated memories mapping dict for user_id={self.user_id}, "
+            f"mem_cube_id={self.mem_cube_id}, "
+            f"total_items={len(mapping_dict)}, "
+            f"source_memory_count={len(self.memories)}"
+        )
+        return mapping_dict
+
+    def get_sorted_mem_monitors(self, reverse=True) -> list[MemoryMonitorItem]:
+        """
+        Retrieve memory monitors sorted by their ranking score in descending order.
+
+        Returns:
+            list[MemoryMonitorItem]: Sorted list of memory monitor items.
+        """
+        return sorted(
+            self.memories,
+            key=lambda item: item.get_importance_score(
+                weight_vector=DEFAULT_WEIGHT_VECTOR_FOR_RANKING
+            ),
+            reverse=reverse,
+        )
+
+    def update_memories(
+        self, new_memory_monitors: list[MemoryMonitorItem], partial_retention_number: int
+    ) -> list[MemoryMonitorItem]:  # Fix: Correct return type
+        """
+        Update memories based on monitor_working_memories.
+        """
+
+        # Validate partial_retention_number
+        if partial_retention_number < 0:
+            raise ValueError("partial_retention_number must be non-negative")
+
+        # Step 1: Update existing memories or add new ones
+        added_count = 0
+        memories_mapping_dict = self.memories_mapping_dict
+        new_mem_set = set()
+        for memory_monitor in new_memory_monitors:
+            if memory_monitor.tree_memory_item_mapping_key in memories_mapping_dict:
+                # Update existing memory
+                item: MemoryMonitorItem = memories_mapping_dict[
+                    memory_monitor.tree_memory_item_mapping_key
+                ]
+                item.recording_count += 1
+                item.keywords_score = memory_monitor.keywords_score
+                item.sorting_score = memory_monitor.sorting_score
+            else:
+                # Add new memory
+                self.memories.append(memory_monitor)
+                added_count += 1
+
+            new_mem_set.add(memory_monitor.tree_memory_item_mapping_key)
+
+        # Step 2: Identify memories to remove
+        old_mem_monitor_list = []
+        for mem_monitor in self.memories:
+            if mem_monitor.tree_memory_item_mapping_key not in new_mem_set:
+                old_mem_monitor_list.append(mem_monitor)
+
+        # Sort memories by recording_count in descending order
+        sorted_old_mem_monitors = sorted(
+            old_mem_monitor_list,
+            key=lambda item: item.get_importance_score(
+                weight_vector=DEFAULT_WEIGHT_VECTOR_FOR_RANKING
+            ),
+            reverse=True,
+        )
+
+        # Fix: Add bounds checking to prevent IndexError
+        if partial_retention_number > len(sorted_old_mem_monitors):
+            partial_retention_number = len(sorted_old_mem_monitors)
+            logger.info(
+                f"partial_retention_number adjusted to {partial_retention_number} to match available old memories"
+            )
+
+        # Keep the top N old memories
+        memories_to_remove = sorted_old_mem_monitors[partial_retention_number:]
+        memories_to_change_score = sorted_old_mem_monitors[:partial_retention_number]
+
+        # Step 3: Remove identified memories and change the scores of left old memories
+        for memory in memories_to_remove:
+            self.memories.remove(memory)
+
+        for memory in memories_to_change_score:
+            memory.sorting_score = 0
+            memory.recording_count = 1
+            memory.keywords_score = 0
+
+        # Step 4: Enforce max_capacity if set
+        # Fix: Handle max_capacity safely
+        if self.max_capacity is not None:
+            sorted_memories = sorted(
+                self.memories,
+                key=lambda item: item.get_importance_score(
+                    weight_vector=DEFAULT_WEIGHT_VECTOR_FOR_RANKING
+                ),
+                reverse=True,
+            )
+            # Keep only the top max_capacity memories
+            self.memories = sorted_memories[: self.max_capacity]
+
+        # Log the update result
+        logger.info(
+            f"Updated monitor manager for user {self.user_id}, mem_cube {self.mem_cube_id}: "
+            f"Total memories: {len(self.memories)}, "
+            f"Added/Updated: {added_count}, "
+            f"Removed: {len(memories_to_remove)} (excluding top {partial_retention_number} by recording_count)"
+        )
+
+        return self.memories

memos/mem_scheduler/schemas/task_schemas.py

@@ -0,0 +1,132 @@
+import os
+
+from datetime import datetime
+from enum import Enum
+from pathlib import Path
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, Field, computed_field
+
+from memos.log import get_logger
+from memos.mem_scheduler.general_modules.misc import DictConversionMixin
+from memos.mem_scheduler.utils.db_utils import get_utc_now
+
+
+logger = get_logger(__name__)
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+
+# ============== Schedule Task Definitaion ==============
+class TaskPriorityLevel(Enum):
+    # priority top
+    LEVEL_1 = 1
+    LEVEL_2 = 2
+    LEVEL_3 = 3
+    # priority bottom
+
+
+QUERY_TASK_LABEL = "query"
+ANSWER_TASK_LABEL = "answer"
+ADD_TASK_LABEL = "add"
+MEM_READ_TASK_LABEL = "mem_read"
+MEM_ORGANIZE_TASK_LABEL = "mem_organize"
+MEM_UPDATE_TASK_LABEL = "mem_update"
+MEM_ARCHIVE_TASK_LABEL = "mem_archive"
+API_MIX_SEARCH_TASK_LABEL = "api_mix_search"
+PREF_ADD_TASK_LABEL = "pref_add"
+MEM_FEEDBACK_TASK_LABEL = "mem_feedback"
+
+# Additional constants moved from general_schemas
+DEFAULT_MAX_QUERY_KEY_WORDS = 1000
+LONG_TERM_MEMORY_TYPE = "LongTermMemory"
+USER_INPUT_TYPE = "UserInput"
+NOT_APPLICABLE_TYPE = "NotApplicable"
+
+
+# scheduler daemon defaults
+# Interval in seconds for periodically releasing stale pending messages
+DEFAULT_PENDING_REQUEUE_INTERVAL_SEC = 30.0
+
+# Interval in seconds for refreshing cached Redis stream keys
+DEFAULT_STREAM_KEYS_REFRESH_INTERVAL_SEC = 30.0
+
+# Interval in seconds for batching and cleaning up deletions (xdel)
+DEFAULT_DELETE_CLEANUP_INTERVAL_SEC = 30.0
+
+# pending claim configuration
+# Only claim pending messages whose idle time exceeds this threshold.
+# Unit: milliseconds. Default: 1 hour.
+DEFAULT_PENDING_CLAIM_MIN_IDLE_MS = 3_600_000
+
+
+# Recency threshold for active streams
+# Consider a stream "active" if its last message is within this window.
+# Unit: seconds. Default: 1 hours.
+DEFAULT_STREAM_RECENT_ACTIVE_SECONDS = 3_600.0
+
+
+# Inactivity threshold for stream deletion
+# Delete streams whose last message ID timestamp is older than this threshold.
+# Unit: seconds. Default: 2 hour.
+DEFAULT_STREAM_INACTIVITY_DELETE_SECONDS = 7_200.0
+
+
+# task queue
+DEFAULT_STREAM_KEY_PREFIX = os.getenv(
+    "MEMSCHEDULER_STREAM_KEY_PREFIX", "scheduler:messages:stream:v2.0"
+)
+
+
+# ============== Running Tasks ==============
+class RunningTaskItem(BaseModel, DictConversionMixin):
+    """Data class for tracking running tasks in SchedulerDispatcher."""
+
+    item_id: str = Field(
+        description="Unique identifier for the task item", default_factory=lambda: str(uuid4())
+    )
+    user_id: str = Field(..., description="Required user identifier", min_length=1)
+    mem_cube_id: str = Field(..., description="Required memory cube identifier", min_length=1)
+    task_info: str = Field(..., description="Information about the task being executed")
+    task_name: str = Field(..., description="Name/type of the task handler")
+    start_time: datetime = Field(description="Task start time", default_factory=get_utc_now)
+    end_time: datetime | None = Field(default=None, description="Task completion time")
+    status: str = Field(default="running", description="Task status: running, completed, failed")
+    result: Any | None = Field(default=None, description="Task execution result")
+    error_message: str | None = Field(default=None, description="Error message if task failed")
+    messages: list[Any] | None = Field(
+        default=None, description="List of messages being processed by this task"
+    )
+
+    def mark_completed(self, result: Any | None = None) -> None:
+        """Mark task as completed with optional result."""
+        self.end_time = get_utc_now()
+        self.status = "completed"
+        self.result = result
+
+    def mark_failed(self, error_message: str) -> None:
+        """Mark task as failed with error message."""
+        self.end_time = get_utc_now()
+        self.status = "failed"
+        self.error_message = error_message
+
+    @computed_field
+    @property
+    def duration_seconds(self) -> float | None:
+        """Calculate task duration in seconds."""
+        if self.end_time:
+            return (self.end_time - self.start_time).total_seconds()
+        return None
+
+    def get_execution_info(self) -> str:
+        """Get formatted execution information for logging."""
+        duration = self.duration_seconds
+        duration_str = f"{duration:.2f}s" if duration else "ongoing"
+
+        return (
+            f"Task {self.task_name} (ID: {self.item_id[:8]}) "
+            f"for user {self.user_id}, cube {self.mem_cube_id} - "
+            f"Status: {self.status}, Duration: {duration_str}"
+        )