MemoryOS 2.0.3__py3-none-any.whl

memos/mem_scheduler/orm_modules/api_redis_model.py

@@ -0,0 +1,518 @@
+import os
+import time
+
+from typing import Any
+
+from sqlalchemy.orm import declarative_base
+
+from memos.log import get_logger
+from memos.mem_scheduler.orm_modules.base_model import DatabaseError
+from memos.mem_scheduler.schemas.api_schemas import (
+    APISearchHistoryManager,
+)
+from memos.mem_scheduler.utils.db_utils import get_utc_now
+
+
+logger = get_logger(__name__)
+
+Base = declarative_base()
+
+
+class APIRedisDBManager:
+    """Redis-based database manager for any serializable object
+
+    This class handles persistence, synchronization, and locking
+    for any object that implements to_json/from_json methods using Redis as the backend storage.
+    """
+
+    # Add orm_class attribute for compatibility
+    orm_class = None
+
+    def __init__(
+        self,
+        user_id: str | None = None,
+        mem_cube_id: str | None = None,
+        obj: APISearchHistoryManager | None = None,
+        lock_timeout: int = 10,
+        redis_client=None,
+        redis_config: dict | None = None,
+        window_size: int = 5,
+    ):
+        """Initialize the Redis database manager
+
+        Args:
+            user_id: Unique identifier for the user
+            mem_cube_id: Unique identifier for the memory cube
+            obj: Optional object instance to manage (must have to_json/from_json methods)
+            lock_timeout: Timeout in seconds for lock acquisition
+            redis_client: Redis client instance (optional)
+            redis_config: Redis configuration dictionary (optional)
+        """
+        # Initialize Redis client
+        self.redis_client = redis_client
+        self.redis_config = redis_config or {}
+
+        if self.redis_client is None:
+            self._init_redis_client()
+
+        # Initialize base attributes without calling parent's init_manager
+        self.user_id = user_id
+        self.mem_cube_id = mem_cube_id
+        self.obj = obj
+        self.lock_timeout = lock_timeout
+        self.engine = None  # Keep for compatibility but not used
+        self.SessionLocal = None  # Not used for Redis
+        self.window_size = window_size
+        self.lock_key = f"{self._get_key_prefix()}:lock"
+
+        logger.info(
+            f"RedisDBManager initialized for user_id: {user_id}, mem_cube_id: {mem_cube_id}"
+        )
+        logger.info(f"Redis client: {type(self.redis_client).__name__}")
+
+        # Test Redis connection
+        try:
+            self.redis_client.ping()
+            logger.info("Redis connection successful")
+        except Exception as e:
+            logger.warning(f"Redis ping failed: {e}")
+            # Don't raise error here as it might be a mock client in tests
+
+    def _get_key_prefix(self) -> str:
+        """Generate Redis key prefix for this user and memory cube
+
+        Returns:
+            Redis key prefix string
+        """
+        return f"redis_api:{self.user_id}:{self.mem_cube_id}"
+
+    def _get_data_key(self) -> str:
+        """Generate Redis key for storing serialized data
+
+        Returns:
+            Redis data key string
+        """
+        return f"{self._get_key_prefix()}:data"
+
+    def _init_redis_client(self):
+        """Initialize Redis client from config or environment"""
+        try:
+            import redis
+        except ImportError:
+            logger.error("Redis package not installed. Install with: pip install redis")
+            raise
+
+        # Try to get Redis client from environment first
+        if not self.redis_client:
+            self.redis_client = APIRedisDBManager.load_redis_engine_from_env()
+
+        # If still no client, try from config
+        if not self.redis_client and self.redis_config:
+            redis_kwargs = {
+                "host": self.redis_config.get("host"),
+                "port": self.redis_config.get("port"),
+                "db": self.redis_config.get("db"),
+                "decode_responses": True,
+            }
+
+            if self.redis_config.get("password"):
+                redis_kwargs["password"] = self.redis_config["password"]
+
+            self.redis_client = redis.Redis(**redis_kwargs)
+
+        # Final fallback to localhost
+        if not self.redis_client:
+            logger.warning("No Redis configuration found, using localhost defaults")
+            self.redis_client = redis.Redis(
+                host="localhost", port=6379, db=0, decode_responses=True
+            )
+
+        # Test connection
+        if not self.redis_client.ping():
+            raise ConnectionError("Redis ping failed")
+
+        logger.info("Redis client initialized successfully")
+
+    def acquire_lock(self, block: bool = True, **kwargs) -> bool:
+        """Acquire a distributed lock using Redis with atomic operations
+
+        Args:
+            block: Whether to block until lock is acquired
+            **kwargs: Additional filter criteria (ignored for Redis)
+
+        Returns:
+            True if lock was acquired, False otherwise
+        """
+
+        now = get_utc_now()
+
+        # Use Redis SET with NX (only if not exists) and EX (expiry) for atomic lock acquisition
+        lock_value = f"{self._get_key_prefix()}:{now.timestamp()}"
+
+        while True:
+            result = self.redis_client.get(self.lock_key)
+            if result:
+                # Wait a bit before retrying
+                logger.info(
+                    f"Waiting for Redis lock to be released for {self.user_id}/{self.mem_cube_id}"
+                )
+                if not block:
+                    logger.warning(
+                        f"Redis lock is held for {self.user_id}/{self.mem_cube_id}, cannot acquire"
+                    )
+                    return False
+                else:
+                    time.sleep(0.1)
+                    continue
+            else:
+                # Try to acquire lock atomically
+                result = self.redis_client.set(
+                    self.lock_key,
+                    lock_value,
+                    ex=self.lock_timeout,  # Set expiry in seconds
+                )
+                logger.info(f"Redis lock acquired for {self._get_key_prefix()}")
+                return True
+
+    def release_locks(self, **kwargs):
+        # Delete the lock key to release the lock
+        result = self.redis_client.delete(self.lock_key)
+
+        # Redis DELETE returns the number of keys deleted (0 or 1)
+        if result > 0:
+            logger.info(f"Redis lock released for {self._get_key_prefix()}")
+        else:
+            logger.info(f"No Redis lock found to release for {self._get_key_prefix()}")
+
+    def merge_items(
+        self,
+        redis_data: str,
+        obj_instance: APISearchHistoryManager,
+        size_limit: int,
+    ):
+        """Merge Redis data with current object instance
+
+        Args:
+            redis_data: JSON string from Redis containing serialized APISearchHistoryManager
+            obj_instance: Current APISearchHistoryManager instance
+            size_limit: Maximum number of completed entries to keep
+
+        Returns:
+            APISearchHistoryManager: Merged and synchronized manager instance
+        """
+
+        # Parse Redis data
+        redis_manager = APISearchHistoryManager.from_json(redis_data)
+        logger.debug(
+            f"Loaded Redis manager with {len(redis_manager.completed_entries)} completed and {len(redis_manager.running_item_ids)} running task IDs"
+        )
+
+        # Create a new merged manager with the original window size from obj_instance
+        # Use size_limit only for limiting entries, not as window_size
+        original_window_size = obj_instance.window_size
+        merged_manager = APISearchHistoryManager(window_size=original_window_size)
+
+        # Merge completed entries - combine both sources and deduplicate by task_id
+        # Ensure all entries are APIMemoryHistoryEntryItem instances
+        from memos.mem_scheduler.schemas.api_schemas import APIMemoryHistoryEntryItem
+
+        all_completed = {}
+
+        # Add Redis completed entries
+        for entry in redis_manager.completed_entries:
+            if isinstance(entry, dict):
+                # Convert dict to APIMemoryHistoryEntryItem instance
+                try:
+                    entry_obj = APIMemoryHistoryEntryItem(**entry)
+                    task_id = entry_obj.item_id
+                    all_completed[task_id] = entry_obj
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to convert dict entry to APIMemoryHistoryEntryItem: {e}"
+                    )
+                    continue
+            else:
+                task_id = entry.item_id
+                all_completed[task_id] = entry
+
+        # Add current instance completed entries (these take priority if duplicated)
+        for entry in obj_instance.completed_entries:
+            if isinstance(entry, dict):
+                # Convert dict to APIMemoryHistoryEntryItem instance
+                try:
+                    entry_obj = APIMemoryHistoryEntryItem(**entry)
+                    task_id = entry_obj.item_id
+                    all_completed[task_id] = entry_obj
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to convert dict entry to APIMemoryHistoryEntryItem: {e}"
+                    )
+                    continue
+            else:
+                task_id = entry.item_id
+                all_completed[task_id] = entry
+
+        # Sort by created_time and apply size limit
+        completed_list = list(all_completed.values())
+
+        def get_created_time(entry):
+            """Helper function to safely extract created_time for sorting"""
+            from datetime import datetime
+
+            # All entries should now be APIMemoryHistoryEntryItem instances
+            return getattr(entry, "created_time", datetime.min)
+
+        completed_list.sort(key=get_created_time, reverse=True)
+        merged_manager.completed_entries = completed_list[:size_limit]
+
+        # Merge running task IDs - combine both sources and deduplicate
+        all_running_item_ids = set()
+
+        # Add Redis running task IDs
+        all_running_item_ids.update(redis_manager.running_item_ids)
+
+        # Add current instance running task IDs
+        all_running_item_ids.update(obj_instance.running_item_ids)
+
+        merged_manager.running_item_ids = list(all_running_item_ids)
+
+        logger.info(
+            f"Merged manager: {len(merged_manager.completed_entries)} completed, {len(merged_manager.running_item_ids)} running task IDs"
+        )
+        return merged_manager
+
+    def sync_with_redis(self, size_limit: int | None = None) -> None:
+        """Synchronize data between Redis and the business object
+
+        Args:
+            size_limit: Optional maximum number of items to keep after synchronization
+        """
+
+        # Use window_size from the object if size_limit is not provided
+        if size_limit is None:
+            size_limit = self.window_size
+
+        # Acquire lock before operations
+        lock_status = self.acquire_lock(block=True)
+        if not lock_status:
+            logger.error("Failed to acquire Redis lock for synchronization")
+            return
+
+        # Load existing data from Redis
+        data_key = self._get_data_key()
+        redis_data = self.redis_client.get(data_key)
+
+        if redis_data:
+            # Merge Redis data with current object
+            merged_obj = self.merge_items(
+                redis_data=redis_data, obj_instance=self.obj, size_limit=size_limit
+            )
+
+            # Update the current object with merged data
+            self.obj = merged_obj
+            logger.info(
+                f"Successfully synchronized with Redis data for {self.user_id}/{self.mem_cube_id}"
+            )
+        else:
+            logger.info(
+                f"No existing Redis data found for {self.user_id}/{self.mem_cube_id}, using current object"
+            )
+
+        # Save the synchronized object back to Redis
+        self.save_to_db(self.obj)
+
+        self.release_locks()
+
+    def save_to_db(self, obj_instance: Any) -> None:
+        """Save the current state of the business object to Redis
+
+        Args:
+            obj_instance: The object instance to save (must have to_json method)
+        """
+
+        data_key = self._get_data_key()
+
+        self.redis_client.set(data_key, obj_instance.to_json())
+
+        logger.info(f"Updated existing Redis record for {data_key}")
+
+    def load_from_db(self) -> Any | None:
+        data_key = self._get_data_key()
+
+        # Load from Redis
+        serialized_data = self.redis_client.get(data_key)
+
+        if not serialized_data:
+            logger.info(f"No Redis record found for {data_key}")
+            return None
+
+        # Deserialize the business object using the actual object type
+        if hasattr(self, "obj_type") and self.obj_type is not None:
+            db_instance = self.obj_type.from_json(serialized_data)
+        else:
+            # Default to APISearchHistoryManager for this class
+            db_instance = APISearchHistoryManager.from_json(serialized_data)
+
+        logger.info(f"Successfully loaded object from Redis for {data_key} ")
+
+        return db_instance
+
+    @classmethod
+    def from_env(
+        cls,
+        user_id: str,
+        mem_cube_id: str,
+        obj: Any | None = None,
+        lock_timeout: int = 10,
+        env_file_path: str | None = None,
+    ) -> "APIRedisDBManager":
+        """Create RedisDBManager from environment variables
+
+        Args:
+            user_id: User identifier
+            mem_cube_id: Memory cube identifier
+            obj: Optional MemoryMonitorManager instance
+            lock_timeout: Lock timeout in seconds
+            env_file_path: Optional path to .env file
+
+        Returns:
+                RedisDBManager instance
+        """
+
+        redis_client = APIRedisDBManager.load_redis_engine_from_env(env_file_path)
+        return cls(
+            user_id=user_id,
+            mem_cube_id=mem_cube_id,
+            obj=obj,
+            lock_timeout=lock_timeout,
+            redis_client=redis_client,
+        )
+
+    def close(self):
+        """Close the Redis connection and clean up resources"""
+        try:
+            if hasattr(self.redis_client, "close"):
+                self.redis_client.close()
+            logger.info(
+                f"Redis connection closed for user_id: {self.user_id}, mem_cube_id: {self.mem_cube_id}"
+            )
+        except Exception as e:
+            logger.warning(f"Error closing Redis connection: {e}")
+
+    @staticmethod
+    def load_redis_engine_from_env(env_file_path: str | None = None) -> Any:
+        """Load Redis connection from environment variables
+
+        Args:
+            env_file_path: Path to .env file (optional, defaults to loading from current environment)
+
+        Returns:
+            Redis connection instance
+
+        Raises:
+            DatabaseError: If required environment variables are missing or connection fails
+        """
+        try:
+            import redis
+        except ImportError as e:
+            error_msg = "Redis package not installed. Install with: pip install redis"
+            logger.error(error_msg)
+            raise DatabaseError(error_msg) from e
+
+        # Load environment variables from file if provided
+        if env_file_path:
+            if os.path.exists(env_file_path):
+                from dotenv import load_dotenv
+
+                load_dotenv(env_file_path)
+                logger.info(f"Loaded environment variables from {env_file_path}")
+            else:
+                logger.warning(
+                    f"Environment file not found: {env_file_path}, using current environment variables",
+                    stack_info=True,
+                )
+        else:
+            logger.info("Using current environment variables (no env_file_path provided)")
+
+        # Get Redis configuration from environment variables
+        redis_host = os.getenv("REDIS_HOST") or os.getenv("MEMSCHEDULER_REDIS_HOST")
+        redis_port_str = os.getenv("REDIS_PORT") or os.getenv("MEMSCHEDULER_REDIS_PORT")
+        redis_db_str = os.getenv("REDIS_DB") or os.getenv("MEMSCHEDULER_REDIS_DB")
+        redis_password = os.getenv("REDIS_PASSWORD") or os.getenv("MEMSCHEDULER_REDIS_PASSWORD")
+
+        # Check required environment variables
+        if not redis_host:
+            error_msg = (
+                "Missing required Redis environment variable: REDIS_HOST or MEMSCHEDULER_REDIS_HOST"
+            )
+            logger.error(error_msg)
+            return None
+
+        # Parse port with validation
+        try:
+            redis_port = int(redis_port_str) if redis_port_str else 6379
+        except ValueError:
+            error_msg = f"Invalid REDIS_PORT value: {redis_port_str}. Must be a valid integer."
+            logger.error(error_msg)
+            return None
+
+        # Parse database with validation
+        try:
+            redis_db = int(redis_db_str) if redis_db_str else 0
+        except ValueError:
+            error_msg = f"Invalid REDIS_DB value: {redis_db_str}. Must be a valid integer."
+            logger.error(error_msg)
+            return None
+
+        # Optional timeout settings
+        socket_timeout = os.getenv(
+            "REDIS_SOCKET_TIMEOUT", os.getenv("MEMSCHEDULER_REDIS_TIMEOUT", None)
+        )
+        socket_connect_timeout = os.getenv(
+            "REDIS_SOCKET_CONNECT_TIMEOUT", os.getenv("MEMSCHEDULER_REDIS_CONNECT_TIMEOUT", None)
+        )
+
+        try:
+            # Build Redis connection parameters
+            redis_kwargs = {
+                "host": redis_host,
+                "port": redis_port,
+                "db": redis_db,
+                "decode_responses": True,
+            }
+
+            if redis_password:
+                redis_kwargs["password"] = redis_password
+
+            if socket_timeout:
+                try:
+                    redis_kwargs["socket_timeout"] = float(socket_timeout)
+                except ValueError:
+                    logger.warning(
+                        f"Invalid REDIS_SOCKET_TIMEOUT value: {socket_timeout}, ignoring"
+                    )
+
+            if socket_connect_timeout:
+                try:
+                    redis_kwargs["socket_connect_timeout"] = float(socket_connect_timeout)
+                except ValueError:
+                    logger.warning(
+                        f"Invalid REDIS_SOCKET_CONNECT_TIMEOUT value: {socket_connect_timeout}, ignoring"
+                    )
+
+            # Create Redis connection
+            redis_client = redis.Redis(**redis_kwargs)
+
+            # Test connection
+            if not redis_client.ping():
+                raise ConnectionError("Redis ping failed")
+
+            logger.info(
+                f"Successfully created Redis connection: {redis_host}:{redis_port}/{redis_db}"
+            )
+            return redis_client
+
+        except Exception as e:
+            error_msg = f"Failed to create Redis connection from environment variables: {e}"
+            logger.error(error_msg, stack_info=True)
+            raise DatabaseError(error_msg) from e