stabilize 0.9.2__py3-none-any.whl

stabilize/context/__init__.py

@@ -0,0 +1,7 @@
+"""Context handling for pipeline execution."""
+
+from stabilize.context.stage_context import StageContext
+
+__all__ = [
+    "StageContext",
+]

stabilize/context/stage_context.py

@@ -0,0 +1,170 @@
+"""
+Stage context with ancestor output lookup.
+
+The StageContext class provides access to stage inputs with automatic
+fallback to ancestor stage outputs. This enables data flow between
+stages in a pipeline.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, MutableMapping
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+
+
+class StageContext(MutableMapping[str, Any]):
+    """
+    Stage context with ancestor output lookup.
+
+    When a key is not found in the stage's own context, it automatically
+    searches ancestor stages' outputs. This enables data to flow from
+    earlier stages to later ones.
+
+    Example:
+        # Stage A outputs: {"deploymentId": "abc123"}
+        # Stage B (downstream) context lookup:
+
+        context = StageContext(stage_b, stage_b.context)
+        deployment_id = context["deploymentId"]  # Returns "abc123" from stage A
+    """
+
+    def __init__(
+        self,
+        stage: StageExecution,
+        delegate: dict[str, Any],
+    ) -> None:
+        """
+        Initialize the context.
+
+        Args:
+            stage: The stage this context belongs to
+            delegate: The underlying context dictionary
+        """
+        self._stage = stage
+        self._delegate = delegate
+
+    @property
+    def stage(self) -> StageExecution:
+        """Get the stage this context belongs to."""
+        return self._stage
+
+    def __getitem__(self, key: str) -> Any:
+        """
+        Get a value by key.
+
+        First checks the stage's own context, then falls back to
+        ancestor outputs.
+        """
+        # Check own context first
+        if key in self._delegate:
+            return self._delegate[key]
+
+        # Search ancestor outputs
+        for ancestor in self._stage.ancestors():
+            if key in ancestor.outputs:
+                return ancestor.outputs[key]
+
+        raise KeyError(key)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Set a value in the context."""
+        self._delegate[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        """Delete a value from the context."""
+        del self._delegate[key]
+
+    def __iter__(self) -> Iterator[str]:
+        """Iterate over keys."""
+        return iter(self._delegate)
+
+    def __len__(self) -> int:
+        """Get number of keys in own context."""
+        return len(self._delegate)
+
+    def __contains__(self, key: object) -> bool:
+        """Check if key exists in context or ancestors."""
+        if key in self._delegate:
+            return True
+        if isinstance(key, str):
+            for ancestor in self._stage.ancestors():
+                if key in ancestor.outputs:
+                    return True
+        return False
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get with default value."""
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def get_current_only(self, key: str, default: Any = None) -> Any:
+        """
+        Get a value only from the current stage's context.
+
+        Does not search ancestor outputs.
+        """
+        return self._delegate.get(key, default)
+
+    def get_all(self, key: str) -> list[Any]:
+        """
+        Get all values of a key from ancestors.
+
+        Returns values sorted by proximity (closest ancestor first).
+        """
+        values = []
+        for ancestor in self._stage.ancestors():
+            if key in ancestor.outputs:
+                values.append(ancestor.outputs[key])
+        return values
+
+    def to_dict(self) -> dict[str, Any]:
+        """Get the underlying dictionary."""
+        return dict(self._delegate)
+
+    def merge_with_ancestors(self) -> dict[str, Any]:
+        """
+        Get merged context including all ancestor outputs.
+
+        Returns a new dictionary with all keys from this context
+        and all ancestor outputs. Own values take precedence.
+        """
+        merged = {}
+
+        # Start with ancestor outputs (reversed so closest overwrites)
+        for ancestor in reversed(self._stage.ancestors()):
+            merged.update(ancestor.outputs)
+
+        # Own context takes precedence
+        merged.update(self._delegate)
+
+        return merged
+
+
+def create_merged_context(stage: StageExecution) -> StageContext:
+    """
+    Create a StageContext with merged ancestor data.
+
+    The returned context will have the stage's own context merged
+    with ancestor outputs, with the stage's values taking precedence.
+
+    Args:
+        stage: The stage to create context for
+
+    Returns:
+        A StageContext with merged data
+    """
+    merged = {}
+
+    # Collect ancestor outputs
+    for ancestor in reversed(stage.ancestors()):
+        merged.update(ancestor.outputs)
+
+    # Own context takes precedence
+    merged.update(stage.context)
+
+    return StageContext(stage, merged)

stabilize/dag/__init__.py

@@ -0,0 +1,15 @@
+"""DAG operations for pipeline execution."""
+
+from stabilize.dag.graph import StageGraphBuilder
+from stabilize.dag.topological import (
+    CircularDependencyError,
+    topological_sort,
+    topological_sort_all_stages,
+)
+
+__all__ = [
+    "topological_sort",
+    "topological_sort_all_stages",
+    "CircularDependencyError",
+    "StageGraphBuilder",
+]

stabilize/dag/graph.py

@@ -0,0 +1,215 @@
+"""
+Stage graph builder for constructing synthetic stages.
+
+This module provides the StageGraphBuilder class for building graphs of
+synthetic stages (before/after stages) that are dynamically injected
+by StageDefinitionBuilders.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution, SyntheticStageOwner
+
+
+class StageGraphBuilder:
+    """
+    Builder for constructing graphs of synthetic stages.
+
+    Used by StageDefinitionBuilder to add before/after stages to a parent stage.
+    Manages dependencies between synthetic stages and connects them properly.
+
+    Example:
+        # In a StageDefinitionBuilder.before_stages():
+        def before_stages(self, stage, graph):
+            # Add setup stage
+            setup = StageExecution.create_synthetic(
+                type="setup",
+                name="Setup",
+                parent=stage,
+                owner=SyntheticStageOwner.STAGE_BEFORE,
+            )
+            graph.add(setup)
+
+            # Add validation stage after setup
+            validation = StageExecution.create_synthetic(
+                type="validate",
+                name="Validate",
+                parent=stage,
+                owner=SyntheticStageOwner.STAGE_BEFORE,
+            )
+            graph.add(validation)
+            graph.connect(setup, validation)
+    """
+
+    def __init__(
+        self,
+        parent: StageExecution,
+        owner: SyntheticStageOwner,
+        requisite_stage_ref_ids: set[str] | None = None,
+    ):
+        """
+        Initialize a StageGraphBuilder.
+
+        Args:
+            parent: The parent stage these synthetic stages belong to
+            owner: STAGE_BEFORE or STAGE_AFTER
+            requisite_stage_ref_ids: Initial requisites for first stage
+        """
+
+        self.parent = parent
+        self.owner = owner
+        self.requisite_stage_ref_ids = requisite_stage_ref_ids or set()
+        self._stages: list[StageExecution] = []
+        self._last_stage: StageExecution | None = None
+
+    @classmethod
+    def before_stages(cls, parent: StageExecution) -> StageGraphBuilder:
+        """
+        Create a builder for before stages.
+
+        Before stages run before the parent's tasks.
+        """
+        from stabilize.models.stage import SyntheticStageOwner
+
+        return cls(parent, SyntheticStageOwner.STAGE_BEFORE)
+
+    @classmethod
+    def after_stages(
+        cls,
+        parent: StageExecution,
+        requisite_stage_ref_ids: set[str] | None = None,
+    ) -> StageGraphBuilder:
+        """
+        Create a builder for after stages.
+
+        After stages run after the parent completes.
+        """
+        from stabilize.models.stage import SyntheticStageOwner
+
+        return cls(
+            parent,
+            SyntheticStageOwner.STAGE_AFTER,
+            requisite_stage_ref_ids or set(),
+        )
+
+    def add(self, stage: StageExecution) -> StageGraphBuilder:
+        """
+        Add a stage to the graph.
+
+        The stage will be configured as a synthetic stage of the parent.
+        If this is the first stage and requisite_stage_ref_ids is set,
+        those requisites will be applied.
+
+        Args:
+            stage: The stage to add
+
+        Returns:
+            self for method chaining
+        """
+        # Configure as synthetic stage
+        stage.parent_stage_id = self.parent.id
+        stage.synthetic_stage_owner = self.owner
+
+        # Set execution reference if parent has one
+        if self.parent.has_execution():
+            stage._execution = self.parent._execution
+
+        # Apply initial requisites to first stage
+        if not self._stages and self.requisite_stage_ref_ids:
+            stage.requisite_stage_ref_ids = stage.requisite_stage_ref_ids.union(self.requisite_stage_ref_ids)
+
+        self._stages.append(stage)
+        self._last_stage = stage
+
+        return self
+
+    def append(self, stage: StageExecution) -> StageGraphBuilder:
+        """
+        Append a stage after the last added stage.
+
+        Creates a sequential dependency from the last stage to this one.
+
+        Args:
+            stage: The stage to append
+
+        Returns:
+            self for method chaining
+        """
+        if self._last_stage:
+            self.connect(self._last_stage, stage)
+        return self.add(stage)
+
+    def connect(self, previous: StageExecution, next_stage: StageExecution) -> None:
+        """
+        Connect two stages with a dependency.
+
+        The next stage will depend on the previous stage.
+
+        Args:
+            previous: The stage that must complete first
+            next_stage: The stage that depends on previous
+        """
+        requisites = set(next_stage.requisite_stage_ref_ids)
+        requisites.add(previous.ref_id)
+        next_stage.requisite_stage_ref_ids = requisites
+
+    def build(self) -> list[StageExecution]:
+        """
+        Build and return the list of stages.
+
+        Returns:
+            List of configured synthetic stages
+        """
+        return list(self._stages)
+
+    @property
+    def stages(self) -> list[StageExecution]:
+        """Get the current list of stages."""
+        return list(self._stages)
+
+    @property
+    def last_stage(self) -> StageExecution | None:
+        """Get the last added stage."""
+        return self._last_stage
+
+    def is_empty(self) -> bool:
+        """Check if no stages have been added."""
+        return len(self._stages) == 0
+
+
+def connect_stages_linearly(stages: list[StageExecution]) -> None:
+    """
+    Connect a list of stages in linear sequence.
+
+    Each stage will depend on the previous one.
+
+    Args:
+        stages: List of stages to connect
+    """
+    for i in range(1, len(stages)):
+        previous = stages[i - 1]
+        current = stages[i]
+        requisites = set(current.requisite_stage_ref_ids)
+        requisites.add(previous.ref_id)
+        current.requisite_stage_ref_ids = requisites
+
+
+def add_stages_to_execution(
+    stages: list[StageExecution],
+    add_stage: Callable[[StageExecution], None],
+) -> None:
+    """
+    Add multiple stages to an execution.
+
+    Calls the provided add_stage function for each stage.
+
+    Args:
+        stages: List of stages to add
+        add_stage: Function to call to add each stage
+    """
+    for stage in stages:
+        add_stage(stage)

stabilize/dag/topological.py

@@ -0,0 +1,199 @@
+"""
+Topological sort for stage execution ordering.
+
+This module implements a topological sort algorithm for stages based on their
+requisite_stage_ref_ids (DAG edges).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+
+
+class CircularDependencyError(Exception):
+    """
+    Raised when a circular dependency is detected in the stage graph.
+
+    This indicates an invalid pipeline configuration where stages depend
+    on each other in a cycle.
+    """
+
+    def __init__(self, message: str, stages: list[StageExecution] | None = None):
+        super().__init__(message)
+        self.stages = stages or []
+
+
+def topological_sort(
+    stages: list[StageExecution],
+    stage_filter: Callable[[StageExecution], bool] = lambda s: s.parent_stage_id is None,
+) -> list[StageExecution]:
+    """
+    Sort stages into topological order based on their dependencies.
+
+    The algorithm:
+
+    1. Starts with all unsorted stages (filtered by predicate)
+    2. Finds stages whose requisites are all in the "processed" set
+    3. Adds those stages to result and their ref_ids to processed set
+    4. Repeats until all stages are sorted
+    5. Raises CircularDependencyError if no progress can be made
+
+    Args:
+        stages: List of stages to sort
+        stage_filter: Predicate to filter stages (default: exclude synthetic stages)
+
+    Returns:
+        List of stages sorted in execution order
+
+    Raises:
+        CircularDependencyError: If stages have circular dependencies
+
+    Example:
+        # Linear: A -> B -> C
+        stages = [stage_c, stage_a, stage_b]
+        sorted_stages = topological_sort(stages)
+        # Result: [stage_a, stage_b, stage_c]
+
+        # Parallel with join: A -> [B, C] -> D
+        # B and C have requisites [A], D has requisites [B, C]
+        sorted_stages = topological_sort(stages)
+        # Result: [A, B, C, D] or [A, C, B, D] (B and C can be in any order)
+    """
+    # Filter stages by predicate
+    unsorted: list[StageExecution] = [s for s in stages if stage_filter(s)]
+    sorted_stages: list[StageExecution] = []
+    ref_ids: set[str] = set()
+
+    while unsorted:
+        # Find all stages whose requisites have been satisfied
+        # A stage is sortable if all its requisite_stage_ref_ids are in ref_ids
+        sortable = [stage for stage in unsorted if ref_ids.issuperset(stage.requisite_stage_ref_ids)]
+
+        if not sortable:
+            # No progress possible - circular dependency
+            relationships = ", ".join(f"{list(stage.requisite_stage_ref_ids)}->{stage.ref_id}" for stage in stages)
+            raise CircularDependencyError(
+                f"Invalid stage relationships found: {relationships}",
+                stages=unsorted,
+            )
+
+        # Add all sortable stages to result
+        for stage in sortable:
+            unsorted.remove(stage)
+            ref_ids.add(stage.ref_id)
+            sorted_stages.append(stage)
+
+    return sorted_stages
+
+
+def topological_sort_all_stages(stages: list[StageExecution]) -> list[StageExecution]:
+    """
+    Sort all stages including synthetic stages.
+
+    Unlike topological_sort(), this does not filter out synthetic stages.
+
+    Args:
+        stages: List of stages to sort
+
+    Returns:
+        List of all stages sorted in execution order
+    """
+    return topological_sort(stages, stage_filter=lambda s: True)
+
+
+def validate_dag(stages: list[StageExecution]) -> bool:
+    """
+    Validate that stages form a valid DAG.
+
+    Returns True if valid, raises CircularDependencyError if invalid.
+
+    Args:
+        stages: List of stages to validate
+
+    Returns:
+        True if DAG is valid
+
+    Raises:
+        CircularDependencyError: If stages have circular dependencies
+    """
+    topological_sort(stages)
+    return True
+
+
+def find_initial_stages(stages: list[StageExecution]) -> list[StageExecution]:
+    """
+    Find all initial stages (those with no dependencies and not synthetic).
+
+    Args:
+        stages: List of stages to search
+
+    Returns:
+        List of initial stages
+    """
+    return [stage for stage in stages if stage.is_initial() and not stage.is_synthetic()]
+
+
+def find_terminal_stages(stages: list[StageExecution]) -> list[StageExecution]:
+    """
+    Find all terminal stages (those with no downstream stages and not synthetic).
+
+    Terminal stages are the last stages in the pipeline - no other stages
+    depend on them.
+
+    Args:
+        stages: List of stages to search
+
+    Returns:
+        List of terminal stages
+    """
+    # Get all ref_ids that are dependencies
+    all_requisites: set[str] = set()
+    for stage in stages:
+        all_requisites.update(stage.requisite_stage_ref_ids)
+
+    # Terminal stages are those whose ref_id is not in any requisite set
+    return [stage for stage in stages if not stage.is_synthetic() and stage.ref_id not in all_requisites]
+
+
+def get_execution_layers(stages: list[StageExecution]) -> list[list[StageExecution]]:
+    """
+    Group stages into execution layers.
+
+    Stages in the same layer can execute in parallel.
+    Each layer depends only on stages in previous layers.
+
+    Args:
+        stages: List of stages to group
+
+    Returns:
+        List of layers, where each layer is a list of stages that can run in parallel
+
+    Example:
+        # A -> [B, C] -> D
+        # Layer 0: [A]
+        # Layer 1: [B, C]
+        # Layer 2: [D]
+    """
+    unsorted: list[StageExecution] = [s for s in stages if s.parent_stage_id is None]
+    layers: list[list[StageExecution]] = []
+    ref_ids: set[str] = set()
+
+    while unsorted:
+        # Find all stages whose requisites are satisfied
+        layer = [stage for stage in unsorted if ref_ids.issuperset(stage.requisite_stage_ref_ids)]
+
+        if not layer:
+            # This shouldn't happen if topological_sort passes
+            break
+
+        layers.append(layer)
+
+        for stage in layer:
+            unsorted.remove(stage)
+            ref_ids.add(stage.ref_id)
+
+    return layers

stabilize/examples/__init__.py

@@ -0,0 +1 @@
+"""Stabilize example pipelines for RAG context."""