MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/scheduler_factory.py

@@ -0,0 +1,23 @@
+from typing import Any, ClassVar
+
+from memos.configs.mem_scheduler import SchedulerConfigFactory
+from memos.mem_scheduler.base_scheduler import BaseScheduler
+from memos.mem_scheduler.general_scheduler import GeneralScheduler
+from memos.mem_scheduler.optimized_scheduler import OptimizedScheduler
+
+
+class SchedulerFactory(BaseScheduler):
+    """Factory class for creating scheduler instances."""
+
+    backend_to_class: ClassVar[dict[str, Any]] = {
+        "general_scheduler": GeneralScheduler,
+        "optimized_scheduler": OptimizedScheduler,
+    }
+
+    @classmethod
+    def from_config(cls, config_factory: SchedulerConfigFactory) -> GeneralScheduler:
+        backend = config_factory.backend
+        if backend not in cls.backend_to_class:
+            raise ValueError(f"Invalid backend: {backend}")
+        mem_scheduler_class = cls.backend_to_class[backend]
+        return mem_scheduler_class(config_factory.config)

memos/mem_scheduler/schemas/analyzer_schemas.py

@@ -0,0 +1,52 @@
+import json
+
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+
+class BasicRecordingCase(BaseModel):
+    # Conversation identification
+    conv_id: str = Field(description="Conversation identifier for this evaluation case")
+    user_id: str = Field(description="User identifier for this evaluation case")
+    memcube_id: str = Field(description="Memcube identifier for this evaluation case")
+
+    # Query and answer information
+    query: str = Field(description="The current question/query being evaluated")
+
+    answer: str = Field(description="The generated answer for the query")
+
+    golden_answer: str | None = Field(
+        default=None, description="Ground truth answer for evaluation"
+    )
+
+    def to_dict(self) -> dict[str, Any]:
+        return self.dict()
+
+    def to_json(self, indent: int = 2) -> str:
+        return self.json(indent=indent, ensure_ascii=False)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "BasicRecordingCase":
+        return cls(**data)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "BasicRecordingCase":
+        data = json.loads(json_str)
+        return cls.from_dict(data)
+
+    class Config:
+        """Pydantic configuration"""
+
+        extra = "allow"  # Allow additional fields not defined in the schema
+        validate_assignment = True  # Validate on assignment
+        use_enum_values = True  # Use enum values instead of enum names

memos/mem_scheduler/schemas/api_schemas.py

@@ -0,0 +1,233 @@
+from datetime import datetime
+from enum import Enum
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field, field_serializer
+
+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
+from memos.memories.textual.item import TextualMemoryItem
+
+
+logger = get_logger(__name__)
+
+
+class TaskRunningStatus(str, Enum):
+    """Enumeration for task running status values."""
+
+    RUNNING = "running"
+    COMPLETED = "completed"
+
+
+class APIMemoryHistoryEntryItem(BaseModel, DictConversionMixin):
+    """Data class for search entry items stored in Redis."""
+
+    item_id: str = Field(
+        description="Unique identifier for the task", default_factory=lambda: str(uuid4())
+    )
+    query: str = Field(..., description="Search query string")
+    formatted_memories: Any = Field(..., description="Formatted search results")
+    memories: list[TextualMemoryItem] = Field(
+        default_factory=list, description="List of TextualMemoryItem objects"
+    )
+    task_status: str = Field(
+        default="running", description="Task status: running, completed, failed"
+    )
+    session_id: str | None = Field(default=None, description="Optional conversation identifier")
+    created_time: datetime = Field(description="Entry creation time", default_factory=get_utc_now)
+    timestamp: datetime | None = Field(default=None, description="Timestamp for the entry")
+    conversation_turn: int = Field(default=0, description="Turn count for the same session_id")
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        validate_assignment=True,
+    )
+
+    @field_serializer("created_time")
+    def serialize_created_time(self, value: datetime) -> str:
+        """Serialize datetime to ISO format string."""
+        return value.isoformat()
+
+    def get(self, key: str, default: Any | None = None) -> Any:
+        """
+        Get attribute value by key name, similar to dict.get().
+
+        Args:
+            key: The attribute name to retrieve
+            default: Default value to return if attribute doesn't exist
+
+        Returns:
+            The attribute value or default if not found
+        """
+        return getattr(self, key, default)
+
+
+class APISearchHistoryManager(BaseModel, DictConversionMixin):
+    """
+    Data structure for managing search history with separate completed and running entries.
+    Supports window_size to limit the number of completed entries.
+    """
+
+    window_size: int = Field(default=5, description="Maximum number of completed entries to keep")
+    completed_entries: list[APIMemoryHistoryEntryItem] = Field(
+        default_factory=list, description="List of completed search entries"
+    )
+    running_item_ids: list[str] = Field(
+        default_factory=list, description="List of running task ids"
+    )
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True,
+        validate_assignment=True,
+    )
+
+    def complete_entry(self, task_id: str) -> bool:
+        """
+        Remove task_id from running list when completed.
+        Note: The actual entry data should be managed separately.
+
+        Args:
+            task_id: The task ID to complete
+
+        Returns:
+            True if task_id was found and removed, False otherwise
+        """
+        if task_id in self.running_item_ids:
+            self.running_item_ids.remove(task_id)
+            logger.debug(f"Completed task_id: {task_id}")
+            return True
+
+        logger.warning(f"Task ID {task_id} not found in running task ids")
+        return False
+
+    def get_running_item_ids(self) -> list[str]:
+        """Get all running task IDs"""
+        return self.running_item_ids.copy()
+
+    def get_completed_entries(self) -> list[APIMemoryHistoryEntryItem]:
+        """Get all completed entries"""
+        return self.completed_entries.copy()
+
+    def get_history_memory_entries(
+        self, turns: int | None = None
+    ) -> list[APIMemoryHistoryEntryItem]:
+        """
+        Get the most recent n completed search entries, sorted by created_time.
+
+        Args:
+            turns: Number of entries to return. If None, returns all completed entries.
+
+        Returns:
+            List of completed search entries, sorted by created_time (newest first)
+        """
+        if not self.completed_entries:
+            return []
+
+        # Sort by created_time (newest first)
+        sorted_entries = sorted(self.completed_entries, key=lambda x: x.created_time, reverse=True)
+
+        if turns is None:
+            return sorted_entries
+
+        return sorted_entries[:turns]
+
+    def get_history_memories(self, turns: int | None = None) -> list[TextualMemoryItem]:
+        """
+        Get the most recent n completed search entries, sorted by created_time.
+
+        Args:
+            turns: Number of entries to return. If None, returns all completed entries.
+
+        Returns:
+            List of TextualMemoryItem objects from completed entries, sorted by created_time (newest first)
+        """
+        sorted_entries = self.get_history_memory_entries(turns=turns)
+
+        memories = []
+        for one in sorted_entries:
+            memories.extend(one.memories)
+        return memories
+
+    def find_entry_by_item_id(self, item_id: str) -> tuple[dict[str, Any] | None, str]:
+        """
+        Find an entry by item_id in completed list only.
+        Running entries are now just task IDs, so we can only search completed entries.
+
+        Args:
+            item_id: The item ID to search for
+
+        Returns:
+            Tuple of (entry_dict, location) where location is 'completed' or 'not_found'
+        """
+        # Check completed entries
+        for entry in self.completed_entries:
+            try:
+                if hasattr(entry, "item_id") and entry.item_id == item_id:
+                    return entry.to_dict(), "completed"
+                elif isinstance(entry, dict) and entry.get("item_id") == item_id:
+                    return entry, "completed"
+            except AttributeError as e:
+                logger.warning(f"Entry missing item_id attribute: {e}, entry type: {type(entry)}")
+                continue
+
+        return None, "not_found"
+
+    def update_entry_by_item_id(
+        self,
+        item_id: str,
+        query: str,
+        formatted_memories: Any,
+        task_status: TaskRunningStatus,
+        session_id: str | None = None,
+        memories: list[TextualMemoryItem] | None = None,
+    ) -> bool:
+        """
+        Update an existing entry by item_id. Since running entries are now just IDs,
+        this method can only update completed entries.
+
+        Args:
+            item_id: The item ID to update
+            query: New query string
+            formatted_memories: New formatted memories
+            task_status: New task status
+            session_id: New conversation ID
+            memories: List of TextualMemoryItem objects
+
+        Returns:
+            True if entry was found and updated, False otherwise
+        """
+        # Find the entry in completed list
+        for entry in self.completed_entries:
+            if entry.item_id == item_id:
+                # Update the entry content
+                entry.query = query
+                entry.formatted_memories = formatted_memories
+                entry.task_status = task_status
+                if session_id is not None:
+                    entry.session_id = session_id
+                if memories is not None:
+                    entry.memories = memories
+
+                logger.debug(f"Updated entry with item_id: {item_id}, new status: {task_status}")
+                return True
+
+        logger.warning(f"Entry with item_id: {item_id} not found in completed entries")
+        return False
+
+    def get_total_count(self) -> dict[str, int]:
+        """Get count of entries by status"""
+        return {
+            "completed": len(self.completed_entries),
+            "running": len(self.running_item_ids),
+            "total": len(self.completed_entries) + len(self.running_item_ids),
+        }
+
+    def __len__(self) -> int:
+        """Return total number of entries (completed + running)"""
+        return len(self.completed_entries) + len(self.running_item_ids)
+
+
+# Alias for easier usage
+SearchHistoryManager = APISearchHistoryManager

memos/mem_scheduler/schemas/general_schemas.py

@@ -0,0 +1,55 @@
+import os
+
+from pathlib import Path
+
+
+FILE_PATH = Path(__file__).absolute()
+BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
+
+TreeTextMemory_SEARCH_METHOD = "tree_text_memory_search"
+TreeTextMemory_FINE_SEARCH_METHOD = "tree_text_memory_fine_search"
+TextMemory_SEARCH_METHOD = "text_memory_search"
+DIRECT_EXCHANGE_TYPE = "direct"
+FANOUT_EXCHANGE_TYPE = "fanout"
+DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT = 30
+DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT = 20
+DEFAULT_ACT_MEM_DUMP_PATH = f"{BASE_DIR}/outputs/mem_scheduler/mem_cube_scheduler_test.kv_cache"
+DEFAULT_THREAD_POOL_MAX_WORKERS = 50
+DEFAULT_CONSUME_INTERVAL_SECONDS = 0.01
+DEFAULT_CONSUME_BATCH = 3
+DEFAULT_DISPATCHER_MONITOR_CHECK_INTERVAL = 300
+DEFAULT_DISPATCHER_MONITOR_MAX_FAILURES = 2
+DEFAULT_STUCK_THREAD_TOLERANCE = 10
+DEFAULT_MAX_INTERNAL_MESSAGE_QUEUE_SIZE = -1
+DEFAULT_TOP_K = 5
+DEFAULT_CONTEXT_WINDOW_SIZE = 5
+DEFAULT_USE_REDIS_QUEUE = os.getenv("MEMSCHEDULER_USE_REDIS_QUEUE", "False").lower() == "true"
+DEFAULT_MULTI_TASK_RUNNING_TIMEOUT = 30
+DEFAULT_SCHEDULER_RETRIEVER_BATCH_SIZE = 20
+DEFAULT_SCHEDULER_RETRIEVER_RETRIES = 1
+DEFAULT_STOP_WAIT = False
+
+# startup mode configuration
+STARTUP_BY_THREAD = "thread"
+STARTUP_BY_PROCESS = "process"
+DEFAULT_STARTUP_MODE = STARTUP_BY_THREAD  # default to thread mode
+
+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]
+DEFAULT_MAX_WEB_LOG_QUEUE_SIZE = 50

memos/mem_scheduler/schemas/message_schemas.py

@@ -0,0 +1,173 @@
+from datetime import datetime
+from typing import Any
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field
+from typing_extensions import TypedDict
+
+from memos.context.context import generate_trace_id
+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
+
+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()))
+    redis_message_id: str = Field(default="", description="the message get from redis stream")
+    stream_key: str = Field("", description="stream_key for identifying the queue in line")
+    user_id: str = Field(..., description="user id")
+    trace_id: str = Field(default_factory=generate_trace_id, description="trace id for logging")
+    mem_cube_id: str = Field(..., description="memcube id")
+    session_id: str = Field(default="", description="Session ID for soft-filtering memories")
+    label: str = Field(..., description="Label of the schedule message")
+    content: str = Field(..., description="Content of the schedule message")
+    timestamp: datetime = Field(
+        default_factory=get_utc_now, description="submit time for schedule_messages"
+    )
+    user_name: str = Field(
+        default="",
+        description="user name / display name (optional)",
+    )
+    info: dict | None = Field(default=None, description="user custom info")
+    task_id: str | None = Field(
+        default=None,
+        description="Optional business-level task ID. Multiple items can share the same task_id.",
+    )
+
+    # 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
+                "content": "sample content",  # Example message content
+                "timestamp": "2024-07-22T12:00:00Z",  # Added timestamp example
+                "user_name": "Alice",  # Added username example
+            }
+        },
+    )
+
+    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,
+            "trace_id": self.trace_id,
+            "label": self.label,
+            "cube": "Not Applicable",  # Custom cube serialization
+            "content": self.content,
+            "timestamp": self.timestamp.isoformat(),
+            "user_name": self.user_name,
+            "task_id": self.task_id if self.task_id is not None else "",
+        }
+
+    @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"],
+            mem_cube_id=data["cube_id"],
+            trace_id=data.get("trace_id", generate_trace_id()),
+            label=data["label"],
+            content=data["content"],
+            timestamp=datetime.fromisoformat(data["timestamp"]),
+            user_name=data.get("user_name"),
+            task_id=data.get("task_id"),
+        )
+
+
+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())
+    )
+    task_id: str | None = Field(default=None, description="Identifier for the parent task")
+    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 | None = Field(None, description="Source memory type")
+    to_memory_type: str | None = Field(None, 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=get_utc_now,
+        description="Timestamp indicating when the log entry was created",
+    )
+    memcube_log_content: list[dict] | None = Field(
+        default=None, description="Structured memcube log content list"
+    )
+    metadata: list[dict] | None = Field(
+        default=None, description="Structured metadata list for each log item"
+    )
+    memcube_name: str | None = Field(default=None, description="Display name for memcube")
+    memory_len: int | None = Field(default=None, description="Count of items involved in the event")
+    status: str | None = Field(
+        default=None, description="Completion status of the task (e.g., 'completed', 'failed')"
+    )
+    source_doc_id: str | None = Field(default=None, description="Source document ID")
+
+    def debug_info(self) -> dict[str, Any]:
+        """Return structured debug information for logging purposes."""
+        return {
+            "content_preview:": self.log_content[:50],
+            "item_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(),
+        }