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

memos/mem_scheduler/modules/schemas.py

@@ -1,35 +1,65 @@
+import json
+
 from datetime import datetime
 from pathlib import Path
-from typing import ClassVar, TypeVar
+from typing import ClassVar, NewType, TypeVar
 from uuid import uuid4
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, computed_field
 from typing_extensions import TypedDict
 
+from memos.log import get_logger
 from memos.mem_cube.general import GeneralMemCube
 
 
+logger = get_logger(__name__)
+
+
 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"
-DEFAULT_ACTIVATION_MEM_SIZE = 5
+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
 BaseModelType = TypeVar("T", bound="BaseModel")
 
+# 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"
+
+
+# new types
+UserID = NewType("UserID", str)
+MemCubeID = NewType("CubeID", str)
 
+
+# ************************* Public *************************
 class DictConversionMixin:
     def to_dict(self) -> dict:
         """Convert the instance to a dictionary."""
         return {
-            **self.dict(),
+            **self.model_dump(),  # 替换 self.dict()
             "timestamp": self.timestamp.isoformat() if hasattr(self, "timestamp") else None,
         }
 
@@ -40,10 +70,25 @@ class DictConversionMixin:
             data["timestamp"] = datetime.fromisoformat(data["timestamp"])
         return cls(**data)
 
+    def __str__(self) -> str:
+        """Convert the instance to a JSON string with indentation of 4 spaces.
+        This will be used when str() or print() is called on the instance.
+
+        Returns:
+            str: A JSON string representation of the instance with 4-space indentation.
+        """
+        return json.dumps(
+            self.to_dict(),
+            indent=4,
+            ensure_ascii=False,
+            default=str,  # 处理无法序列化的对象
+        )
+
     class Config:
         json_encoders: ClassVar[dict[type, object]] = {datetime: lambda v: v.isoformat()}
 
 
+# ************************* Messages *************************
 class ScheduleMessageItem(BaseModel, DictConversionMixin):
     item_id: str = Field(description="uuid", default_factory=lambda: str(uuid4()))
     user_id: str = Field(..., description="user id")
@@ -68,7 +113,6 @@ class ScheduleMessageItem(BaseModel, DictConversionMixin):
             "item_id": self.item_id,
             "user_id": self.user_id,
             "cube_id": self.mem_cube_id,
-            "message_id": self.message_id,
             "label": self.label,
             "cube": "Not Applicable",  # Custom cube serialization
             "content": self.content,
@@ -82,7 +126,6 @@ class ScheduleMessageItem(BaseModel, DictConversionMixin):
             item_id=data.get("item_id", str(uuid4())),
             user_id=data["user_id"],
             cube_id=data["cube_id"],
-            message_id=data.get("message_id", str(uuid4())),
             label=data["label"],
             cube="Not Applicable",  # Custom cube deserialization
             content=data["content"],
@@ -130,7 +173,8 @@ class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
         ..., description="Identifier for the memcube associated with this log entry"
     )
     label: str = Field(..., description="Label categorizing the type of log")
-    log_title: str = Field(..., description="Title or brief summary of the log content")
+    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),
@@ -144,3 +188,141 @@ class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
         default_factory=datetime.now,
         description="Timestamp indicating when the log entry was created",
     )
+
+
+# ************************* Monitor *************************
+class MemoryMonitorItem(BaseModel, DictConversionMixin):
+    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,
+        max_length=10000,  # Prevent excessively large memory texts
+    )
+    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,  # Greater than or equal to 1
+    )
+
+    def get_score(self) -> float:
+        """
+        Calculate the effective score for the memory item.
+
+        Returns:
+            float: The importance_score if it has been initialized (>=0),
+                   otherwise the recording_count converted to float.
+
+        Note:
+            This method provides a unified way to retrieve a comparable score
+            for memory items, regardless of whether their importance has been explicitly set.
+        """
+        if self.importance_score == NOT_INITIALIZED:
+            # Return recording_count as float when importance_score is not initialized
+            return float(self.recording_count)
+        else:
+            # Return the initialized importance_score
+            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)
+
+    def update_memories(
+        self, text_working_memories: list[str], partial_retention_number: int
+    ) -> MemoryMonitorItem:
+        """
+        Update memories based on text_working_memories.
+
+        Args:
+            text_working_memories: List of memory texts to update
+            partial_retention_number: Number of top memories to keep by recording count
+
+        Returns:
+            List of added or updated MemoryMonitorItem instances
+        """
+
+        # Validate partial_retention_number
+        if partial_retention_number < 0:
+            raise ValueError("partial_retention_number must be non-negative")
+
+        # Create text lookup set
+        working_memory_set = set(text_working_memories)
+
+        # Step 1: Update existing memories or add new ones
+        added_or_updated = []
+        memory_text_map = {item.memory_text: item for item in self.memories}
+
+        for text in text_working_memories:
+            if text in memory_text_map:
+                # Update existing memory
+                memory = memory_text_map[text]
+                memory.recording_count += 1
+                added_or_updated.append(memory)
+            else:
+                # Add new memory
+                new_memory = MemoryMonitorItem(memory_text=text, recording_count=1)
+                self.memories.append(new_memory)
+                added_or_updated.append(new_memory)
+
+        # Step 2: Identify memories to remove
+        # Sort memories by recording_count in descending order
+        sorted_memories = sorted(self.memories, key=lambda item: item.recording_count, reverse=True)
+
+        # Keep the top N memories by recording_count
+        records_to_keep = {
+            memory.memory_text for memory in sorted_memories[:partial_retention_number]
+        }
+
+        # Collect memories to remove: not in current working memory and not in top N
+        memories_to_remove = [
+            memory
+            for memory in self.memories
+            if memory.memory_text not in working_memory_set
+            and memory.memory_text not in records_to_keep
+        ]
+
+        # Step 3: Remove identified memories
+        for memory in memories_to_remove:
+            self.memories.remove(memory)
+
+        # Step 4: Enforce max_capacity if set
+        if self.max_capacity is not None and len(self.memories) > self.max_capacity:
+            # Sort by importance and then recording count
+            sorted_memories = sorted(
+                self.memories,
+                key=lambda item: (item.importance_score, item.recording_count),
+                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: {len(added_or_updated)}, "
+            f"Removed: {len(memories_to_remove)} (excluding top {partial_retention_number} by recording_count)"
+        )
+
+        return added_or_updated

memos/mem_scheduler/mos_for_test_scheduler.py

@@ -0,0 +1,143 @@
+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 (
+    ANSWER_LABEL,
+    MONITOR_WORKING_MEMORY_TYPE,
+    QUERY_LABEL,
+    ScheduleMessageItem,
+)
+
+
+logger = get_logger(__name__)
+
+
+class MOSForTestScheduler(MOS):
+    """This class is only to test abilities of mem scheduler"""
+
+    def __init__(self, config: MOSConfig):
+        super().__init__(config)
+
+    def _str_memories(self, memories: list[str]) -> str:
+        """Format memories for display."""
+        if not memories:
+            return "No memories."
+        return "\n".join(f"{i + 1}. {memory}" for i, memory in enumerate(memories))
+
+    def chat(self, query: str, user_id: str | None = None) -> str:
+        """
+        Chat with the MOS.
+
+        Args:
+            query (str): The user's query.
+
+        Returns:
+            str: The response from the MOS.
+        """
+        target_user_id = user_id if user_id is not None else self.user_id
+        accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
+        user_cube_ids = [cube.cube_id for cube in accessible_cubes]
+        if target_user_id not in self.chat_history_manager:
+            self._register_chat_history(target_user_id)
+
+        chat_history = self.chat_history_manager[target_user_id]
+
+        topk_for_scheduler = 2
+
+        if self.config.enable_textual_memory and self.mem_cubes:
+            memories_all = []
+            for mem_cube_id, mem_cube in self.mem_cubes.items():
+                if mem_cube_id not in user_cube_ids:
+                    continue
+                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,
+                    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,
+                )
+
+                # from scheduler
+                scheduler_memories = self.mem_scheduler.monitor.get_monitor_memories(
+                    user_id=target_user_id,
+                    mem_cube_id=mem_cube_id,
+                    memory_type=MONITOR_WORKING_MEMORY_TYPE,
+                    top_k=topk_for_scheduler,
+                )
+                memories_all.extend(scheduler_memories)
+
+                # from mem_cube
+                memories = mem_cube.text_mem.search(
+                    query, top_k=self.config.top_k - topk_for_scheduler
+                )
+                text_memories = [m.memory for m in memories]
+                memories_all.extend(text_memories)
+
+                memories_all = list(set(memories_all))
+
+            logger.info(f"🧠 [Memory] Searched memories:\n{self._str_memories(memories_all)}\n")
+            system_prompt = self._build_system_prompt(memories_all)
+        else:
+            system_prompt = self._build_system_prompt()
+        current_messages = [
+            {"role": "system", "content": system_prompt},
+            *chat_history.chat_history,
+            {"role": "user", "content": query},
+        ]
+        past_key_values = None
+
+        if self.config.enable_activation_memory:
+            assert self.config.chat_model.backend == "huggingface", (
+                "Activation memory only used for huggingface backend."
+            )
+            # TODO this only one cubes
+            for mem_cube_id, mem_cube in self.mem_cubes.items():
+                if mem_cube_id not in user_cube_ids:
+                    continue
+                if mem_cube.act_mem:
+                    kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
+                    past_key_values = (
+                        kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
+                    )
+                    break
+            # Generate response
+            response = self.chat_llm.generate(current_messages, past_key_values=past_key_values)
+        else:
+            response = self.chat_llm.generate(current_messages)
+        logger.info(f"🤖 [Assistant] {response}\n")
+        chat_history.chat_history.append({"role": "user", "content": query})
+        chat_history.chat_history.append({"role": "assistant", "content": response})
+        self.chat_history_manager[user_id] = chat_history
+
+        # submit message to scheduler
+        for accessible_mem_cube in accessible_cubes:
+            mem_cube_id = accessible_mem_cube.cube_id
+            mem_cube = self.mem_cubes[mem_cube_id]
+            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=ANSWER_LABEL,
+                    content=response,
+                    timestamp=datetime.now(),
+                )
+                self.mem_scheduler.submit_messages(messages=[message_item])
+
+        return response

memos/mem_scheduler/utils.py

@@ -1,4 +1,5 @@
 import json
+import re
 
 from pathlib import Path
 
@@ -7,13 +8,41 @@ import yaml
 
 def extract_json_dict(text: str):
     text = text.strip()
-    patterns_to_remove = ["json'''", "latex'''", "'''"]
+    patterns_to_remove = ["json```", "```json", "latex```", "```latex", "```"]
     for pattern in patterns_to_remove:
         text = text.replace(pattern, "")
-    res = json.loads(text)
+    res = json.loads(text.strip())
     return res
 
 
+def transform_name_to_key(name):
+    """
+    Normalize text by removing all punctuation marks, keeping only letters, numbers, and word characters.
+
+    Args:
+        name (str): Input text to be processed
+
+    Returns:
+        str: Processed text with all punctuation removed
+    """
+    # Match all characters that are NOT:
+    # \w - word characters (letters, digits, underscore)
+    # \u4e00-\u9fff - Chinese/Japanese/Korean characters
+    # \s - whitespace
+    pattern = r"[^\w\u4e00-\u9fff\s]"
+
+    # Substitute all matched punctuation marks with empty string
+    # re.UNICODE flag ensures proper handling of Unicode characters
+    normalized = re.sub(pattern, "", name, flags=re.UNICODE)
+
+    # Optional: Collapse multiple whitespaces into single space
+    normalized = "_".join(normalized.split())
+
+    normalized = normalized.lower()
+
+    return normalized
+
+
 def parse_yaml(yaml_file):
     yaml_path = Path(yaml_file)
     yaml_path = Path(yaml_file)
@@ -24,3 +53,23 @@ def parse_yaml(yaml_file):
         data = yaml.safe_load(fr)
 
     return data
+
+
+def is_all_english(input_string: str) -> bool:
+    """Determine if the string consists entirely of English characters (including spaces)"""
+    return all(char.isascii() or char.isspace() for char in input_string)
+
+
+def is_all_chinese(input_string: str) -> bool:
+    """Determine if the string consists entirely of Chinese characters (including Chinese punctuation and spaces)"""
+    return all(
+        ("\u4e00" <= char <= "\u9fff")  # Basic Chinese characters
+        or ("\u3400" <= char <= "\u4dbf")  # Extension A
+        or ("\u20000" <= char <= "\u2a6df")  # Extension B
+        or ("\u2a700" <= char <= "\u2b73f")  # Extension C
+        or ("\u2b740" <= char <= "\u2b81f")  # Extension D
+        or ("\u2b820" <= char <= "\u2ceaf")  # Extension E
+        or ("\u2f800" <= char <= "\u2fa1f")  # Extension F
+        or char.isspace()  # Spaces
+        for char in input_string
+    )