MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/task_schedule_modules/task_queue.py

@@ -0,0 +1,162 @@
+"""
+Redis Queue implementation for SchedulerMessageItem objects.
+
+This module provides a Redis-based queue implementation that can replace
+the local memos_message_queue functionality in BaseScheduler.
+"""
+
+from memos.context.context import get_current_trace_id
+from memos.log import get_logger
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+from memos.mem_scheduler.task_schedule_modules.local_queue import SchedulerLocalQueue
+from memos.mem_scheduler.task_schedule_modules.orchestrator import SchedulerOrchestrator
+from memos.mem_scheduler.task_schedule_modules.redis_queue import SchedulerRedisQueue
+from memos.mem_scheduler.utils.db_utils import get_utc_now
+from memos.mem_scheduler.utils.misc_utils import group_messages_by_user_and_mem_cube
+from memos.mem_scheduler.utils.monitor_event_utils import emit_monitor_event, to_iso
+from memos.mem_scheduler.utils.status_tracker import TaskStatusTracker
+
+
+logger = get_logger(__name__)
+
+
+class ScheduleTaskQueue:
+    def __init__(
+        self,
+        use_redis_queue: bool,
+        maxsize: int,
+        disabled_handlers: list | None = None,
+        orchestrator: SchedulerOrchestrator | None = None,
+        status_tracker: TaskStatusTracker | None = None,
+    ):
+        self.use_redis_queue = use_redis_queue
+        self.maxsize = maxsize
+        self.orchestrator = SchedulerOrchestrator() if orchestrator is None else orchestrator
+        self.status_tracker = status_tracker
+
+        if self.use_redis_queue:
+            if maxsize is None or not isinstance(maxsize, int) or maxsize <= 0:
+                maxsize = None
+            self.memos_message_queue = SchedulerRedisQueue(
+                max_len=maxsize,
+                consumer_group="scheduler_group",
+                consumer_name="scheduler_consumer",
+                orchestrator=self.orchestrator,
+                status_tracker=self.status_tracker,  # Propagate status_tracker
+            )
+        else:
+            self.memos_message_queue = SchedulerLocalQueue(maxsize=self.maxsize)
+
+        self.disabled_handlers = disabled_handlers
+
+    def set_status_tracker(self, status_tracker: TaskStatusTracker) -> None:
+        """
+        Set the status tracker for this queue and propagate it to the underlying queue implementation.
+
+        This allows the tracker to be injected after initialization (e.g., when Redis connection becomes available).
+        """
+        self.status_tracker = status_tracker
+        if self.memos_message_queue and hasattr(self.memos_message_queue, "status_tracker"):
+            # SchedulerRedisQueue has status_tracker attribute (from our previous fix)
+            # SchedulerLocalQueue can also accept it dynamically if it doesn't use __slots__
+            self.memos_message_queue.status_tracker = status_tracker
+            logger.info("Propagated status_tracker to underlying message queue")
+
+    def ack_message(
+        self,
+        user_id: str,
+        mem_cube_id: str,
+        task_label: str,
+        redis_message_id,
+        message: ScheduleMessageItem | None,
+    ) -> None:
+        if not isinstance(self.memos_message_queue, SchedulerRedisQueue):
+            logger.warning("ack_message is only supported for Redis queues")
+            return
+
+        self.memos_message_queue.ack_message(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            task_label=task_label,
+            redis_message_id=redis_message_id,
+            message=message,
+        )
+
+    def get_stream_keys(self) -> list[str]:
+        if isinstance(self.memos_message_queue, SchedulerRedisQueue):
+            stream_keys = self.memos_message_queue.get_stream_keys()
+        else:
+            stream_keys = list(self.memos_message_queue.queue_streams.keys())
+        return stream_keys
+
+    def submit_messages(self, messages: ScheduleMessageItem | list[ScheduleMessageItem]):
+        """Submit messages to the message queue (either local queue or Redis)."""
+        if isinstance(messages, ScheduleMessageItem):
+            messages = [messages]
+
+        current_trace_id = get_current_trace_id()
+
+        for msg in messages:
+            if current_trace_id:
+                # Prefer current request trace_id so logs can be correlated
+                msg.trace_id = current_trace_id
+            msg.stream_key = self.memos_message_queue.get_stream_key(
+                user_id=msg.user_id, mem_cube_id=msg.mem_cube_id, task_label=msg.label
+            )
+
+        if len(messages) < 1:
+            logger.error("Submit empty")
+        elif len(messages) == 1:
+            if getattr(messages[0], "timestamp", None) is None:
+                messages[0].timestamp = get_utc_now()
+            enqueue_ts = to_iso(getattr(messages[0], "timestamp", None))
+            emit_monitor_event(
+                "enqueue",
+                messages[0],
+                {"enqueue_ts": enqueue_ts, "event_duration_ms": 0, "total_duration_ms": 0},
+            )
+            self.memos_message_queue.put(messages[0])
+        else:
+            user_cube_groups = group_messages_by_user_and_mem_cube(messages)
+
+            # Process each user and mem_cube combination
+            for _user_id, cube_groups in user_cube_groups.items():
+                for _mem_cube_id, user_cube_msgs in cube_groups.items():
+                    for message in user_cube_msgs:
+                        if not isinstance(message, ScheduleMessageItem):
+                            error_msg = f"Invalid message type: {type(message)}, expected ScheduleMessageItem"
+                            logger.error(error_msg)
+                            raise TypeError(error_msg)
+
+                        if getattr(message, "timestamp", None) is None:
+                            message.timestamp = get_utc_now()
+
+                        if self.disabled_handlers and message.label in self.disabled_handlers:
+                            logger.info(
+                                f"Skipping disabled handler: {message.label} - {message.content}"
+                            )
+                            continue
+
+                        enqueue_ts = to_iso(getattr(message, "timestamp", None))
+                        emit_monitor_event(
+                            "enqueue",
+                            message,
+                            {
+                                "enqueue_ts": enqueue_ts,
+                                "event_duration_ms": 0,
+                                "total_duration_ms": 0,
+                            },
+                        )
+                        self.memos_message_queue.put(message)
+                        logger.info(
+                            f"Submitted message to local queue: {message.label} - {message.content}"
+                        )
+
+    def get_messages(self, batch_size: int) -> list[ScheduleMessageItem]:
+        return self.memos_message_queue.get_messages(batch_size=batch_size)
+
+    def clear(self):
+        self.memos_message_queue.clear()
+
+    def qsize(self):
+        return self.memos_message_queue.qsize()

memos/mem_scheduler/utils/api_utils.py

@@ -0,0 +1,77 @@
+import uuid
+
+from typing import Any
+
+from memos.memories.textual.item import TreeNodeTextualMemoryMetadata
+from memos.memories.textual.tree import TextualMemoryItem
+
+
+def format_textual_memory_item(memory_data: Any, include_embedding: bool = False) -> dict[str, Any]:
+    """Format a single memory item for API response."""
+    memory = memory_data.model_dump()
+    memory_id = memory["id"]
+    ref_id = f"[{memory_id.split('-')[0]}]"
+
+    memory["ref_id"] = ref_id
+    if not include_embedding:
+        memory["metadata"]["embedding"] = []
+    memory["metadata"]["sources"] = []
+    memory["metadata"]["ref_id"] = ref_id
+    memory["metadata"]["id"] = memory_id
+    memory["metadata"]["memory"] = memory["memory"]
+
+    return memory
+
+
+def make_textual_item(memory_data):
+    return memory_data
+
+
+def text_to_textual_memory_item(
+    text: str,
+    user_id: str | None = None,
+    session_id: str | None = None,
+    memory_type: str = "WorkingMemory",
+    tags: list[str] | None = None,
+    key: str | None = None,
+    sources: list | None = None,
+    background: str = "",
+    confidence: float = 0.99,
+    embedding: list[float] | None = None,
+) -> TextualMemoryItem:
+    """
+    Convert text into a TextualMemoryItem object.
+
+    Args:
+        text: Memory content text
+        user_id: User ID
+        session_id: Session ID
+        memory_type: Memory type, defaults to "WorkingMemory"
+        tags: List of tags
+        key: Memory key or title
+        sources: List of sources
+        background: Background information
+        confidence: Confidence score (0-1)
+        embedding: Vector embedding
+
+    Returns:
+        TextualMemoryItem: Wrapped memory item
+    """
+    return TextualMemoryItem(
+        id=str(uuid.uuid4()),
+        memory=text,
+        metadata=TreeNodeTextualMemoryMetadata(
+            user_id=user_id,
+            session_id=session_id,
+            memory_type=memory_type,
+            status="activated",
+            tags=tags or [],
+            key=key,
+            embedding=embedding or [],
+            usage=[],
+            sources=sources or [],
+            background=background,
+            confidence=confidence,
+            type="fact",
+        ),
+    )

memos/mem_scheduler/utils/config_utils.py

@@ -0,0 +1,100 @@
+import json
+import os
+
+from typing import Any
+
+import yaml
+
+
+def flatten_dict(
+    data: dict[str, Any], parent_keys: list[str] | None = None, prefix: str = ""
+) -> dict[str, str]:
+    """
+    Recursively flattens a nested dictionary to generate environment variable keys following the specified format.
+    Combines nested keys with underscores, converts to uppercase, and prepends a custom prefix if provided.
+
+    Args:
+        data: Nested dictionary to be flattened (parsed from JSON/YAML)
+        parent_keys: List to track nested keys during recursion
+        prefix: Custom prefix to be added to all generated keys
+
+    Returns:
+        Flattened dictionary with keys in PREFIX_KEY1_KEY2... format and string values
+    """
+    parent_keys = parent_keys or []
+    flat_data = {}
+
+    for key, value in data.items():
+        # Clean and standardize key: convert to uppercase, replace spaces/hyphens with underscores
+        clean_key = key.upper().replace(" ", "_").replace("-", "_")
+        current_keys = [*parent_keys, clean_key]
+
+        if isinstance(value, dict):
+            # Recursively process nested dictionaries
+            nested_flat = flatten_dict(value, current_keys, prefix)
+            flat_data.update(nested_flat)
+        else:
+            # Construct full key name with prefix (if provided) and nested keys
+            if prefix:
+                full_key = f"{prefix.upper()}_{'_'.join(current_keys)}"
+            else:
+                full_key = "_".join(current_keys)
+
+            # Process value: ensure string type, convert None to empty string
+            flat_value = "" if value is None else str(value).strip()
+
+            flat_data[full_key] = flat_value
+
+    return flat_data
+
+
+def convert_config_to_env(input_file: str, output_file: str = ".env", prefix: str = "") -> None:
+    """
+    Converts a JSON or YAML configuration file to a .env file with standardized environment variables.
+    Uses the flatten_dict function to generate keys in PREFIX_KEY1_KEY2... format.
+
+    Args:
+        input_file: Path to input configuration file (.json, .yaml, or .yml)
+        output_file: Path to output .env file (default: .env)
+        prefix: Custom prefix for all environment variable keys
+
+    Raises:
+        FileNotFoundError: If input file does not exist
+        ValueError: If file format is unsupported or parsing fails
+    """
+    # Check if input file exists
+    if not os.path.exists(input_file):
+        raise FileNotFoundError(f"Input file not found: {input_file}")
+
+    # Parse input file based on extension
+    file_ext = os.path.splitext(input_file)[1].lower()
+    config_data: dict[str, Any] = {}
+
+    try:
+        with open(input_file, encoding="utf-8") as f:
+            if file_ext in (".json",):
+                config_data = json.load(f)
+            elif file_ext in (".yaml", ".yml"):
+                config_data = yaml.safe_load(f)
+            else:
+                raise ValueError(
+                    f"Unsupported file format: {file_ext}. Supported formats: .json, .yaml, .yml"
+                )
+    except (json.JSONDecodeError, yaml.YAMLError) as e:
+        raise ValueError(f"Error parsing file: {e!s}") from e
+
+    # Flatten configuration and generate environment variable key-value pairs
+    flat_config = flatten_dict(config_data, prefix=prefix)
+
+    # Write to .env file
+    with open(output_file, "w", encoding="utf-8") as f:
+        for key, value in flat_config.items():
+            # Handle values containing double quotes (use no surrounding quotes)
+            if '"' in value:
+                f.write(f"{key}={value}\n")
+            else:
+                f.write(f'{key}="{value}"\n')  # Enclose regular values in double quotes
+
+    print(
+        f"Conversion complete! Generated {output_file} with {len(flat_config)} environment variables"
+    )

memos/mem_scheduler/utils/db_utils.py

@@ -0,0 +1,50 @@
+import os
+import sqlite3
+import sys
+
+from datetime import datetime, timezone
+
+
+# Compatibility handling: Python 3.11+ supports UTC, earlier versions use timezone.utc
+if sys.version_info >= (3, 11):
+    from datetime import UTC
+
+    def get_utc_now():
+        """Get current UTC datetime with compatibility for different Python versions"""
+        return datetime.now(UTC)
+else:
+
+    def get_utc_now():
+        """Get current UTC datetime with compatibility for different Python versions"""
+        return datetime.now(timezone.utc)
+
+
+def print_db_tables(db_path: str):
+    """Print all table names and structures in the SQLite database"""
+    print(f"\n🔍 Checking database file: {db_path}")
+
+    if not os.path.exists(db_path):
+        print(f"❌ File does not exist! Path: {db_path}")
+        return
+
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    # List all tables
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+    tables = cursor.fetchall()
+    if not tables:
+        print("❌ Database is empty, no tables created")
+    else:
+        print(f"✅ Database contains {len(tables)} table(s):")
+        for (table_name,) in tables:
+            print(f"  📂 Table name: {table_name}")
+
+            # Print table structure
+            cursor.execute(f"PRAGMA table_info({table_name});")
+            columns = cursor.fetchall()
+            print("    🧩 Structure:")
+            for col in columns:
+                print(f"      {col[1]} ({col[2]}) {'(PK)' if col[5] else ''}")
+
+    conn.close()

memos/mem_scheduler/utils/filter_utils.py

@@ -0,0 +1,176 @@
+import re
+
+from memos.dependency import require_python_package
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+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 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
+    )
+
+
+@require_python_package(
+    import_name="sklearn",
+    install_command="pip install scikit-learn",
+    install_link="https://scikit-learn.org/stable/install.html",
+)
+def filter_vector_based_similar_memories(
+    text_memories: list[str], similarity_threshold: float = 0.75
+) -> list[str]:
+    """
+    Filters out low-quality or duplicate memories based on text similarity.
+
+    Args:
+        text_memories: List of text memories to filter
+        similarity_threshold: Threshold for considering memories duplicates (0.0-1.0)
+                            Higher values mean stricter filtering
+
+    Returns:
+        List of filtered memories with duplicates removed
+    """
+    from sklearn.feature_extraction.text import TfidfVectorizer
+    from sklearn.metrics.pairwise import cosine_similarity
+
+    if not text_memories:
+        logger.warning("Received empty memories list - nothing to filter")
+        return []
+
+    for idx in range(len(text_memories)):
+        if not isinstance(text_memories[idx], str):
+            logger.error(
+                f"{text_memories[idx]} in memories is not a string,"
+                f" and now has been transformed to be a string."
+            )
+            text_memories[idx] = str(text_memories[idx])
+
+    try:
+        # Step 1: Vectorize texts using TF-IDF
+        vectorizer = TfidfVectorizer()
+        tfidf_matrix = vectorizer.fit_transform(text_memories)
+
+        # Step 2: Calculate pairwise similarity matrix
+        similarity_matrix = cosine_similarity(tfidf_matrix)
+
+        # Step 3: Identify duplicates
+        to_keep = set(range(len(text_memories)))  # Start with all indices
+        for i in range(len(similarity_matrix)):
+            if i not in to_keep:
+                continue  # Already marked for removal
+
+            # Find all similar items to this one (excluding self and already removed)
+            similar_indices = [
+                j
+                for j in range(i + 1, len(similarity_matrix))
+                if similarity_matrix[i][j] >= similarity_threshold and j in to_keep
+            ]
+            similar_indices = set(similar_indices)
+
+            # Remove all similar items (keeping the first one - i)
+            to_keep -= similar_indices
+
+        # Return filtered memories
+        filtered_memories = [text_memories[i] for i in sorted(to_keep)]
+        logger.debug(f"filtered_memories: {filtered_memories}")
+        return filtered_memories
+
+    except Exception as e:
+        logger.error(f"Error filtering memories: {e!s}")
+        return text_memories  # Return original list if error occurs
+
+
+def filter_too_short_memories(
+    text_memories: list[str], min_length_threshold: int = 20
+) -> list[str]:
+    """
+    Filters out text memories that fall below the minimum length requirement.
+    Handles both English (word count) and Chinese (character count) differently.
+
+    Args:
+        text_memories: List of text memories to be filtered
+        min_length_threshold: Minimum length required to keep a memory.
+                            For English: word count, for Chinese: character count.
+
+    Returns:
+        List of filtered memories meeting the length requirement
+    """
+    if not text_memories:
+        logger.debug("Empty memories list received in short memory filter")
+        return []
+
+    filtered_memories = []
+    removed_count = 0
+
+    for memory in text_memories:
+        stripped_memory = memory.strip()
+        if not stripped_memory:  # Skip empty/whitespace memories
+            removed_count += 1
+            continue
+
+        # Determine measurement method based on language
+        if is_all_english(stripped_memory):
+            length = len(stripped_memory.split())  # Word count for English
+        elif is_all_chinese(stripped_memory):
+            length = len(stripped_memory)  # Character count for Chinese
+        else:
+            logger.debug(f"Mixed-language memory, using character count: {stripped_memory[:50]}...")
+            length = len(stripped_memory)  # Default to character count
+
+        if length >= min_length_threshold:
+            filtered_memories.append(memory)
+        else:
+            removed_count += 1
+
+    if removed_count > 0:
+        logger.info(
+            f"Filtered out {removed_count} short memories "
+            f"(below {min_length_threshold} units). "
+            f"Total remaining: {len(filtered_memories)}"
+        )
+
+    return filtered_memories

memos/mem_scheduler/utils/metrics.py

@@ -0,0 +1,125 @@
+# src/memos/mem_scheduler/utils/metrics.py
+import time
+
+from contextlib import ContextDecorator
+
+from prometheus_client import Counter, Gauge, Histogram, Summary
+
+
+# --- Metric Definitions ---
+
+TASKS_ENQUEUED_TOTAL = Counter(
+    "memos_scheduler_tasks_enqueued_total",
+    "Total number of tasks enqueued",
+    ["user_id", "task_type"],
+)
+
+TASKS_DEQUEUED_TOTAL = Counter(
+    "memos_scheduler_tasks_dequeued_total",
+    "Total number of tasks dequeued",
+    ["user_id", "task_type"],
+)
+
+TASK_DURATION_SECONDS = Summary(
+    "memos_scheduler_task_duration_seconds",
+    "Task processing duration in seconds",
+    ["user_id", "task_type"],
+)
+
+TASK_WAIT_DURATION_SECONDS = Summary(
+    "memos_scheduler_task_wait_duration_seconds",
+    "Task waiting duration in seconds",
+    ["user_id", "task_type"],
+)
+
+TASKS_FAILED_TOTAL = Counter(
+    "memos_scheduler_tasks_failed_total",
+    "Total number of failed tasks",
+    ["user_id", "task_type", "error_type"],
+)
+
+TASKS_COMPLETED_TOTAL = Counter(
+    "memos_scheduler_tasks_completed_total",
+    "Total number of successfully completed tasks",
+    ["user_id", "task_type"],
+)
+
+QUEUE_LENGTH = Gauge(
+    "memos_scheduler_queue_length", "Current length of the task queue", ["user_id"]
+)
+
+INTERNAL_SPAN_DURATION = Histogram(
+    "memos_scheduler_internal_span_duration_seconds",
+    "Duration of internal operations",
+    ["span_name", "user_id", "task_id"],
+)
+
+
+# --- Instrumentation Functions ---
+
+
+def task_enqueued(user_id: str, task_type: str, count: int = 1):
+    TASKS_ENQUEUED_TOTAL.labels(user_id=user_id, task_type=task_type).inc(count)
+
+
+def task_dequeued(user_id: str, task_type: str, count: int = 1):
+    TASKS_DEQUEUED_TOTAL.labels(user_id=user_id, task_type=task_type).inc(count)
+
+
+def observe_task_duration(duration: float, user_id: str, task_type: str):
+    TASK_DURATION_SECONDS.labels(user_id=user_id, task_type=task_type).observe(duration)
+
+
+def observe_task_wait_duration(duration: float, user_id: str, task_type: str):
+    TASK_WAIT_DURATION_SECONDS.labels(user_id=user_id, task_type=task_type).observe(duration)
+
+
+def task_failed(user_id: str, task_type: str, error_type: str):
+    TASKS_FAILED_TOTAL.labels(user_id=user_id, task_type=task_type, error_type=error_type).inc()
+
+
+def task_completed(user_id: str, task_type: str, count: int = 1):
+    TASKS_COMPLETED_TOTAL.labels(user_id=user_id, task_type=task_type).inc(count)
+
+
+def update_queue_length(length: int, user_id: str):
+    QUEUE_LENGTH.labels(user_id=user_id).set(length)
+
+
+def observe_internal_span(duration: float, span_name: str, user_id: str, task_id: str):
+    INTERNAL_SPAN_DURATION.labels(span_name=span_name, user_id=user_id, task_id=task_id).observe(
+        duration
+    )
+
+
+# --- TimingSpan Context Manager ---
+
+
+class TimingSpan(ContextDecorator):
+    """
+    A context manager/decorator to measure the duration of a code block and record it
+    as a Prometheus histogram observation.
+
+    Usage as a decorator:
+    @TimingSpan("expensive_operation", user_id="user123")
+    def my_function():
+        time.sleep(2)
+
+    Usage as a context manager:
+    with TimingSpan("another_op", user_id="user456", task_id="t1"):
+        ...
+    """
+
+    def __init__(self, span_name: str, user_id: str = "unknown", task_id: str = "unknown"):
+        self.span_name = span_name
+        self.user_id = user_id
+        self.task_id = task_id
+        self.start_time = 0
+
+    def __enter__(self):
+        self.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        duration = time.perf_counter() - self.start_time
+        observe_internal_span(duration, self.span_name, self.user_id, self.task_id)