PraisonAI 3.0.0__py3-none-any.whl

praisonai/persistence/state/base.py

@@ -0,0 +1,107 @@
+"""
+Base interfaces for StateStore.
+
+StateStore handles fast key-value state and caching.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+import json
+
+
+class StateStore(ABC):
+    """
+    Abstract base class for key-value state persistence.
+    
+    Implementations handle fast state/cache storage for different backends:
+    - Redis (dedicated KV store)
+    - DynamoDB, Firestore (serverless document stores)
+    - MongoDB (document store)
+    - Upstash (serverless Redis)
+    - In-memory / JSON file (zero-dependency fallback)
+    """
+    
+    @abstractmethod
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value by key. Returns None if not found."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def set(
+        self, 
+        key: str, 
+        value: Any, 
+        ttl: Optional[int] = None
+    ) -> None:
+        """Set a value. TTL is in seconds."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def delete(self, key: str) -> bool:
+        """Delete a key. Returns True if key existed."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def keys(self, pattern: str = "*") -> List[str]:
+        """List keys matching pattern."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def ttl(self, key: str) -> Optional[int]:
+        """Get remaining TTL in seconds. Returns None if no TTL or key doesn't exist."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def expire(self, key: str, ttl: int) -> bool:
+        """Set TTL on existing key. Returns True if key exists."""
+        raise NotImplementedError
+    
+    def get_json(self, key: str) -> Optional[Any]:
+        """Get and deserialize JSON value."""
+        value = self.get(key)
+        if value is None:
+            return None
+        if isinstance(value, str):
+            return json.loads(value)
+        return value
+    
+    def set_json(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
+        """Serialize and set JSON value."""
+        self.set(key, json.dumps(value), ttl)
+    
+    @abstractmethod
+    def hget(self, key: str, field: str) -> Optional[Any]:
+        """Get a field from a hash."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def hset(self, key: str, field: str, value: Any) -> None:
+        """Set a field in a hash."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def hgetall(self, key: str) -> Dict[str, Any]:
+        """Get all fields from a hash."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def hdel(self, key: str, *fields: str) -> int:
+        """Delete fields from a hash. Returns count deleted."""
+        raise NotImplementedError
+    
+    @abstractmethod
+    def close(self) -> None:
+        """Close the store and release resources."""
+        raise NotImplementedError
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False

praisonai/persistence/state/dynamodb.py

@@ -0,0 +1,226 @@
+"""
+DynamoDB implementation of StateStore.
+
+Requires: boto3
+Install: pip install boto3
+"""
+
+import json
+import logging
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+class DynamoDBStateStore(StateStore):
+    """
+    DynamoDB-based state store.
+    
+    Example:
+        store = DynamoDBStateStore(
+            table_name="praisonai_state",
+            region="us-east-1"
+        )
+    """
+    
+    def __init__(
+        self,
+        table_name: str = "praisonai_state",
+        region: Optional[str] = None,
+        endpoint_url: Optional[str] = None,
+        auto_create_table: bool = True,
+    ):
+        try:
+            import boto3
+        except ImportError:
+            raise ImportError(
+                "boto3 is required for DynamoDB support. "
+                "Install with: pip install boto3"
+            )
+        
+        region = region or os.getenv("AWS_REGION", "us-east-1")
+        
+        self._dynamodb = boto3.resource(
+            "dynamodb",
+            region_name=region,
+            endpoint_url=endpoint_url,
+        )
+        self._client = boto3.client(
+            "dynamodb",
+            region_name=region,
+            endpoint_url=endpoint_url,
+        )
+        self.table_name = table_name
+        
+        if auto_create_table:
+            self._create_table()
+        
+        self._table = self._dynamodb.Table(table_name)
+        logger.info(f"Connected to DynamoDB table: {table_name}")
+    
+    def _create_table(self) -> None:
+        """Create table if not exists."""
+        try:
+            self._client.describe_table(TableName=self.table_name)
+        except self._client.exceptions.ResourceNotFoundException:
+            self._client.create_table(
+                TableName=self.table_name,
+                KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}],
+                AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}],
+                BillingMode="PAY_PER_REQUEST",
+            )
+            waiter = self._client.get_waiter("table_exists")
+            waiter.wait(TableName=self.table_name)
+            
+            # Enable TTL
+            self._client.update_time_to_live(
+                TableName=self.table_name,
+                TimeToLiveSpecification={"Enabled": True, "AttributeName": "ttl"}
+            )
+            logger.info(f"Created DynamoDB table: {self.table_name}")
+    
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value by key."""
+        response = self._table.get_item(Key={"pk": key})
+        item = response.get("Item")
+        
+        if not item:
+            return None
+        
+        # Check TTL
+        if item.get("ttl") and item["ttl"] <= int(time.time()):
+            return None
+        
+        value = item.get("value")
+        if isinstance(value, str):
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                return value
+        return value
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        ttl: Optional[int] = None
+    ) -> None:
+        """Set a value with optional TTL."""
+        item = {
+            "pk": key,
+            "value": json.dumps(value) if not isinstance(value, str) else value,
+            "updated_at": int(time.time()),
+        }
+        
+        if ttl:
+            item["ttl"] = int(time.time()) + ttl
+        
+        self._table.put_item(Item=item)
+    
+    def delete(self, key: str) -> bool:
+        """Delete a key."""
+        response = self._table.delete_item(
+            Key={"pk": key},
+            ReturnValues="ALL_OLD"
+        )
+        return "Attributes" in response
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        response = self._table.get_item(
+            Key={"pk": key},
+            ProjectionExpression="pk,#t",
+            ExpressionAttributeNames={"#t": "ttl"}
+        )
+        item = response.get("Item")
+        if not item:
+            return False
+        if item.get("ttl") and item["ttl"] <= int(time.time()):
+            return False
+        return True
+    
+    def keys(self, pattern: str = "*") -> List[str]:
+        """List keys matching pattern."""
+        # DynamoDB scan is expensive, use with caution
+        response = self._table.scan(ProjectionExpression="pk")
+        keys = [item["pk"] for item in response.get("Items", [])]
+        
+        if pattern != "*":
+            import fnmatch
+            keys = [k for k in keys if fnmatch.fnmatch(k, pattern)]
+        
+        return keys
+    
+    def ttl(self, key: str) -> Optional[int]:
+        """Get remaining TTL in seconds."""
+        response = self._table.get_item(
+            Key={"pk": key},
+            ProjectionExpression="#t",
+            ExpressionAttributeNames={"#t": "ttl"}
+        )
+        item = response.get("Item")
+        if not item or "ttl" not in item:
+            return None
+        
+        remaining = item["ttl"] - int(time.time())
+        if remaining <= 0:
+            return None
+        return remaining
+    
+    def expire(self, key: str, ttl: int) -> bool:
+        """Set TTL on existing key."""
+        try:
+            self._table.update_item(
+                Key={"pk": key},
+                UpdateExpression="SET #t = :ttl",
+                ExpressionAttributeNames={"#t": "ttl"},
+                ExpressionAttributeValues={":ttl": int(time.time()) + ttl},
+                ConditionExpression="attribute_exists(pk)"
+            )
+            return True
+        except self._client.exceptions.ConditionalCheckFailedException:
+            return False
+    
+    def hget(self, key: str, field: str) -> Optional[Any]:
+        """Get a field from a hash."""
+        value = self.get(key)
+        if not isinstance(value, dict):
+            return None
+        return value.get(field)
+    
+    def hset(self, key: str, field: str, value: Any) -> None:
+        """Set a field in a hash."""
+        current = self.get(key)
+        if not isinstance(current, dict):
+            current = {}
+        current[field] = value
+        self.set(key, current)
+    
+    def hgetall(self, key: str) -> Dict[str, Any]:
+        """Get all fields from a hash."""
+        value = self.get(key)
+        if not isinstance(value, dict):
+            return {}
+        return value
+    
+    def hdel(self, key: str, *fields: str) -> int:
+        """Delete fields from a hash."""
+        current = self.get(key)
+        if not isinstance(current, dict):
+            return 0
+        count = 0
+        for field in fields:
+            if field in current:
+                del current[field]
+                count += 1
+        if count > 0:
+            self.set(key, current)
+        return count
+    
+    def close(self) -> None:
+        """Close the store."""
+        pass  # boto3 handles cleanup

praisonai/persistence/state/firestore.py

@@ -0,0 +1,175 @@
+"""
+Firestore implementation of StateStore.
+
+Requires: google-cloud-firestore
+Install: pip install google-cloud-firestore
+"""
+
+import json
+import logging
+import os
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+class FirestoreStateStore(StateStore):
+    """
+    Firestore-based state store.
+    
+    Example:
+        store = FirestoreStateStore(
+            project="my-project",
+            collection="praisonai_state"
+        )
+    """
+    
+    def __init__(
+        self,
+        project: Optional[str] = None,
+        collection: str = "praisonai_state",
+        credentials_path: Optional[str] = None,
+    ):
+        try:
+            from google.cloud import firestore
+        except ImportError:
+            raise ImportError(
+                "google-cloud-firestore is required for Firestore support. "
+                "Install with: pip install google-cloud-firestore"
+            )
+        
+        if credentials_path:
+            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credentials_path
+        
+        project = project or os.getenv("FIRESTORE_PROJECT")
+        self._client = firestore.Client(project=project)
+        self._collection = self._client.collection(collection)
+        
+        logger.info(f"Connected to Firestore collection: {collection}")
+    
+    def get(self, key: str) -> Optional[Any]:
+        """Get a value by key."""
+        doc = self._collection.document(key).get()
+        if not doc.exists:
+            return None
+        
+        data = doc.to_dict()
+        
+        # Check TTL
+        if data.get("expires_at") and data["expires_at"] <= time.time():
+            self._collection.document(key).delete()
+            return None
+        
+        return data.get("value")
+    
+    def set(
+        self,
+        key: str,
+        value: Any,
+        ttl: Optional[int] = None
+    ) -> None:
+        """Set a value with optional TTL."""
+        data = {"value": value, "updated_at": time.time()}
+        
+        if ttl:
+            data["expires_at"] = time.time() + ttl
+        
+        self._collection.document(key).set(data)
+    
+    def delete(self, key: str) -> bool:
+        """Delete a key."""
+        doc = self._collection.document(key)
+        if doc.get().exists:
+            doc.delete()
+            return True
+        return False
+    
+    def exists(self, key: str) -> bool:
+        """Check if a key exists."""
+        doc = self._collection.document(key).get()
+        if not doc.exists:
+            return False
+        
+        data = doc.to_dict()
+        if data.get("expires_at") and data["expires_at"] <= time.time():
+            return False
+        
+        return True
+    
+    def keys(self, pattern: str = "*") -> List[str]:
+        """List keys matching pattern."""
+        docs = self._collection.stream()
+        keys = [doc.id for doc in docs]
+        
+        if pattern != "*":
+            import fnmatch
+            keys = [k for k in keys if fnmatch.fnmatch(k, pattern)]
+        
+        return keys
+    
+    def ttl(self, key: str) -> Optional[int]:
+        """Get remaining TTL in seconds."""
+        doc = self._collection.document(key).get()
+        if not doc.exists:
+            return None
+        
+        data = doc.to_dict()
+        if "expires_at" not in data:
+            return None
+        
+        remaining = data["expires_at"] - time.time()
+        if remaining <= 0:
+            return None
+        return int(remaining)
+    
+    def expire(self, key: str, ttl: int) -> bool:
+        """Set TTL on existing key."""
+        doc = self._collection.document(key)
+        if not doc.get().exists:
+            return False
+        
+        doc.update({"expires_at": time.time() + ttl})
+        return True
+    
+    def hget(self, key: str, field: str) -> Optional[Any]:
+        """Get a field from a hash."""
+        value = self.get(key)
+        if not isinstance(value, dict):
+            return None
+        return value.get(field)
+    
+    def hset(self, key: str, field: str, value: Any) -> None:
+        """Set a field in a hash."""
+        current = self.get(key)
+        if not isinstance(current, dict):
+            current = {}
+        current[field] = value
+        self.set(key, current)
+    
+    def hgetall(self, key: str) -> Dict[str, Any]:
+        """Get all fields from a hash."""
+        value = self.get(key)
+        if not isinstance(value, dict):
+            return {}
+        return value
+    
+    def hdel(self, key: str, *fields: str) -> int:
+        """Delete fields from a hash."""
+        current = self.get(key)
+        if not isinstance(current, dict):
+            return 0
+        count = 0
+        for field in fields:
+            if field in current:
+                del current[field]
+                count += 1
+        if count > 0:
+            self.set(key, current)
+        return count
+    
+    def close(self) -> None:
+        """Close the store."""
+        pass  # Firestore client handles cleanup

praisonai/persistence/state/gcs.py

@@ -0,0 +1,155 @@
+"""
+Google Cloud Storage implementation of StateStore.
+
+Requires: google-cloud-storage
+Install: pip install google-cloud-storage
+"""
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+from .base import StateStore
+
+logger = logging.getLogger(__name__)
+
+
+class GCSStateStore(StateStore):
+    """
+    Google Cloud Storage-based state store.
+    
+    Stores state as JSON objects in GCS buckets.
+    
+    Example:
+        store = GCSStateStore(
+            bucket_name="my-praisonai-state",
+            prefix="state/"
+        )
+    """
+    
+    def __init__(
+        self,
+        bucket_name: str,
+        prefix: str = "praisonai_state/",
+        project: Optional[str] = None,
+        credentials_path: Optional[str] = None,
+    ):
+        """
+        Initialize GCS state store.
+        
+        Args:
+            bucket_name: GCS bucket name
+            prefix: Object prefix for state keys
+            project: GCP project ID (optional)
+            credentials_path: Path to service account JSON (optional)
+        """
+        try:
+            from google.cloud import storage
+            if credentials_path:
+                self._client = storage.Client.from_service_account_json(
+                    credentials_path, project=project
+                )
+            else:
+                self._client = storage.Client(project=project)
+        except ImportError:
+            raise ImportError(
+                "google-cloud-storage is required for GCS support. "
+                "Install with: pip install google-cloud-storage"
+            )
+        
+        self.bucket_name = bucket_name
+        self.prefix = prefix
+        self._bucket = self._client.bucket(bucket_name)
+    
+    def _key_to_path(self, key: str) -> str:
+        """Convert key to GCS object path."""
+        return f"{self.prefix}{key}.json"
+    
+    def get(self, key: str) -> Optional[Dict[str, Any]]:
+        """Get state by key."""
+        try:
+            blob = self._bucket.blob(self._key_to_path(key))
+            if blob.exists():
+                data = json.loads(blob.download_as_text())
+                # Check TTL
+                if "_ttl_expires" in data:
+                    if time.time() > data["_ttl_expires"]:
+                        self.delete(key)
+                        return None
+                    del data["_ttl_expires"]
+                return data
+            return None
+        except Exception as e:
+            logger.error(f"Error getting state {key}: {e}")
+            return None
+    
+    def set(self, key: str, value: Dict[str, Any], ttl: Optional[int] = None) -> bool:
+        """Set state by key with optional TTL in seconds."""
+        try:
+            data = value.copy()
+            if ttl:
+                data["_ttl_expires"] = time.time() + ttl
+            
+            blob = self._bucket.blob(self._key_to_path(key))
+            blob.upload_from_string(
+                json.dumps(data, default=str),
+                content_type="application/json"
+            )
+            return True
+        except Exception as e:
+            logger.error(f"Error setting state {key}: {e}")
+            return False
+    
+    def delete(self, key: str) -> bool:
+        """Delete state by key."""
+        try:
+            blob = self._bucket.blob(self._key_to_path(key))
+            if blob.exists():
+                blob.delete()
+                return True
+            return False
+        except Exception as e:
+            logger.error(f"Error deleting state {key}: {e}")
+            return False
+    
+    def exists(self, key: str) -> bool:
+        """Check if key exists."""
+        try:
+            blob = self._bucket.blob(self._key_to_path(key))
+            return blob.exists()
+        except Exception:
+            return False
+    
+    def list_keys(self, prefix: Optional[str] = None) -> List[str]:
+        """List all keys with optional prefix filter."""
+        try:
+            search_prefix = self.prefix
+            if prefix:
+                search_prefix = f"{self.prefix}{prefix}"
+            
+            blobs = self._client.list_blobs(self.bucket_name, prefix=search_prefix)
+            keys = []
+            for blob in blobs:
+                # Remove prefix and .json suffix
+                key = blob.name[len(self.prefix):]
+                if key.endswith(".json"):
+                    key = key[:-5]
+                keys.append(key)
+            return keys
+        except Exception as e:
+            logger.error(f"Error listing keys: {e}")
+            return []
+    
+    def clear(self, prefix: Optional[str] = None) -> int:
+        """Clear all keys with optional prefix filter."""
+        keys = self.list_keys(prefix)
+        count = 0
+        for key in keys:
+            if self.delete(key):
+                count += 1
+        return count
+    
+    def close(self) -> None:
+        """Close the store."""
+        self._client.close()