MemoryOS 2.0.3__py3-none-any.whl

memos/mem_user/persistent_user_manager.py

@@ -0,0 +1,260 @@
+"""Persistent user management system for MemOS with configuration storage.
+
+This module extends the base UserManager to provide persistent storage
+for user configurations and MOS instances.
+"""
+
+import json
+
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy import Column, String, Text
+
+from memos.configs.mem_os import MOSConfig
+from memos.log import get_logger
+from memos.mem_user.user_manager import Base, UserManager
+
+
+logger = get_logger(__name__)
+
+
+class UserConfig(Base):
+    """User configuration model for the database."""
+
+    __tablename__ = "user_configs"
+
+    user_id = Column(String, primary_key=True)
+    config_data = Column(Text, nullable=False)  # JSON string of MOSConfig
+    created_at = Column(String, nullable=False)  # ISO format timestamp
+    updated_at = Column(String, nullable=False)  # ISO format timestamp
+
+    def __repr__(self):
+        return f"<UserConfig(user_id='{self.user_id}')>"
+
+
+class PersistentUserManager(UserManager):
+    """Extended UserManager with configuration persistence."""
+
+    def __init__(self, db_path: str | None = None, user_id: str = "root"):
+        """Initialize the persistent user manager.
+
+        Args:
+            db_path (str, optional): Path to the SQLite database file.
+                If None, uses default path in MEMOS_DIR.
+            user_id (str, optional): User ID. If None, uses default user ID.
+        """
+        super().__init__(db_path, user_id)
+
+        # Create user_configs table
+        Base.metadata.create_all(bind=self.engine)
+        logger.info("PersistentUserManager initialized with configuration storage")
+
+    def _convert_datetime_strings(self, obj: Any) -> Any:
+        """Recursively convert datetime strings back to datetime objects in config dict.
+
+        Args:
+            obj: The object to process (dict, list, or primitive type)
+
+        Returns:
+            The object with datetime strings converted to datetime objects
+        """
+        if isinstance(obj, dict):
+            result = {}
+            for key, value in obj.items():
+                if key == "created_at" and isinstance(value, str):
+                    try:
+                        result[key] = datetime.fromisoformat(value)
+                    except ValueError:
+                        # If parsing fails, keep the original string
+                        result[key] = value
+                else:
+                    result[key] = self._convert_datetime_strings(value)
+            return result
+        elif isinstance(obj, list):
+            return [self._convert_datetime_strings(item) for item in obj]
+        else:
+            return obj
+
+    def save_user_config(self, user_id: str, config: MOSConfig) -> bool:
+        """Save user configuration to database.
+
+        Args:
+            user_id (str): The user ID.
+            config (MOSConfig): The user's MOS configuration.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            # Convert config to JSON string with proper datetime handling
+            config_dict = config.model_dump(mode="json")
+            config_json = json.dumps(config_dict, indent=2)
+
+            from datetime import datetime
+
+            now = datetime.now().isoformat()
+
+            # Check if config already exists
+            existing_config = (
+                session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+            )
+
+            if existing_config:
+                # Update existing config
+                existing_config.config_data = config_json
+                existing_config.updated_at = now
+                logger.info(f"Updated configuration for user {user_id}")
+            else:
+                # Create new config
+                user_config = UserConfig(
+                    user_id=user_id, config_data=config_json, created_at=now, updated_at=now
+                )
+                session.add(user_config)
+                logger.info(f"Saved new configuration for user {user_id}")
+
+            session.commit()
+            return True
+
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error saving user config for {user_id}: {e}")
+            return False
+        finally:
+            session.close()
+
+    def get_user_config(self, user_id: str) -> MOSConfig | None:
+        """Get user configuration from database.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            MOSConfig | None: The user's configuration or None if not found.
+        """
+        session = self._get_session()
+        try:
+            user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+
+            if user_config:
+                config_dict = json.loads(user_config.config_data)
+                # Convert datetime strings back to datetime objects
+                config_dict = self._convert_datetime_strings(config_dict)
+                return MOSConfig(**config_dict)
+            return None
+
+        except Exception as e:
+            logger.error(f"Error loading user config for {user_id}: {e}")
+            return None
+        finally:
+            session.close()
+
+    def delete_user_config(self, user_id: str) -> bool:
+        """Delete user configuration from database.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        session = self._get_session()
+        try:
+            user_config = session.query(UserConfig).filter(UserConfig.user_id == user_id).first()
+
+            if user_config:
+                session.delete(user_config)
+                session.commit()
+                logger.info(f"Deleted configuration for user {user_id}")
+                return True
+            return False
+
+        except Exception as e:
+            session.rollback()
+            logger.error(f"Error deleting user config for {user_id}: {e}")
+            return False
+        finally:
+            session.close()
+
+    def list_user_configs(self, limit: int = 1) -> dict[str, MOSConfig]:
+        """List all user configurations.
+
+        Returns:
+            Dict[str, MOSConfig]: Dictionary mapping user_id to MOSConfig.
+        """
+        session = self._get_session()
+        try:
+            user_configs = session.query(UserConfig).limit(limit).all()
+            result = {}
+
+            for user_config in user_configs:
+                try:
+                    config_dict = json.loads(user_config.config_data)
+                    # Convert datetime strings back to datetime objects
+                    config_dict = self._convert_datetime_strings(config_dict)
+                    result[user_config.user_id] = MOSConfig(**config_dict)
+                except Exception as e:
+                    logger.error(f"Error parsing config for user {user_config.user_id}: {e}")
+                    continue
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Error listing user configs: {e}")
+            return {}
+        finally:
+            session.close()
+
+    def create_user_with_config(
+        self, user_name: str, config: MOSConfig, role=None, user_id: str | None = None
+    ) -> str:
+        """Create a new user with configuration.
+
+        Args:
+            user_name (str): Name of the user.
+            config (MOSConfig): The user's configuration.
+            role: User role (optional, uses default from UserManager).
+            user_id (str, optional): Custom user ID.
+
+        Returns:
+            str: The created user ID.
+
+        Raises:
+            ValueError: If user_name already exists.
+        """
+        # Create user using parent method
+        created_user_id = self.create_user(user_name, role, user_id)
+
+        # Save configuration
+        if not self.save_user_config(created_user_id, config):
+            logger.error(f"Failed to save configuration for user {created_user_id}")
+
+        return created_user_id
+
+    def delete_user(self, user_id: str) -> bool:
+        """Delete a user and their configuration.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        # Delete configuration first
+        self.delete_user_config(user_id)
+
+        # Delete user using parent method
+        return super().delete_user(user_id)
+
+    def get_user_cube_access(self, user_id: str) -> list[str]:
+        """Get list of cube IDs that a user has access to.
+
+        Args:
+            user_id (str): The user ID.
+
+        Returns:
+            list[str]: List of cube IDs the user can access.
+        """
+        cubes = self.get_user_cubes(user_id)
+        return [cube.cube_id for cube in cubes]

memos/mem_user/redis_persistent_user_manager.py

@@ -0,0 +1,225 @@
+"""Redis-based persistent user management system for MemOS with configuration storage.
+
+This module provides persistent storage for user configurations using Redis.
+"""
+
+import json
+
+from memos.configs.mem_os import MOSConfig
+from memos.dependency import require_python_package
+from memos.log import get_logger
+
+
+logger = get_logger(__name__)
+
+
+class RedisPersistentUserManager:
+    """Redis-based user configuration manager with persistence."""
+
+    @require_python_package(
+        import_name="redis",
+        install_command="pip install redis",
+        install_link="https://redis.readthedocs.io/en/stable/",
+    )
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        password: str = "",
+        db: int = 0,
+        decode_responses: bool = True,
+    ):
+        """Initialize the Redis persistent user manager.
+
+        Args:
+            user_id (str, optional): User ID. Defaults to "root".
+            host (str): Redis server host. Defaults to "localhost".
+            port (int): Redis server port. Defaults to 6379.
+            password (str): Redis password. Defaults to "".
+            db (int): Redis database number. Defaults to 0.
+            decode_responses (bool): Whether to decode responses to strings. Defaults to True.
+        """
+        import redis
+
+        self.host = host
+        self.port = port
+        self.db = db
+
+        try:
+            # Create Redis connection
+            self._redis_client = redis.Redis(
+                host=host,
+                port=port,
+                password=password if password else None,
+                db=db,
+                decode_responses=decode_responses,
+            )
+
+            # Test connection
+            if not self._redis_client.ping():
+                raise ConnectionError("Redis connection failed")
+
+            logger.info(
+                f"RedisPersistentUserManager initialized successfully, connected to {host}:{port}/{db}"
+            )
+
+        except Exception as e:
+            logger.error(f"Redis connection error: {e}")
+            raise
+
+    def _get_config_key(self, user_id: str) -> str:
+        """Generate Redis key for user configuration.
+
+        Args:
+            user_id (str): User ID.
+
+        Returns:
+            str: Redis key name.
+        """
+        return user_id
+
+    def save_user_config(self, user_id: str, config: MOSConfig) -> bool:
+        """Save user configuration to Redis.
+
+        Args:
+            user_id (str): User ID.
+            config (MOSConfig): User's MOS configuration.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        try:
+            # Convert config to JSON string
+            config_dict = config.model_dump(mode="json")
+            config_json = json.dumps(config_dict, ensure_ascii=False, indent=2)
+
+            # Save to Redis
+            key = self._get_config_key(user_id)
+            self._redis_client.set(key, config_json)
+
+            logger.info(f"Successfully saved configuration for user {user_id} to Redis")
+            return True
+
+        except Exception as e:
+            logger.error(f"Error saving configuration for user {user_id}: {e}")
+            return False
+
+    def get_user_config(self, user_id: str) -> dict | None:
+        """Get user configuration from Redis (search interface).
+
+        Args:
+            user_id (str): User ID.
+
+        Returns:
+            MOSConfig | None: User's configuration object, or None if not found.
+        """
+        try:
+            # Get configuration from Redis
+            key = self._get_config_key(user_id)
+            config_json = self._redis_client.get(key)
+
+            if config_json is None:
+                logger.info(f"Configuration for user {user_id} does not exist")
+                return None
+
+            # Parse JSON and create MOSConfig object
+            config_dict = json.loads(config_json)
+
+            logger.info(f"Successfully retrieved configuration for user {user_id}")
+            return config_dict
+
+        except json.JSONDecodeError as e:
+            logger.error(f"Error parsing JSON configuration for user {user_id}: {e}")
+            return None
+        except Exception as e:
+            logger.error(f"Error retrieving configuration for user {user_id}: {e}")
+            return None
+
+    def delete_user_config(self, user_id: str) -> bool:
+        """Delete user configuration from Redis.
+
+        Args:
+            user_id (str): User ID.
+
+        Returns:
+            bool: True if successful, False otherwise.
+        """
+        try:
+            key = self._get_config_key(user_id)
+            result = self._redis_client.delete(key)
+
+            if result > 0:
+                logger.info(f"Successfully deleted configuration for user {user_id}")
+                return True
+            else:
+                logger.warning(f"Configuration for user {user_id} does not exist, cannot delete")
+                return False
+
+        except Exception as e:
+            logger.error(f"Error deleting configuration for user {user_id}: {e}")
+            return False
+
+    def exists_user_config(self, user_id: str) -> bool:
+        """Check if user configuration exists.
+
+        Args:
+            user_id (str): User ID.
+
+        Returns:
+            bool: True if exists, False otherwise.
+        """
+        try:
+            key = self._get_config_key(user_id)
+            return self._redis_client.exists(key) > 0
+        except Exception as e:
+            logger.error(f"Error checking if configuration exists for user {user_id}: {e}")
+            return False
+
+    def list_user_configs(
+        self, pattern: str = "user_config:*", count: int = 100
+    ) -> dict[str, dict]:
+        """List all user configurations.
+
+        Args:
+            pattern (str): Redis key matching pattern. Defaults to "user_config:*".
+            count (int): Number of keys to return per scan. Defaults to 100.
+
+        Returns:
+            dict[str, dict]: Dictionary mapping user_id to dict objects.
+        """
+        result = {}
+        try:
+            # Use SCAN command to iterate through all matching keys
+            cursor = 0
+            while True:
+                cursor, keys = self._redis_client.scan(cursor, match=pattern, count=count)
+
+                for key in keys:
+                    # Extract user_id (remove "user_config:" prefix)
+                    user_id = key.replace("user_config:", "")
+                    config = self.get_user_config(user_id)
+                    if config:
+                        result[user_id] = config
+
+                if cursor == 0:
+                    break
+
+            logger.info(f"Successfully listed {len(result)} user configurations")
+            return result
+
+        except Exception as e:
+            logger.error(f"Error listing user configurations: {e}")
+            return {}
+
+    def close(self) -> None:
+        """Close Redis connection.
+
+        This method should be called when the RedisPersistentUserManager is no longer needed
+        to ensure proper cleanup of Redis connections.
+        """
+        try:
+            if hasattr(self, "_redis_client") and self._redis_client:
+                self._redis_client.close()
+                logger.info("Redis connection closed")
+        except Exception as e:
+            logger.error(f"Error closing Redis connection: {e}")