PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/state/memory.py

@@ -0,0 +1,245 @@
+"""
+In-memory / JSON file implementation of StateStore.
+
+Zero external dependencies - uses Python built-ins.
+"""
+
+import json
+import logging
+import os
+import threading
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryStateStore(StateStore):
+    """
+    In-memory state store with optional JSON file persistence.
+    
+    Zero external dependencies - ideal for development and testing.
+    
+    Example:
+        # In-memory only
+        store = MemoryStateStore()
+        
+        # With file persistence
+        store = MemoryStateStore(path="./state.json")
+    """
+    
+    def __init__(
+        self,
+        path: Optional[str] = None,
+        auto_save: bool = True,
+        save_interval: int = 60,
+    ):
+        """
+        Initialize memory state store.
+        
+        Args:
+            path: Path to JSON file for persistence (None = in-memory only)
+            auto_save: Auto-save to file periodically
+            save_interval: Save interval in seconds
+        """
+        self.path = path
+        self.auto_save = auto_save
+        self.save_interval = save_interval
+        
+        self._data: Dict[str, Any] = {}
+        self._ttls: Dict[str, float] = {}  # key -> expiry timestamp
+        self._lock = threading.RLock()
+        self._last_save = time.time()
+        
+        # Load from file if exists
+        if path and os.path.exists(path):
+            self._load()
+        
+        logger.info(f"MemoryStateStore initialized (path={path})")
+    
+    def _load(self) -> None:
+        """Load state from JSON file."""
+        if not self.path:
+            return
+        
+        try:
+            with open(self.path, "r") as f:
+                saved = json.load(f)
+                self._data = saved.get("data", {})
+                self._ttls = saved.get("ttls", {})
+                
+                # Clean expired keys
+                now = time.time()
+                expired = [k for k, exp in self._ttls.items() if exp <= now]
+                for k in expired:
+                    self._data.pop(k, None)
+                    self._ttls.pop(k, None)
+                
+            logger.debug(f"Loaded {len(self._data)} keys from {self.path}")
+        except Exception as e:
+            logger.warning(f"Failed to load state from {self.path}: {e}")
+    
+    def _save(self) -> None:
+        """Save state to JSON file."""
+        if not self.path:
+            return
+        
+        try:
+            os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True)
+            with open(self.path, "w") as f:
+                json.dump({"data": self._data, "ttls": self._ttls}, f)
+            self._last_save = time.time()
+        except Exception as e:
+            logger.warning(f"Failed to save state to {self.path}: {e}")
+    
+    def _maybe_save(self) -> None:
+        """Save if auto_save is enabled and interval has passed."""
+        if self.auto_save and self.path:
+            if time.time() - self._last_save >= self.save_interval:
+                self._save()
+    
+    def _check_ttl(self, key: str) -> bool:
+        """Check if key is expired. Returns True if valid, False if expired."""
+        if key not in self._ttls:
+            return True
+        if self._ttls[key] <= time.time():
+            self._data.pop(key, None)
+            self._ttls.pop(key, None)
+            return False
+        return True
+    
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value by key."""
+        with self._lock:
+            if not self._check_ttl(key):
+                return None
+            return self._data.get(key)
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        ttl: Optional[int] = None
+    ) -> None:
+        """Set a value with optional TTL."""
+        with self._lock:
+            self._data[key] = value
+            if ttl:
+                self._ttls[key] = time.time() + ttl
+            elif key in self._ttls:
+                del self._ttls[key]
+            self._maybe_save()
+    
+    def delete(self, key: str) -> bool:
+        """Delete a key."""
+        with self._lock:
+            existed = key in self._data
+            self._data.pop(key, None)
+            self._ttls.pop(key, None)
+            self._maybe_save()
+            return existed
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        with self._lock:
+            if not self._check_ttl(key):
+                return False
+            return key in self._data
+    
+    def keys(self, pattern: str = "*") -> List[str]:
+        """List keys matching pattern."""
+        import fnmatch
+        
+        with self._lock:
+            # Clean expired keys first
+            now = time.time()
+            expired = [k for k, exp in self._ttls.items() if exp <= now]
+            for k in expired:
+                self._data.pop(k, None)
+                self._ttls.pop(k, None)
+            
+            if pattern == "*":
+                return list(self._data.keys())
+            return [k for k in self._data.keys() if fnmatch.fnmatch(k, pattern)]
+    
+    def ttl(self, key: str) -> Optional[int]:
+        """Get remaining TTL in seconds."""
+        with self._lock:
+            if key not in self._ttls:
+                return None
+            remaining = self._ttls[key] - time.time()
+            if remaining <= 0:
+                self._data.pop(key, None)
+                self._ttls.pop(key, None)
+                return None
+            return int(remaining)
+    
+    def expire(self, key: str, ttl: int) -> bool:
+        """Set TTL on existing key."""
+        with self._lock:
+            if key not in self._data:
+                return False
+            self._ttls[key] = time.time() + ttl
+            self._maybe_save()
+            return True
+    
+    def hget(self, key: str, field: str) -> Optional[Any]:
+        """Get a field from a hash."""
+        with self._lock:
+            if not self._check_ttl(key):
+                return None
+            data = self._data.get(key)
+            if not isinstance(data, dict):
+                return None
+            return data.get(field)
+    
+    def hset(self, key: str, field: str, value: Any) -> None:
+        """Set a field in a hash."""
+        with self._lock:
+            if key not in self._data or not isinstance(self._data[key], dict):
+                self._data[key] = {}
+            self._data[key][field] = value
+            self._maybe_save()
+    
+    def hgetall(self, key: str) -> Dict[str, Any]:
+        """Get all fields from a hash."""
+        with self._lock:
+            if not self._check_ttl(key):
+                return {}
+            data = self._data.get(key)
+            if not isinstance(data, dict):
+                return {}
+            return dict(data)
+    
+    def hdel(self, key: str, *fields: str) -> int:
+        """Delete fields from a hash."""
+        with self._lock:
+            if key not in self._data or not isinstance(self._data[key], dict):
+                return 0
+            count = 0
+            for field in fields:
+                if field in self._data[key]:
+                    del self._data[key][field]
+                    count += 1
+            self._maybe_save()
+            return count
+    
+    def flush(self) -> None:
+        """Force save to file."""
+        with self._lock:
+            self._save()
+    
+    def clear(self) -> None:
+        """Clear all data."""
+        with self._lock:
+            self._data.clear()
+            self._ttls.clear()
+            self._maybe_save()
+    
+    def close(self) -> None:
+        """Close the store and save."""
+        with self._lock:
+            if self.path:
+                self._save()

praisonai/persistence/state/mongodb.py

@@ -0,0 +1,158 @@
+"""
+MongoDB implementation of StateStore.
+
+Requires: pymongo
+Install: pip install pymongo
+"""
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDBStateStore(StateStore):
+    """
+    MongoDB-based state store.
+    
+    Example:
+        store = MongoDBStateStore(
+            url="mongodb://localhost:27017",
+            database="praisonai"
+        )
+    """
+    
+    def __init__(
+        self,
+        url: str = "mongodb://localhost:27017",
+        database: str = "praisonai",
+        collection: str = "state",
+    ):
+        try:
+            from pymongo import MongoClient
+        except ImportError:
+            raise ImportError(
+                "pymongo is required for MongoDB support. "
+                "Install with: pip install pymongo"
+            )
+        
+        self._client = MongoClient(url)
+        self._db = self._client[database]
+        self._collection = self._db[collection]
+        
+        # Create TTL index for automatic expiration
+        self._collection.create_index("expires_at", expireAfterSeconds=0)
+        
+        logger.info(f"Connected to MongoDB: {database}.{collection}")
+    
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value by key."""
+        doc = self._collection.find_one({"_id": key})
+        if not doc:
+            return None
+        
+        # Check TTL
+        if doc.get("expires_at") and doc["expires_at"] <= time.time():
+            self._collection.delete_one({"_id": key})
+            return None
+        
+        return doc.get("value")
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        ttl: Optional[int] = None
+    ) -> None:
+        """Set a value with optional TTL."""
+        doc = {"_id": key, "value": value, "updated_at": time.time()}
+        
+        if ttl:
+            from datetime import datetime, timedelta
+            doc["expires_at"] = datetime.utcnow() + timedelta(seconds=ttl)
+        
+        self._collection.replace_one({"_id": key}, doc, upsert=True)
+    
+    def delete(self, key: str) -> bool:
+        """Delete a key."""
+        result = self._collection.delete_one({"_id": key})
+        return result.deleted_count > 0
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        doc = self._collection.find_one({"_id": key}, {"_id": 1, "expires_at": 1})
+        if not doc:
+            return False
+        if doc.get("expires_at") and doc["expires_at"] <= time.time():
+            return False
+        return True
+    
+    def keys(self, pattern: str = "*") -> List[str]:
+        """List keys matching pattern."""
+        if pattern == "*":
+            cursor = self._collection.find({}, {"_id": 1})
+        else:
+            # Convert glob pattern to regex
+            import re
+            regex = pattern.replace("*", ".*").replace("?", ".")
+            cursor = self._collection.find({"_id": {"$regex": f"^{regex}$"}}, {"_id": 1})
+        
+        return [doc["_id"] for doc in cursor]
+    
+    def ttl(self, key: str) -> Optional[int]:
+        """Get remaining TTL in seconds."""
+        doc = self._collection.find_one({"_id": key}, {"expires_at": 1})
+        if not doc or "expires_at" not in doc:
+            return None
+        
+        remaining = doc["expires_at"].timestamp() - time.time()
+        if remaining <= 0:
+            return None
+        return int(remaining)
+    
+    def expire(self, key: str, ttl: int) -> bool:
+        """Set TTL on existing key."""
+        from datetime import datetime, timedelta
+        
+        result = self._collection.update_one(
+            {"_id": key},
+            {"$set": {"expires_at": datetime.utcnow() + timedelta(seconds=ttl)}}
+        )
+        return result.modified_count > 0
+    
+    def hget(self, key: str, field: str) -> Optional[Any]:
+        """Get a field from a hash."""
+        doc = self._collection.find_one({"_id": key})
+        if not doc or not isinstance(doc.get("value"), dict):
+            return None
+        return doc["value"].get(field)
+    
+    def hset(self, key: str, field: str, value: Any) -> None:
+        """Set a field in a hash."""
+        self._collection.update_one(
+            {"_id": key},
+            {"$set": {f"value.{field}": value, "updated_at": time.time()}},
+            upsert=True
+        )
+    
+    def hgetall(self, key: str) -> Dict[str, Any]:
+        """Get all fields from a hash."""
+        doc = self._collection.find_one({"_id": key})
+        if not doc or not isinstance(doc.get("value"), dict):
+            return {}
+        return doc["value"]
+    
+    def hdel(self, key: str, *fields: str) -> int:
+        """Delete fields from a hash."""
+        unset = {f"value.{field}": "" for field in fields}
+        result = self._collection.update_one({"_id": key}, {"$unset": unset})
+        return len(fields) if result.modified_count > 0 else 0
+    
+    def close(self) -> None:
+        """Close the store."""
+        if self._client:
+            self._client.close()

praisonai/persistence/state/redis.py

@@ -0,0 +1,190 @@
+"""
+Redis implementation of StateStore.
+
+Requires: redis
+Install: pip install redis
+"""
+
+import json
+import logging
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+class RedisStateStore(StateStore):
+    """
+    Redis-based state store for fast key-value operations.
+    
+    Example:
+        store = RedisStateStore(
+            url="redis://localhost:6379"
+        )
+    """
+    
+    def __init__(
+        self,
+        url: Optional[str] = None,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password: Optional[str] = None,
+        prefix: str = "praison:",
+        decode_responses: bool = True,
+        socket_timeout: int = 5,
+        max_connections: int = 10,
+    ):
+        """
+        Initialize Redis state store.
+        
+        Args:
+            url: Full Redis URL (overrides host/port/db/password)
+            host: Redis host
+            port: Redis port
+            db: Redis database number
+            password: Redis password
+            prefix: Key prefix for namespacing
+            decode_responses: Decode bytes to strings
+            socket_timeout: Socket timeout in seconds
+            max_connections: Max connections in pool
+        """
+        try:
+            import redis as redis_lib
+        except ImportError:
+            raise ImportError(
+                "redis is required for Redis support. "
+                "Install with: pip install redis"
+            )
+        
+        self._redis_lib = redis_lib
+        self.prefix = prefix
+        
+        if url:
+            self._client = redis_lib.from_url(
+                url,
+                decode_responses=decode_responses,
+                socket_timeout=socket_timeout,
+                max_connections=max_connections,
+            )
+        else:
+            pool = redis_lib.ConnectionPool(
+                host=host,
+                port=port,
+                db=db,
+                password=password,
+                decode_responses=decode_responses,
+                socket_timeout=socket_timeout,
+                max_connections=max_connections,
+            )
+            self._client = redis_lib.Redis(connection_pool=pool)
+        
+        # Test connection
+        self._client.ping()
+        logger.info(f"Connected to Redis at {url or f'{host}:{port}'}")
+    
+    def _key(self, key: str) -> str:
+        """Add prefix to key."""
+        return f"{self.prefix}{key}"
+    
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value by key."""
+        value = self._client.get(self._key(key))
+        if value is None:
+            return None
+        # Try to deserialize JSON
+        try:
+            return json.loads(value)
+        except (json.JSONDecodeError, TypeError):
+            return value
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        ttl: Optional[int] = None
+    ) -> None:
+        """Set a value with optional TTL."""
+        # Serialize non-string values
+        if not isinstance(value, str):
+            value = json.dumps(value)
+        
+        if ttl:
+            self._client.setex(self._key(key), ttl, value)
+        else:
+            self._client.set(self._key(key), value)
+    
+    def delete(self, key: str) -> bool:
+        """Delete a key."""
+        return self._client.delete(self._key(key)) > 0
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        return self._client.exists(self._key(key)) > 0
+    
+    def keys(self, pattern: str = "*") -> List[str]:
+        """List keys matching pattern."""
+        full_pattern = self._key(pattern)
+        keys = self._client.keys(full_pattern)
+        # Remove prefix from returned keys
+        prefix_len = len(self.prefix)
+        return [k[prefix_len:] if k.startswith(self.prefix) else k for k in keys]
+    
+    def ttl(self, key: str) -> Optional[int]:
+        """Get remaining TTL in seconds."""
+        result = self._client.ttl(self._key(key))
+        if result < 0:  # -1 = no TTL, -2 = key doesn't exist
+            return None
+        return result
+    
+    def expire(self, key: str, ttl: int) -> bool:
+        """Set TTL on existing key."""
+        return self._client.expire(self._key(key), ttl)
+    
+    def hget(self, key: str, field: str) -> Optional[Any]:
+        """Get a field from a hash."""
+        value = self._client.hget(self._key(key), field)
+        if value is None:
+            return None
+        try:
+            return json.loads(value)
+        except (json.JSONDecodeError, TypeError):
+            return value
+    
+    def hset(self, key: str, field: str, value: Any) -> None:
+        """Set a field in a hash."""
+        if not isinstance(value, str):
+            value = json.dumps(value)
+        self._client.hset(self._key(key), field, value)
+    
+    def hgetall(self, key: str) -> Dict[str, Any]:
+        """Get all fields from a hash."""
+        data = self._client.hgetall(self._key(key))
+        result = {}
+        for k, v in data.items():
+            try:
+                result[k] = json.loads(v)
+            except (json.JSONDecodeError, TypeError):
+                result[k] = v
+        return result
+    
+    def hdel(self, key: str, *fields: str) -> int:
+        """Delete fields from a hash."""
+        if not fields:
+            return 0
+        return self._client.hdel(self._key(key), *fields)
+    
+    def incr(self, key: str, amount: int = 1) -> int:
+        """Increment a counter."""
+        return self._client.incrby(self._key(key), amount)
+    
+    def decr(self, key: str, amount: int = 1) -> int:
+        """Decrement a counter."""
+        return self._client.decrby(self._key(key), amount)
+    
+    def close(self) -> None:
+        """Close the store."""
+        if self._client:
+            self._client.close()
+            self._client = None