hexdag 0.5.0.dev1__py3-none-any.whl

hexdag/core/orchestration/components/checkpoint_manager.py

@@ -0,0 +1,134 @@
+"""CheckpointManager component for orchestrator state persistence.
+
+Storage-agnostic checkpoint manager using Memory Port for maximum flexibility.
+Supports any backend: SQL databases, files (JSON/YAML), Redis, S3, etc.
+"""
+
+from hexdag.core.domain.dag import DirectedGraph, NodeSpec
+from hexdag.core.orchestration.models import CheckpointState
+from hexdag.core.ports.memory import Memory
+
+
+class CheckpointManager:
+    """Manages orchestrator checkpoints using Memory Port abstraction.
+
+    This implementation is storage-agnostic and works with any Memory backend:
+    - SQL databases (via SQLiteMemoryAdapter)
+    - File storage (JSON, YAML, pickle via FileMemoryAdapter)
+    - In-memory storage (for testing)
+    - Redis, S3, etc.
+
+    Responsibilities:
+    - Save/restore execution state
+    - Filter graphs for resume
+    - Automatic serialization via Pydantic
+
+    Parameters
+    ----------
+    storage : Memory
+        Memory port implementation for storage backend
+    key_prefix : str, default="checkpoint:"
+        Prefix for checkpoint keys (useful for namespacing)
+    auto_checkpoint : bool, default=True
+        Auto-save after nodes complete
+
+    Examples
+    --------
+    In-memory storage (testing)::
+
+        storage = InMemoryMemory()
+        mgr = CheckpointManager(storage=storage)
+        await mgr.save(state)
+        restored = await mgr.load("run-123")
+
+    File-based storage (production)::
+
+        storage = FileMemoryAdapter(base_path="./checkpoints", format="json")
+        mgr = CheckpointManager(storage=storage)
+
+    Database storage (enterprise)::
+
+        db = SQLiteAdapter(db_path="hexdag.db")
+        storage = SQLiteMemoryAdapter(database=db)
+        mgr = CheckpointManager(storage=storage)
+    """
+
+    def __init__(
+        self,
+        storage: Memory,
+        key_prefix: str = "checkpoint:",
+        auto_checkpoint: bool = True,
+    ):
+        self.storage = storage
+        self.key_prefix = key_prefix
+        self.auto_checkpoint = auto_checkpoint
+
+    def _make_key(self, run_id: str) -> str:
+        """Generate storage key for a run_id."""
+        return f"{self.key_prefix}{run_id}"
+
+    async def save(self, state: CheckpointState) -> None:
+        """Save checkpoint state.
+
+        Uses Pydantic's model_dump_json() for automatic serialization.
+        All complex types (datetime, nested models) are handled automatically.
+
+        Parameters
+        ----------
+        state : CheckpointState
+            Complete checkpoint state to persist
+        """
+        key = self._make_key(state.run_id)
+        # Pydantic handles all serialization including datetime, nested models, etc.
+        serialized = state.model_dump_json()
+        await self.storage.aset(key, serialized)
+
+    async def load(self, run_id: str) -> CheckpointState | None:
+        """Load checkpoint state by run_id.
+
+        Uses Pydantic's model_validate_json() for automatic deserialization.
+
+        Parameters
+        ----------
+        run_id : str
+            Run identifier to load
+
+        Returns
+        -------
+        CheckpointState | None
+            Restored checkpoint state, or None if not found
+        """
+        key = self._make_key(run_id)
+        serialized = await self.storage.aget(key)
+
+        if serialized is None:
+            return None
+
+        # Pydantic handles all deserialization and validation
+        return CheckpointState.model_validate_json(serialized)
+
+    def filter_completed(self, graph: DirectedGraph, completed: set[str]) -> DirectedGraph:
+        """Create graph with only pending nodes.
+
+        Parameters
+        ----------
+        graph : DirectedGraph
+            Original DAG
+        completed : set[str]
+            Set of completed node names
+
+        Returns
+        -------
+        DirectedGraph
+            New graph with only pending nodes and updated dependencies
+        """
+        pending = DirectedGraph()
+        for spec in graph:  # Using iterator instead of .nodes.items()
+            if spec.name not in completed:
+                pending += NodeSpec(  # Using += operator instead of .add()
+                    name=spec.name,
+                    fn=spec.fn,
+                    deps=frozenset(d for d in spec.deps if d not in completed),
+                    timeout=spec.timeout,
+                )
+        return pending

hexdag/core/orchestration/components/execution_coordinator.py

@@ -0,0 +1,360 @@
+"""Execution coordinator for observer notifications and input mapping.
+
+This module provides execution coordination functionality:
+
+- Observer notifications during execution
+- Input preparation and dependency mapping
+- Input mapping transformation (including $input syntax)
+"""
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from hexdag.core.ports.observer_manager import ObserverManagerPort
+else:
+    ObserverManagerPort = Any
+
+from hexdag.core.domain.dag import NodeSpec
+from hexdag.core.logging import get_logger
+
+__all__ = ["ExecutionCoordinator"]
+
+logger = get_logger(__name__)
+
+
+class ExecutionCoordinator:
+    """Coordinates execution context: observer notifications and input mapping.
+
+    This component handles two responsibilities:
+
+    1. **Observer Notifications**: Notifying observers of events during DAG execution.
+
+    2. **Input Mapping**: Preparing input data for nodes based on their dependencies.
+       Uses a smart mapping strategy:
+       - No dependencies → initial input
+       - Single dependency → pass through that result
+       - Multiple dependencies → dict of results
+
+    Examples
+    --------
+    Basic usage::
+
+        coordinator = ExecutionCoordinator()
+
+        # Notify observer of an event
+        await coordinator.notify_observer(observer_manager, NodeStarted(...))
+
+        # Prepare input for a node
+        input_data = coordinator.prepare_node_input(
+            node_spec, node_results, initial_input
+        )
+    """
+
+    # ========================================================================
+    # Observer Notifications (from PolicyCoordinator)
+    # ========================================================================
+
+    async def notify_observer(
+        self, observer_manager: ObserverManagerPort | None, event: Any
+    ) -> None:
+        """Notify observer manager of an event if it exists.
+
+        Parameters
+        ----------
+        observer_manager : ObserverManagerPort | None
+            Observer manager to notify (None if no observer configured)
+        event : Any
+            Event to send (typically NodeStarted, NodeCompleted, etc.)
+
+        Examples
+        --------
+        >>> from hexdag.core.orchestration.events import NodeStarted
+        >>> event = NodeStarted(name="my_node", wave_index=0)
+        >>> await coordinator.notify_observer(observer_manager, event)  # doctest: +SKIP
+        """
+        if observer_manager:
+            await observer_manager.notify(event)
+
+    # ========================================================================
+    # Input Mapping
+    # ========================================================================
+
+    def prepare_node_input(
+        self, node_spec: NodeSpec, node_results: dict[str, Any], initial_input: Any
+    ) -> Any:
+        """Prepare input data for node execution with simplified data mapping.
+
+        The mapping strategy is:
+        1. **No dependencies** → initial_input (entry point)
+        2. **Single dependency** → results[dependency_name] (pass-through)
+        3. **Multiple dependencies** → {dep1: result1, dep2: result2, ...} (namespace)
+
+        This approach balances simplicity (pass-through for single deps) with
+        clarity (named dict for multiple deps).
+
+        Parameters
+        ----------
+        node_spec : NodeSpec
+            Node specification containing dependencies
+        node_results : dict[str, Any]
+            Results from previously executed nodes
+        initial_input : Any
+            Initial input data for the pipeline
+
+        Returns
+        -------
+        Any
+            Prepared input data for the node:
+            - initial_input if no dependencies
+            - dependency result if single dependency
+            - dict of dependency results if multiple dependencies
+
+        Examples
+        --------
+        >>> coordinator = ExecutionCoordinator()
+        >>>
+        >>> # No dependencies - gets initial input
+        >>> # start_input = coordinator.prepare_node_input(
+        >>> #     NodeSpec("start", lambda x: x.upper()),
+        >>> #     node_results={},
+        >>> #     initial_input="hello"
+        >>> # )
+        >>> # start_input == "hello"
+        >>>
+        >>> # Single dependency - gets that result directly
+        >>> # process_input = coordinator.prepare_node_input(
+        >>> #     NodeSpec("process", lambda x: x + "!", deps={"start"}),
+        >>> #     node_results={"start": "HELLO"},
+        >>> #     initial_input="hello"
+        >>> # )
+        >>> # process_input == "HELLO"
+
+        Notes
+        -----
+        The multi-dependency dict preserves node names as keys, making it clear
+        where each piece of data came from. This is especially useful for
+        debugging and for nodes that need to treat different dependencies
+        differently.
+
+        If the node has an ``input_mapping`` in its params, the prepared input
+        will be transformed according to the mapping. This supports:
+        - ``$input.field`` - Reference the initial pipeline input
+        - ``node_name.field`` - Reference a specific dependency's output
+        """
+        # Prepare base input from dependencies
+        if not node_spec.deps:
+            base_input = initial_input
+        elif len(node_spec.deps) == 1:
+            dep_name = next(iter(node_spec.deps))
+            base_input = node_results.get(dep_name, initial_input)
+        else:
+            # Multiple dependencies - preserve namespace structure
+            base_input = {}
+            for dep_name in node_spec.deps:
+                if dep_name in node_results:
+                    base_input[dep_name] = node_results[dep_name]
+
+        # Apply input_mapping if present in node params
+        input_mapping = node_spec.params.get("input_mapping") if node_spec.params else None
+        if input_mapping:
+            return self._apply_input_mapping(base_input, input_mapping, initial_input, node_results)
+
+        return base_input
+
+    def _is_expression(self, source: str) -> bool:
+        """Check if a source string is an expression (contains function calls or operators).
+
+        Parameters
+        ----------
+        source : str
+            The source string to check
+
+        Returns
+        -------
+        bool
+            True if the source appears to be an expression
+        """
+        from hexdag.core.expression_parser import ALLOWED_FUNCTIONS
+
+        # Check for function call patterns (function_name followed by parenthesis)
+        for func_name in ALLOWED_FUNCTIONS:
+            if f"{func_name}(" in source:
+                return True
+
+        # Check for arithmetic/comparison operators (but not dots which are field paths)
+        # Be careful not to match operators in simple field paths
+        expression_indicators = [
+            "==",
+            "!=",
+            "<=",
+            ">=",
+            " < ",
+            " > ",
+            " + ",
+            " - ",
+            " * ",
+            " / ",
+            " % ",
+            " and ",
+            " or ",
+            " not ",
+            " in ",
+        ]
+        return any(op in source for op in expression_indicators)
+
+    def _apply_input_mapping(
+        self,
+        base_input: Any,
+        input_mapping: dict[str, str],
+        initial_input: Any,
+        node_results: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Apply field mapping to transform input data.
+
+        Supports multiple syntaxes:
+        - ``$input.field`` - Extract from the initial pipeline input
+        - ``node_name.field`` - Extract from a specific node's output
+        - Expression syntax - Use allowed functions and operators
+
+        Parameters
+        ----------
+        base_input : Any
+            The prepared input from dependencies (may be single value or dict)
+        input_mapping : dict[str, str]
+            Mapping of {target_field: "source_path"} or {target_field: "expression"}
+        initial_input : Any
+            The original pipeline input (for $input references)
+        node_results : dict[str, Any]
+            Results from all previously executed nodes
+
+        Returns
+        -------
+        dict[str, Any]
+            Transformed input with mapped fields
+
+        Examples
+        --------
+        >>> coordinator = ExecutionCoordinator()
+        >>> mapping = {"load_id": "$input.load_id", "result": "analyzer.output"}
+        >>> # This would extract load_id from initial input and result from analyzer node
+
+        Expression examples::
+
+            mapping = {
+                "is_valid": "len(items) > 0",
+                "name_upper": "upper(user.name)",
+                "total": "price * quantity",
+            }
+        """
+        from hexdag.builtin.nodes.mapped_input import FieldExtractor
+
+        result: dict[str, Any] = {}
+
+        for target_field, source_path in input_mapping.items():
+            # Check if this is an expression that needs evaluation
+            if self._is_expression(source_path):
+                value = self._evaluate_expression(
+                    source_path, base_input, initial_input, node_results
+                )
+            elif source_path.startswith("$input."):
+                # Extract from initial pipeline input
+                actual_path = source_path[7:]  # Remove "$input." prefix
+                if actual_path:
+                    # Has a field path like "$input.my_field"
+                    if isinstance(initial_input, dict):
+                        value = FieldExtractor.extract(initial_input, actual_path)
+                    else:
+                        # Non-dict input - wrap and extract
+                        value = FieldExtractor.extract({"_root": initial_input}, "_root")
+                else:
+                    # Just "$input." with no field - return entire initial input
+                    value = initial_input
+            elif source_path == "$input":
+                # Reference the entire initial input
+                value = initial_input
+            elif "." in source_path:
+                # Check if it's a node_name.field pattern
+                parts = source_path.split(".", 1)
+                node_name, field_path = parts[0], parts[1]
+                if node_name in node_results:
+                    # Extract from specific node's result
+                    value = FieldExtractor.extract(node_results[node_name], field_path)
+                else:
+                    # Fall back to extracting from base_input
+                    value = FieldExtractor.extract(
+                        base_input if isinstance(base_input, dict) else {}, source_path
+                    )
+            else:
+                # Simple field name - extract from base_input
+                value = FieldExtractor.extract(
+                    base_input if isinstance(base_input, dict) else {}, source_path
+                )
+
+            if value is None:
+                logger.warning(
+                    f"input_mapping: '{source_path}' resolved to None for target '{target_field}'"
+                )
+
+            result[target_field] = value
+
+        return result
+
+    def _evaluate_expression(
+        self,
+        expression: str,
+        base_input: Any,
+        initial_input: Any,
+        node_results: dict[str, Any],
+    ) -> Any:
+        """Evaluate an expression against available data.
+
+        Parameters
+        ----------
+        expression : str
+            The expression to evaluate (e.g., "len(items) > 0")
+        base_input : Any
+            The prepared input from dependencies
+        initial_input : Any
+            The original pipeline input
+        node_results : dict[str, Any]
+            Results from all previously executed nodes
+
+        Returns
+        -------
+        Any
+            The result of evaluating the expression
+        """
+        from hexdag.core.expression_parser import ExpressionError, evaluate_expression
+
+        # Build the data context for expression evaluation
+        # Merge all available data sources into a single dict
+        data_context: dict[str, Any] = {}
+
+        # Add node results
+        data_context.update(node_results)
+
+        # Add base_input (either as-is if dict, or wrapped)
+        if isinstance(base_input, dict):
+            data_context.update(base_input)
+        elif base_input is not None:
+            data_context["_input"] = base_input
+
+        # Add initial input with $input prefix removed (accessible as 'input')
+        if isinstance(initial_input, dict):
+            data_context["input"] = initial_input
+            # Also add initial_input fields at top level for convenience
+            for key, val in initial_input.items():
+                if key not in data_context:
+                    data_context[key] = val
+        elif initial_input is not None:
+            data_context["input"] = initial_input
+
+        try:
+            # Use evaluate_expression to get the actual value, not a boolean
+            return evaluate_expression(expression, data_context, {})
+        except ExpressionError as e:
+            logger.error(f"Expression evaluation failed for '{expression}': {e}")
+            return None
+        except Exception as e:
+            logger.error(f"Unexpected error evaluating expression '{expression}': {e}")
+            return None

hexdag/core/orchestration/components/health_check_manager.py

@@ -0,0 +1,176 @@
+"""Health check manager for pre-DAG adapter validation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from hexdag.core.ports.observer_manager import ObserverManagerPort
+else:
+    ObserverManagerPort = Any
+
+from hexdag.core.logging import get_logger
+from hexdag.core.orchestration.events import HealthCheckCompleted
+from hexdag.core.ports.healthcheck import HealthStatus
+from hexdag.core.protocols import HealthCheckable
+
+logger = get_logger(__name__)
+
+# Constants
+MANAGER_PORT_NAMES = frozenset({"observer_manager"})
+LATENCY_PRECISION = 1  # Decimal places for latency display
+
+
+class HealthCheckManager:
+    """Manages health checks on adapters before DAG execution.
+
+    Responsibilities:
+    - Check adapter health via ahealth_check() method
+    - Emit HealthCheckCompleted events
+    - Determine if unhealthy adapters should block execution
+
+    Examples
+    --------
+    Example usage::
+
+        manager = HealthCheckManager(fail_fast=True, warn_only=False)
+        health_results = await manager.check_all_adapters(
+            ports={"llm": openai, "database": postgres},
+            observer_manager=observer,
+            pipeline_name="my_pipeline"
+        )
+    """
+
+    def __init__(self, fail_fast: bool = False, warn_only: bool = True):
+        """Initialize health check manager.
+
+        Parameters
+        ----------
+        fail_fast : bool, default=False
+            If True, unhealthy adapters block pipeline execution
+        warn_only : bool, default=True
+            If True, log warnings for unhealthy adapters but don't block
+        """
+        self.fail_fast = fail_fast
+        self.warn_only = warn_only
+
+    async def check_all_adapters(
+        self,
+        ports: dict[str, Any],
+        observer_manager: ObserverManagerPort | None,
+        pipeline_name: str,
+    ) -> list[HealthStatus]:
+        """Run health checks on all adapters that implement ahealth_check().
+
+        Parameters
+        ----------
+        ports : dict[str, Any]
+            All available ports
+        observer_manager : ObserverManagerPort | None
+            Optional observer for event emission
+        pipeline_name : str
+            Name of the pipeline
+
+        Returns
+        -------
+        list[HealthStatus]
+            Health status results from all adapters
+        """
+        health_results = []
+
+        for port_name, adapter in ports.items():
+            # Skip non-adapter ports
+            if port_name in MANAGER_PORT_NAMES:
+                continue
+
+            if isinstance(adapter, HealthCheckable):
+                status = await self._check_single_adapter(port_name, adapter, observer_manager)
+                health_results.append(status)
+
+        return health_results
+
+    async def _check_single_adapter(
+        self,
+        port_name: str,
+        adapter: Any,
+        observer_manager: ObserverManagerPort | None,
+    ) -> HealthStatus:
+        """Check health of a single adapter.
+
+        Parameters
+        ----------
+        port_name : str
+            Name of the port
+        adapter : Any
+            Adapter instance
+        observer_manager : ObserverManagerPort | None
+            Optional observer for event emission
+
+        Returns
+        -------
+        HealthStatus
+            Health status of the adapter
+        """
+        try:
+            logger.debug(f"Running health check for {port_name}")
+            health_check = adapter.ahealth_check
+            status: HealthStatus = await health_check()  # pyright: ignore[reportGeneralTypeIssues]
+            status.port_name = port_name  # Ensure port name is set
+
+            # Emit event
+            if observer_manager:
+                event = HealthCheckCompleted(
+                    adapter_name=status.adapter_name,
+                    port_name=port_name,
+                    status=status,
+                )
+                await observer_manager.notify(event)
+
+            # Log result
+            self._log_health_result(port_name, status)
+
+            return status
+
+        except (RuntimeError, ConnectionError, TimeoutError, ValueError) as e:
+            # Health check errors - mark adapter as unhealthy
+            logger.error(f"Health check failed for {port_name}: {e}", exc_info=True)
+            adapter_name = getattr(adapter, "_hexdag_name", port_name)
+            return HealthStatus(
+                status="unhealthy",
+                adapter_name=adapter_name,
+                port_name=port_name,
+                error=e,
+            )
+
+    def _log_health_result(self, port_name: str, status: HealthStatus) -> None:
+        """Log health check result.
+
+        Parameters
+        ----------
+        port_name : str
+            Name of the port
+        status : HealthStatus
+            Health status result
+        """
+        if status.status == "healthy":
+            latency_info = (
+                f" ({status.latency_ms:.{LATENCY_PRECISION}f}ms)" if status.latency_ms else ""
+            )
+            logger.info(f"✅ {port_name} health check: {status.status}{latency_info}")
+        else:
+            logger.warning(f"⚠️ {port_name} health check: {status.status} - {status.error}")
+
+    def get_unhealthy_adapters(self, health_results: list[HealthStatus]) -> list[HealthStatus]:
+        """Filter health results to only unhealthy adapters.
+
+        Parameters
+        ----------
+        health_results : list[HealthStatus]
+            All health check results
+
+        Returns
+        -------
+        list[HealthStatus]
+            Only the unhealthy adapters
+        """
+        return [h for h in health_results if h.status == "unhealthy"]