MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/utils/misc_utils.py

@@ -0,0 +1,290 @@
+import json
+import os
+import re
+import traceback
+
+from collections import defaultdict
+from functools import wraps
+from pathlib import Path
+
+import yaml
+
+from memos.log import get_logger
+from memos.mem_scheduler.schemas.message_schemas import (
+    ScheduleMessageItem,
+)
+
+
+logger = get_logger(__name__)
+
+
+def _normalize_env_value(value: str | None) -> str:
+    """Normalize environment variable values for comparison."""
+    return value.strip().lower() if isinstance(value, str) else ""
+
+
+def is_playground_env() -> bool:
+    """Return True when ENV_NAME indicates a Playground environment."""
+    env_name = _normalize_env_value(os.getenv("ENV_NAME"))
+    return env_name.startswith("playground")
+
+
+def is_cloud_env() -> bool:
+    """
+    Determine whether the scheduler should treat the runtime as a cloud environment.
+
+    Rules:
+    - Any Playground ENV_NAME is explicitly NOT cloud.
+    - MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME must be set to enable cloud behavior.
+    - The default memos-fanout/fanout combination is treated as non-cloud.
+    """
+    if is_playground_env():
+        return False
+
+    exchange_name = _normalize_env_value(os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_NAME"))
+    exchange_type = _normalize_env_value(os.getenv("MEMSCHEDULER_RABBITMQ_EXCHANGE_TYPE"))
+
+    if not exchange_name:
+        return False
+
+    return not (
+        exchange_name == "memos-fanout" and (not exchange_type or exchange_type == "fanout")
+    )
+
+
+def extract_json_obj(text: str):
+    """
+    Safely extracts JSON from LLM response text with robust error handling.
+
+    Args:
+        text: Raw text response from LLM that may contain JSON
+
+    Returns:
+        Parsed JSON data (dict or list)
+
+    Raises:
+        ValueError: If no valid JSON can be extracted
+    """
+    if not text:
+        raise ValueError("Empty input text")
+
+    # Normalize the text
+    text = text.strip()
+
+    # Remove common code block markers
+    patterns_to_remove = ["json```", "```python", "```json", "latex```", "```latex", "```"]
+    for pattern in patterns_to_remove:
+        text = text.replace(pattern, "")
+
+    # Try: direct JSON parse first
+    try:
+        return json.loads(text.strip())
+    except json.JSONDecodeError as e:
+        logger.info(f"Failed to parse JSON from text: {text}. Error: {e!s}", exc_info=True)
+
+    # Fallback 1: Extract JSON using regex
+    json_pattern = r"\{[\s\S]*\}|\[[\s\S]*\]"
+    matches = re.findall(json_pattern, text)
+    if matches:
+        try:
+            return json.loads(matches[0])
+        except json.JSONDecodeError as e:
+            logger.info(f"Failed to parse JSON from text: {text}. Error: {e!s}", exc_info=True)
+
+    # Fallback 2: Handle malformed JSON (common LLM issues)
+    try:
+        # Try adding missing quotes around keys
+        text = re.sub(r"([\{\s,])(\w+)(:)", r'\1"\2"\3', text)
+        return json.loads(text)
+    except json.JSONDecodeError as e:
+        logger.error(f"Failed to parse JSON from text: {text}. Error: {e!s}")
+        logger.error("Full traceback:\n" + traceback.format_exc())
+        raise ValueError(text) from e
+
+
+def extract_list_items(text: str, bullet_prefixes: tuple[str, ...] = ("- ",)) -> list[str]:
+    """
+    Extract bullet list items from LLM output where each item is on a single line
+    starting with a given bullet prefix (default: "- ").
+
+    This function is designed to be robust to common LLM formatting variations,
+    following similar normalization practices as `extract_json_obj`.
+
+    Behavior:
+    - Strips common code-fence markers (```json, ```python, ``` etc.).
+    - Collects all lines that start with any of the provided `bullet_prefixes`.
+    - Tolerates the "• " bullet as a loose fallback.
+    - Unescapes common sequences like "\\n" and "\\t" within items.
+    - If no bullet lines are found, falls back to attempting to parse a JSON array
+      (using `extract_json_obj`) and returns its string elements.
+
+    Args:
+        text: Raw text response from LLM.
+        bullet_prefixes: Tuple of accepted bullet line prefixes.
+
+    Returns:
+        List of extracted items (strings). Returns an empty list if none can be parsed.
+    """
+    if not text:
+        return []
+
+    # Normalize the text similar to extract_json_obj
+    normalized = text.strip()
+    patterns_to_remove = ["json```", "```python", "```json", "latex```", "```latex", "```"]
+    for pattern in patterns_to_remove:
+        normalized = normalized.replace(pattern, "")
+    normalized = normalized.replace("\r\n", "\n")
+
+    lines = normalized.splitlines()
+    items: list[str] = []
+    seen: set[str] = set()
+
+    for raw in lines:
+        line = raw.strip()
+        if not line:
+            continue
+
+        matched = False
+        for prefix in bullet_prefixes:
+            if line.startswith(prefix):
+                content = line[len(prefix) :].strip()
+                content = content.replace("\\n", "\n").replace("\\t", "\t").replace("\\r", "\r")
+                if content and content not in seen:
+                    items.append(content)
+                    seen.add(content)
+                matched = True
+                break
+
+        if matched:
+            continue
+
+    if items:
+        return items
+    else:
+        logger.error(f"Fail to parse {text}")
+
+    return []
+
+
+def extract_list_items_in_answer(
+    text: str, bullet_prefixes: tuple[str, ...] = ("- ",)
+) -> list[str]:
+    """
+    Extract list items specifically from content enclosed within `<answer>...</answer>` tags.
+
+    - When one or more `<answer>...</answer>` blocks are present, concatenates their inner
+      contents with newlines and parses using `extract_list_items`.
+    - When no `<answer>` block is found, falls back to parsing the entire input with
+      `extract_list_items`.
+    - Case-insensitive matching of the `<answer>` tag.
+
+    Args:
+        text: Raw text that may contain `<answer>...</answer>` blocks.
+        bullet_prefixes: Accepted bullet prefixes (default: strictly `"- "`).
+
+    Returns:
+        List of extracted items (strings), or an empty list when nothing is parseable.
+    """
+    if not text:
+        return []
+
+    try:
+        normalized = text.strip().replace("\r\n", "\n")
+        # Ordered, exact-case matching for <answer> blocks: answer -> Answer -> ANSWER
+        tag_variants = ["answer", "Answer", "ANSWER"]
+        matches: list[str] = []
+        for tag in tag_variants:
+            matches = re.findall(rf"<{tag}>([\\s\\S]*?)</{tag}>", normalized)
+            if matches:
+                break
+        # Fallback: case-insensitive matching if none of the exact-case variants matched
+        if not matches:
+            matches = re.findall(r"<answer>([\\s\\S]*?)</answer>", normalized, flags=re.IGNORECASE)
+
+        if matches:
+            combined = "\n".join(m.strip() for m in matches if m is not None)
+            return extract_list_items(combined, bullet_prefixes=bullet_prefixes)
+
+        # Fallback: parse the whole text if tags are absent
+        return extract_list_items(normalized, bullet_prefixes=bullet_prefixes)
+    except Exception as e:
+        logger.info(f"Failed to extract items within <answer> tags: {e!s}", exc_info=True)
+        # Final fallback: attempt direct list extraction
+        try:
+            return extract_list_items(text, bullet_prefixes=bullet_prefixes)
+        except Exception:
+            return []
+
+
+def parse_yaml(yaml_file: str | Path):
+    yaml_path = Path(yaml_file)
+    if not yaml_path.is_file():
+        raise FileNotFoundError(f"No such file: {yaml_file}")
+
+    with yaml_path.open("r", encoding="utf-8") as fr:
+        data = yaml.safe_load(fr)
+
+    return data
+
+
+def log_exceptions(logger=logger):
+    """
+    Exception-catching decorator that automatically logs errors (including stack traces)
+
+    Args:
+        logger: Optional logger object (default: module-level logger)
+
+    Example:
+        @log_exceptions()
+        def risky_function():
+            raise ValueError("Oops!")
+
+        @log_exceptions(logger=custom_logger)
+        def another_risky_function():
+            might_fail()
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                logger.error(f"Error in {func.__name__}: {e}", stack_info=True)
+
+        return wrapper
+
+    return decorator
+
+
+def group_messages_by_user_and_mem_cube(
+    messages: list[ScheduleMessageItem],
+) -> dict[str, dict[str, list[ScheduleMessageItem]]]:
+    """
+    Groups messages into a nested dictionary structure first by user_id, then by mem_cube_id.
+
+    Args:
+        messages: List of ScheduleMessageItem objects to be grouped
+
+    Returns:
+        A nested dictionary with the structure:
+        {
+            "user_id_1": {
+                "mem_cube_id_1": [msg1, msg2, ...],
+                "mem_cube_id_2": [msg3, msg4, ...],
+                ...
+            },
+            "user_id_2": {
+                ...
+            },
+            ...
+        }
+        Where each msg is the original ScheduleMessageItem object
+    """
+    grouped_dict = defaultdict(lambda: defaultdict(list))
+
+    for msg in messages:
+        grouped_dict[msg.user_id][msg.mem_cube_id].append(msg)
+
+    # Convert defaultdict to regular dict for cleaner output
+    return {user_id: dict(cube_groups) for user_id, cube_groups in grouped_dict.items()}

memos/mem_scheduler/utils/monitor_event_utils.py

@@ -0,0 +1,67 @@
+import json
+import os
+import socket
+
+from datetime import datetime, timezone
+from typing import Any
+
+from memos.log import get_logger
+from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
+
+
+logger = get_logger(__name__)
+
+
+def _iso_ts_now() -> str:
+    """Return current UTC timestamp in ISO format with milliseconds."""
+    return datetime.now(timezone.utc).isoformat()
+
+
+def to_iso(ts) -> str | None:
+    """Convert datetime to ISO string; return None if not convertible."""
+    if ts is None:
+        return None
+    if isinstance(ts, datetime):
+        dt = ts
+        if dt.tzinfo is None:
+            dt = dt.replace(tzinfo=timezone.utc)
+        return dt.isoformat()
+    try:
+        return datetime.fromtimestamp(float(ts), tz=timezone.utc).isoformat()
+    except Exception:
+        return None
+
+
+def emit_monitor_event(event: str, msg: ScheduleMessageItem, extra: dict[str, Any] | None = None):
+    """
+    Emit a structured MONITOR_EVENT log line for SLS consumption.
+
+    This must be fire-and-forget: any exception here should never break the scheduler flow.
+    """
+    try:
+        payload: dict[str, Any] = {
+            "event": event,
+            "ts": _iso_ts_now(),
+            "label": getattr(msg, "label", None),
+            "user_id": getattr(msg, "user_id", None),
+            "mem_cube_id": getattr(msg, "mem_cube_id", None),
+            "item_id": getattr(msg, "item_id", None),
+            "task_id": getattr(msg, "task_id", "") or "",
+            "trace_id": getattr(msg, "trace_id", None),
+            "stream_key": getattr(msg, "stream_key", None),
+            "redis_message_id": getattr(msg, "redis_message_id", None),
+            "monitor_flag": None,
+            "host": socket.gethostname(),
+            "env": os.getenv("ENV") or os.getenv("ENVIRONMENT") or "",
+        }
+
+        info = getattr(msg, "info", None)
+        if isinstance(info, dict):
+            payload["monitor_flag"] = info.get("monitor_flag")
+
+        if extra:
+            payload.update(extra)
+
+        logger.info("MONITOR_EVENT " + json.dumps(payload, ensure_ascii=False))
+    except Exception:
+        logger.debug("Failed to emit MONITOR_EVENT", exc_info=True)

memos/mem_scheduler/utils/status_tracker.py

@@ -0,0 +1,229 @@
+# src/memos/mem_scheduler/utils/status_tracker.py
+import json
+
+from datetime import datetime, timedelta, timezone
+from typing import TYPE_CHECKING
+
+from memos.dependency import require_python_package
+
+
+if TYPE_CHECKING:
+    import redis
+
+
+class TaskStatusTracker:
+    @require_python_package(import_name="redis", install_command="pip install redis")
+    def __init__(self, redis_client: "redis.Redis | None"):
+        self.redis = redis_client
+
+    def _get_key(self, user_id: str) -> str:
+        if not self.redis:
+            return
+
+        return f"memos:task_meta:{user_id}"
+
+    def _get_task_items_key(self, user_id: str, task_id: str) -> str:
+        """Get Redis key for task_id → [item_id] mapping."""
+        return f"memos:task_items:{user_id}:{task_id}"
+
+    def task_submitted(
+        self,
+        task_id: str,
+        user_id: str,
+        task_type: str,
+        mem_cube_id: str,
+        business_task_id: str | None = None,
+    ):
+        """
+        Submit a new task for tracking.
+
+        Args:
+            task_id: Internal item_id (UUID)
+            user_id: User identifier
+            task_type: Type of task (label)
+            mem_cube_id: Memory cube identifier
+            business_task_id: Optional business-level task ID (one task_id can have multiple item_ids)
+        """
+        if not self.redis:
+            return
+
+        key = self._get_key(user_id)
+        payload = {
+            "status": "waiting",
+            "task_type": task_type,
+            "mem_cube_id": mem_cube_id,
+            "submitted_at": datetime.now(timezone.utc).isoformat(),
+        }
+
+        # Add business_task_id to payload if provided
+        if business_task_id:
+            payload["business_task_id"] = business_task_id
+            # Add item_id to the task_id → [item_ids] set
+            task_items_key = self._get_task_items_key(user_id, business_task_id)
+            self.redis.sadd(task_items_key, task_id)
+            self.redis.expire(task_items_key, timedelta(days=7))
+
+        self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
+
+    def task_started(self, task_id: str, user_id: str):
+        if not self.redis:
+            return
+
+        key = self._get_key(user_id)
+        existing_data_json = self.redis.hget(key, task_id)
+        if not existing_data_json:
+            # 容错处理: 如果任务不存在, 也创建一个
+            payload = {
+                "status": "in_progress",
+                "started_at": datetime.now(timezone.utc).isoformat(),
+            }
+        else:
+            payload = json.loads(existing_data_json)
+            payload["status"] = "in_progress"
+            payload["started_at"] = datetime.now(timezone.utc).isoformat()
+        self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
+
+    def task_completed(self, task_id: str, user_id: str):
+        if not self.redis:
+            return
+
+        key = self._get_key(user_id)
+        existing_data_json = self.redis.hget(key, task_id)
+        if not existing_data_json:
+            return
+        payload = json.loads(existing_data_json)
+        payload["status"] = "completed"
+        payload["completed_at"] = datetime.now(timezone.utc).isoformat()
+        # 设置该任务条目的过期时间, 例如 24 小时
+        # 注意: Redis Hash 不能为单个 field 设置 TTL, 这里我们可以 通过后台任务清理或在获取时判断时间戳
+        # 简单起见, 我们暂时依赖一个后台清理任务
+        self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
+
+    def task_failed(self, task_id: str, user_id: str, error_message: str):
+        if not self.redis:
+            return
+
+        key = self._get_key(user_id)
+        existing_data_json = self.redis.hget(key, task_id)
+        if not existing_data_json:
+            payload = {
+                "status": "failed",
+                "error": error_message,
+                "failed_at": datetime.now(timezone.utc).isoformat(),
+            }
+        else:
+            payload = json.loads(existing_data_json)
+            payload["status"] = "failed"
+            payload["error"] = error_message
+            payload["failed_at"] = datetime.now(timezone.utc).isoformat()
+        self.redis.hset(key, task_id, json.dumps(payload))
+        self.redis.expire(key, timedelta(days=7))
+
+    def get_task_status(self, task_id: str, user_id: str) -> dict | None:
+        if not self.redis:
+            return None
+
+        key = self._get_key(user_id)
+        data = self.redis.hget(key, task_id)
+        return json.loads(data) if data else None
+
+    def get_all_tasks_for_user(self, user_id: str) -> dict[str, dict]:
+        if not self.redis:
+            return {}
+
+        key = self._get_key(user_id)
+        all_tasks = self.redis.hgetall(key)
+        return {tid: json.loads(t_data) for tid, t_data in all_tasks.items()}
+
+    def get_task_status_by_business_id(self, business_task_id: str, user_id: str) -> dict | None:
+        """
+        Get aggregated status for a business-level task_id.
+
+        Args:
+            business_task_id: Business-level task ID
+            user_id: User identifier
+
+        Returns:
+            Aggregated status dict with status determined by all item statuses:
+            - If any item is 'waiting' or 'in_progress' → 'in_progress'
+            - If all items are 'completed' → 'completed'
+            - If any item is 'failed' → 'failed'
+            Returns None if task_id not found.
+        """
+        if not self.redis:
+            return None
+
+        # Get all item_ids for this task_id
+        task_items_key = self._get_task_items_key(user_id, business_task_id)
+        item_ids = self.redis.smembers(task_items_key)
+
+        if not item_ids:
+            return None
+
+        # Get statuses for all items
+        key = self._get_key(user_id)
+        item_statuses = []
+        errors = []
+        for item_id in item_ids:
+            item_data_json = self.redis.hget(key, item_id)
+            if item_data_json:
+                item_data = json.loads(item_data_json)
+                item_statuses.append(item_data["status"])
+                if item_data.get("status") == "failed" and "error" in item_data:
+                    errors.append(item_data["error"])
+
+        if not item_statuses:
+            return None
+
+        # Aggregate status
+        if "failed" in item_statuses:
+            aggregated_status = "failed"
+        elif "in_progress" in item_statuses or "waiting" in item_statuses:
+            aggregated_status = "in_progress"
+        elif all(s == "completed" for s in item_statuses):
+            aggregated_status = "completed"
+        else:
+            # Fallback
+            aggregated_status = "unknown"
+
+        return {
+            "status": aggregated_status,
+            "business_task_id": business_task_id,
+            "item_count": len(item_ids),
+            "item_statuses": item_statuses,
+            "errors": errors,
+        }
+
+    def get_all_tasks_global(self) -> dict[str, dict[str, dict]]:
+        """
+        Retrieve all tasks for all users from Redis.
+
+        Returns:
+            dict: {user_id: {task_id: task_data, ...}, ...}
+        """
+        if not self.redis:
+            return {}
+
+        all_users_tasks = {}
+        cursor: int | str = 0
+        while True:
+            cursor, keys = self.redis.scan(cursor=cursor, match="memos:task_meta:*", count=100)
+            for key in keys:
+                # key format: memos:task_meta:{user_id}
+                parts = key.split(":")
+                if len(parts) < 3:
+                    continue
+                user_id = parts[2]
+
+                tasks = self.redis.hgetall(key)
+                if tasks:
+                    user_tasks = {tid: json.loads(t_data) for tid, t_data in tasks.items()}
+                    all_users_tasks[user_id] = user_tasks
+
+            if cursor == 0 or cursor == "0":
+                break
+
+        return all_users_tasks