flatagents 0.4.1__py3-none-any.whl

flatagents/gcp/__init__.py

@@ -0,0 +1,25 @@
+"""
+Google Cloud Platform backends for FlatAgents.
+
+Provides Firestore-based persistence and result storage for
+Cloud Functions and Firebase deployments.
+
+Usage:
+    from flatagents.gcp import FirestoreBackend
+    
+    backend = FirestoreBackend(collection="flatagents")
+    machine = FlatMachine(
+        config_file="machine.yml",
+        persistence=backend,
+        result_backend=backend
+    )
+
+Requirements:
+    pip install google-cloud-firestore
+"""
+
+from .firestore import FirestoreBackend
+
+__all__ = [
+    "FirestoreBackend",
+]

flatagents/gcp/firestore.py

@@ -0,0 +1,227 @@
+"""
+Firestore backend for FlatAgents persistence and results.
+
+Implements both PersistenceBackend and ResultBackend using Firestore.
+Compatible with Cloud Functions, Firebase, and local emulator.
+
+Document structure:
+    Collection: flatagents (configurable)
+    Document ID: {execution_id}
+    Subcollections:
+        - checkpoints/{step}_{event}
+        - results/{path}
+"""
+
+import asyncio
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+# Lazy import to avoid hard dependency
+_firestore = None
+_async_client = None
+
+
+def _get_firestore():
+    global _firestore
+    if _firestore is None:
+        try:
+            from google.cloud import firestore
+            _firestore = firestore
+        except ImportError:
+            raise ImportError(
+                "google-cloud-firestore is required for GCP backends. "
+                "Install with: pip install google-cloud-firestore"
+            )
+    return _firestore
+
+
+def _get_async_client():
+    """Get or create an async Firestore client."""
+    global _async_client
+    if _async_client is None:
+        firestore = _get_firestore()
+        _async_client = firestore.AsyncClient()
+    return _async_client
+
+
+class FirestoreBackend:
+    """
+    Combined Persistence and Result backend using Firestore.
+    
+    Implements both PersistenceBackend and ResultBackend interfaces.
+    Uses a single collection with subcollections for organization.
+    
+    Args:
+        collection: Root collection name (default: "flatagents")
+        project: GCP project ID (optional, uses default)
+    
+    Document Layout:
+        flatagents/{execution_id}/checkpoints/{step_event} = checkpoint data
+        flatagents/{execution_id}/results/{path} = result data
+        flatagents/{execution_id}/_meta = metadata (created_at, etc.)
+    """
+    
+    def __init__(
+        self,
+        collection: str = "flatagents",
+        project: Optional[str] = None
+    ):
+        self.collection = collection
+        self.project = project
+        self._db = None
+    
+    @property
+    def db(self):
+        """Lazy-load Firestore client."""
+        if self._db is None:
+            firestore = _get_firestore()
+            if self.project:
+                self._db = firestore.AsyncClient(project=self.project)
+            else:
+                self._db = firestore.AsyncClient()
+        return self._db
+    
+    def _doc_ref(self, execution_id: str, subcollection: str, doc_id: str):
+        """Get document reference."""
+        return (
+            self.db.collection(self.collection)
+            .document(execution_id)
+            .collection(subcollection)
+            .document(doc_id)
+        )
+    
+    # =========================================================================
+    # PersistenceBackend Interface
+    # =========================================================================
+    
+    async def save(self, key: str, value: bytes) -> None:
+        """Save checkpoint data.
+        
+        Args:
+            key: Format "{execution_id}/step_{step}_{event}"
+            value: JSON-encoded checkpoint bytes
+        """
+        parts = key.split("/", 1)
+        execution_id = parts[0]
+        doc_id = parts[1] if len(parts) > 1 else "latest"
+        
+        doc_ref = self._doc_ref(execution_id, "checkpoints", doc_id)
+        
+        await doc_ref.set({
+            "data": value.decode("utf-8"),
+            "created_at": datetime.now(timezone.utc).isoformat(),
+        })
+        
+        logger.debug(f"Firestore: saved checkpoint {execution_id}/{doc_id}")
+    
+    async def load(self, key: str) -> Optional[bytes]:
+        """Load checkpoint data."""
+        parts = key.split("/", 1)
+        execution_id = parts[0]
+        doc_id = parts[1] if len(parts) > 1 else "latest"
+        
+        doc_ref = self._doc_ref(execution_id, "checkpoints", doc_id)
+        doc = await doc_ref.get()
+        
+        if not doc.exists:
+            return None
+        
+        return doc.to_dict()["data"].encode("utf-8")
+    
+    async def delete(self, key: str) -> None:
+        """Delete checkpoint data."""
+        parts = key.split("/", 1)
+        execution_id = parts[0]
+        doc_id = parts[1] if len(parts) > 1 else "latest"
+        
+        doc_ref = self._doc_ref(execution_id, "checkpoints", doc_id)
+        await doc_ref.delete()
+    
+    async def list(self, prefix: str) -> List[str]:
+        """List all keys matching prefix."""
+        execution_id = prefix.rstrip("/")
+        
+        collection_ref = (
+            self.db.collection(self.collection)
+            .document(execution_id)
+            .collection("checkpoints")
+        )
+        
+        docs = collection_ref.stream()
+        keys = []
+        async for doc in docs:
+            keys.append(f"{execution_id}/{doc.id}")
+        
+        return sorted(keys)
+    
+    # =========================================================================
+    # ResultBackend Interface
+    # =========================================================================
+    
+    async def write(self, uri: str, data: Any) -> None:
+        """Write result to a URI."""
+        from ..backends import parse_uri
+        
+        execution_id, path = parse_uri(uri)
+        doc_ref = self._doc_ref(execution_id, "results", path)
+        
+        await doc_ref.set({
+            "data": json.dumps(data),
+            "created_at": datetime.now(timezone.utc).isoformat(),
+        })
+        
+        logger.debug(f"Firestore: wrote result {execution_id}/{path}")
+    
+    async def read(
+        self,
+        uri: str,
+        block: bool = True,
+        timeout: Optional[float] = None
+    ) -> Any:
+        """Read result from a URI, optionally blocking until available."""
+        from ..backends import parse_uri
+        
+        execution_id, path = parse_uri(uri)
+        doc_ref = self._doc_ref(execution_id, "results", path)
+        
+        start_time = datetime.now(timezone.utc).timestamp()
+        poll_interval = 0.5
+        
+        while True:
+            doc = await doc_ref.get()
+            
+            if doc.exists:
+                return json.loads(doc.to_dict()["data"])
+            
+            if not block:
+                return None
+            
+            if timeout:
+                elapsed = datetime.now(timezone.utc).timestamp() - start_time
+                if elapsed >= timeout:
+                    raise TimeoutError(f"Timeout waiting for result at {uri}")
+            
+            await asyncio.sleep(poll_interval)
+            poll_interval = min(poll_interval * 1.5, 5.0)
+    
+    async def exists(self, uri: str) -> bool:
+        """Check if result exists at URI."""
+        from ..backends import parse_uri
+        
+        execution_id, path = parse_uri(uri)
+        doc_ref = self._doc_ref(execution_id, "results", path)
+        doc = await doc_ref.get()
+        
+        return doc.exists
+    
+    async def delete(self, uri: str) -> None:
+        """Delete result at URI."""
+        from ..backends import parse_uri
+        
+        execution_id, path = parse_uri(uri)
+        doc_ref = self._doc_ref(execution_id, "results", path)
+        await doc_ref.delete()

flatagents/hooks.py

@@ -0,0 +1,380 @@
+"""
+MachineHooks - Extensibility points for FlatMachine.
+
+Hooks allow custom logic at key points in machine execution:
+- Before/after state entry/exit
+- Before/after agent calls
+- On transitions
+- On errors
+
+Includes built-in LoggingHooks and MetricsHooks implementations.
+"""
+
+import logging
+import time
+from abc import ABC
+from typing import Any, Dict, Optional
+
+from .monitoring import get_logger
+
+logger = get_logger(__name__)
+
+try:
+    import httpx
+except ImportError:
+    httpx = None
+
+
+class MachineHooks(ABC):
+    """
+    Base class for machine hooks.
+    
+    Override methods to customize machine behavior.
+    All methods have default implementations that pass through unchanged.
+    
+    Example:
+        from flatagents import get_logger
+        logger = get_logger(__name__)
+
+        class MyHooks(MachineHooks):
+            def on_state_enter(self, state_name, context):
+                logger.info(f"Entering state: {state_name}")
+                return context
+                
+        machine = FlatMachine(config_file="...", hooks=MyHooks())
+    """
+
+    def on_machine_start(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Called when machine execution starts.
+        
+        Args:
+            context: Initial context
+            
+        Returns:
+            Modified context
+        """
+        return context
+
+    def on_machine_end(self, context: Dict[str, Any], final_output: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Called when machine execution ends.
+        
+        Args:
+            context: Final context
+            final_output: Output from final state
+            
+        Returns:
+            Modified final output
+        """
+        return final_output
+
+    def on_state_enter(self, state_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Called before executing a state.
+        
+        Args:
+            state_name: Name of the state being entered
+            context: Current context
+            
+        Returns:
+            Modified context
+        """
+        return context
+
+    def on_state_exit(
+        self,
+        state_name: str,
+        context: Dict[str, Any],
+        output: Optional[Dict[str, Any]]
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Called after executing a state.
+        
+        Args:
+            state_name: Name of the state that was executed
+            context: Current context
+            output: Output from the state (agent output or None)
+            
+        Returns:
+            Modified output
+        """
+        return output
+
+    def on_transition(
+        self,
+        from_state: str,
+        to_state: str,
+        context: Dict[str, Any]
+    ) -> str:
+        """
+        Called when transitioning between states.
+        
+        Can override the target state.
+        
+        Args:
+            from_state: Source state name
+            to_state: Target state name (from transition evaluation)
+            context: Current context
+            
+        Returns:
+            Actual target state name (can override)
+        """
+        return to_state
+
+    def on_error(
+        self,
+        state_name: str,
+        error: Exception,
+        context: Dict[str, Any]
+    ) -> Optional[str]:
+        """
+        Called when an error occurs during state execution.
+        
+        Args:
+            state_name: Name of the state where error occurred
+            error: The exception that was raised
+            context: Current context
+            
+        Returns:
+            State to transition to, or None to re-raise the error
+        """
+        return None  # Re-raise by default
+
+    def on_action(
+        self,
+        action_name: str,
+        context: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Called for custom hook actions defined in states.
+        
+        Args:
+            action_name: Name of the action to execute
+            context: Current context
+            
+        Returns:
+            Modified context
+        """
+        logger.warning(f"Unhandled action: {action_name}")
+        return context
+
+
+class LoggingHooks(MachineHooks):
+    """Hooks that log all state transitions."""
+
+    def __init__(self, log_level: int = logging.INFO):
+        self.log_level = log_level
+
+    def on_machine_start(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        logger.log(self.log_level, "Machine execution started")
+        return context
+
+    def on_machine_end(self, context: Dict[str, Any], final_output: Dict[str, Any]) -> Dict[str, Any]:
+        logger.log(self.log_level, f"Machine execution ended with output: {final_output}")
+        return final_output
+
+    def on_state_enter(self, state_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        logger.log(self.log_level, f"Entering state: {state_name}")
+        return context
+
+    def on_state_exit(
+        self,
+        state_name: str,
+        context: Dict[str, Any],
+        output: Optional[Dict[str, Any]]
+    ) -> Optional[Dict[str, Any]]:
+        logger.log(self.log_level, f"Exiting state: {state_name}")
+        return output
+
+    def on_transition(self, from_state: str, to_state: str, context: Dict[str, Any]) -> str:
+        logger.log(self.log_level, f"Transition: {from_state} -> {to_state}")
+        return to_state
+
+
+class MetricsHooks(MachineHooks):
+    """Hooks that track execution metrics."""
+
+    def __init__(self):
+        self.state_counts: Dict[str, int] = {}
+        self.transition_counts: Dict[str, int] = {}
+        self.total_states_executed = 0
+        self.error_count = 0
+
+    def on_state_enter(self, state_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        self.state_counts[state_name] = self.state_counts.get(state_name, 0) + 1
+        self.total_states_executed += 1
+        return context
+
+    def on_transition(self, from_state: str, to_state: str, context: Dict[str, Any]) -> str:
+        key = f"{from_state}->{to_state}"
+        self.transition_counts[key] = self.transition_counts.get(key, 0) + 1
+        return to_state
+
+    def on_error(self, state_name: str, error: Exception, context: Dict[str, Any]) -> Optional[str]:
+        self.error_count += 1
+        return None
+
+    def get_metrics(self) -> Dict[str, Any]:
+        """Get collected metrics."""
+        return {
+            "state_counts": self.state_counts,
+            "transition_counts": self.transition_counts,
+            "total_states_executed": self.total_states_executed,
+            "error_count": self.error_count,
+        }
+
+
+class CompositeHooks(MachineHooks):
+    """Compose multiple hooks together."""
+
+    def __init__(self, *hooks: MachineHooks):
+        self.hooks = list(hooks)
+
+    def on_machine_start(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        for hook in self.hooks:
+            context = hook.on_machine_start(context)
+        return context
+
+    def on_machine_end(self, context: Dict[str, Any], final_output: Dict[str, Any]) -> Dict[str, Any]:
+        for hook in self.hooks:
+            final_output = hook.on_machine_end(context, final_output)
+        return final_output
+
+    def on_state_enter(self, state_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        for hook in self.hooks:
+            context = hook.on_state_enter(state_name, context)
+        return context
+
+    def on_state_exit(
+        self,
+        state_name: str,
+        context: Dict[str, Any],
+        output: Optional[Dict[str, Any]]
+    ) -> Optional[Dict[str, Any]]:
+        for hook in self.hooks:
+            output = hook.on_state_exit(state_name, context, output)
+        return output
+
+    def on_transition(self, from_state: str, to_state: str, context: Dict[str, Any]) -> str:
+        for hook in self.hooks:
+            to_state = hook.on_transition(from_state, to_state, context)
+        return to_state
+
+    def on_error(self, state_name: str, error: Exception, context: Dict[str, Any]) -> Optional[str]:
+        for hook in self.hooks:
+            result = hook.on_error(state_name, error, context)
+            if result is not None:
+                return result
+        return None
+
+    def on_action(self, action_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        for hook in self.hooks:
+            context = hook.on_action(action_name, context)
+        return context
+
+
+class WebhookHooks(MachineHooks):
+    """
+    Hooks that dispatch events to an HTTP endpoint.
+    
+    Requires 'httpx' installed.
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        timeout: float = 5.0,
+        api_key: Optional[str] = None
+    ):
+        if httpx is None:
+            raise ImportError("httpx is required for WebhookHooks")
+            
+        self.endpoint = endpoint
+        self.timeout = timeout
+        self.headers = {
+            "Content-Type": "application/json",
+            "User-Agent": "FlatAgents/0.1.0"
+        }
+        if api_key:
+            self.headers["Authorization"] = f"Bearer {api_key}"
+
+    async def _send(self, event: str, payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Send event to webhook."""
+        data = {"event": event, **payload}
+        try:
+            async with httpx.AsyncClient() as client:
+                response = await client.post(
+                    self.endpoint,
+                    json=data,
+                    headers=self.headers,
+                    timeout=self.timeout
+                )
+                response.raise_for_status()
+                if response.status_code == 204:
+                    return None
+                return response.json()
+        except Exception as e:
+            logger.error(f"Webhook error ({event}): {e}")
+            return None
+
+    async def on_machine_start(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        resp = await self._send("machine_start", {"context": context})
+        if resp and "context" in resp:
+            return resp["context"]
+        return context
+
+    async def on_machine_end(self, context: Dict[str, Any], final_output: Dict[str, Any]) -> Dict[str, Any]:
+        resp = await self._send("machine_end", {"context": context, "output": final_output})
+        if resp and "output" in resp:
+            return resp["output"]
+        return final_output
+
+    async def on_state_enter(self, state_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        resp = await self._send("state_enter", {"state": state_name, "context": context})
+        if resp and "context" in resp:
+            return resp["context"]
+        return context
+
+    async def on_state_exit(
+        self,
+        state_name: str,
+        context: Dict[str, Any],
+        output: Optional[Dict[str, Any]]
+    ) -> Optional[Dict[str, Any]]:
+        resp = await self._send("state_exit", {"state": state_name, "context": context, "output": output})
+        if resp and "output" in resp:
+            return resp["output"]
+        return output
+
+    async def on_transition(self, from_state: str, to_state: str, context: Dict[str, Any]) -> str:
+        resp = await self._send("transition", {"from": from_state, "to": to_state, "context": context})
+        if resp and "to_state" in resp:
+            return resp["to_state"]
+        return to_state
+
+    async def on_error(self, state_name: str, error: Exception, context: Dict[str, Any]) -> Optional[str]:
+        resp = await self._send("error", {
+            "state": state_name,
+            "error": str(error),
+            "error_type": type(error).__name__,
+            "context": context
+        })
+        if resp and "recovery_state" in resp:
+            return resp["recovery_state"]
+        return None  # Re-raise
+
+    async def on_action(self, action_name: str, context: Dict[str, Any]) -> Dict[str, Any]:
+        resp = await self._send("action", {"action": action_name, "context": context})
+        if resp and "context" in resp:
+            return resp["context"]
+        return context
+
+
+__all__ = [
+    "MachineHooks",
+    "LoggingHooks",
+    "MetricsHooks",
+    "CompositeHooks",
+    "WebhookHooks",
+]

flatagents/locking.py

@@ -0,0 +1,69 @@
+import fcntl
+import asyncio
+import os
+from abc import ABC, abstractmethod
+from typing import Optional
+from pathlib import Path
+import contextlib
+
+class ExecutionLock(ABC):
+    """Abstract interface for concurrency control."""
+    
+    @abstractmethod
+    async def acquire(self, key: str) -> bool:
+        """Acquire lock for key. Returns True if successful."""
+        pass
+        
+    @abstractmethod
+    async def release(self, key: str) -> None:
+        """Release lock for key."""
+        pass
+
+class LocalFileLock(ExecutionLock):
+    """
+    File-based lock using fcntl.flock.
+    Works on local filesystems and NFS (mostly).
+    NOT suited for distributed cloud storage (S3/GCS).
+    """
+    
+    def __init__(self, lock_dir: str = ".locks"):
+        self.lock_dir = Path(lock_dir)
+        self.lock_dir.mkdir(parents=True, exist_ok=True)
+        self._files = {}
+        
+    async def acquire(self, key: str) -> bool:
+        """Attempts to acquire a non-blocking exclusive lock."""
+        path = self.lock_dir / f"{key}.lock"
+        
+        try:
+            # Keep file handle open while locked
+            f = open(path, 'a+')
+            try:
+                # LOCK_EX | LOCK_NB = Exclusive, Non-Blocking
+                fcntl.flock(f.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                self._files[key] = f
+                return True
+            except (IOError, OSError):
+                f.close()
+                return False
+        except Exception:
+            return False
+            
+    async def release(self, key: str) -> None:
+        if key in self._files:
+            f = self._files.pop(key)
+            try:
+                fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+            finally:
+                f.close()
+                # Optional: unlink file? Usually simpler to leave it empty
+                # Path(f.name).unlink(missing_ok=True)
+
+class NoOpLock(ExecutionLock):
+    """Used when concurrency control is disabled or managed externally."""
+    
+    async def acquire(self, key: str) -> bool:
+        return True
+        
+    async def release(self, key: str) -> None:
+        pass