kailash 0.1.0__py3-none-any.whl

kailash/workflow/runner.py

@@ -0,0 +1,302 @@
+"""Workflow runner for executing connected workflows.
+
+This module provides tools for connecting and executing multiple workflows,
+allowing for complex multi-stage processing pipelines.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Tuple
+
+from pydantic import BaseModel
+
+from kailash.sdk_exceptions import WorkflowExecutionError
+from kailash.tracking import TaskManager
+from kailash.workflow.graph import Workflow
+
+logger = logging.getLogger(__name__)
+
+
+class WorkflowConnection:
+    """Defines a connection between two workflows."""
+
+    def __init__(
+        self,
+        source_workflow_id: str,
+        target_workflow_id: str,
+        condition: Optional[Dict[str, Any]] = None,
+        state_mapping: Optional[Dict[str, str]] = None,
+    ):
+        """Initialize a workflow connection.
+
+        Args:
+            source_workflow_id: ID of the source workflow
+            target_workflow_id: ID of the target workflow
+            condition: Optional condition for when this connection should be followed
+            state_mapping: Optional mapping of state fields between workflows
+        """
+        self.source_workflow_id = source_workflow_id
+        self.target_workflow_id = target_workflow_id
+        self.condition = condition or {}
+        self.state_mapping = state_mapping or {}
+
+    def should_follow(self, state: BaseModel) -> bool:
+        """Check if this connection should be followed based on state.
+
+        Args:
+            state: The current state object
+
+        Returns:
+            True if the connection should be followed, False otherwise
+        """
+        if not self.condition:
+            # If no condition is specified, always follow the connection
+            return True
+
+        # Extract condition field and value from the state
+        field_name = self.condition.get("field")
+        operator = self.condition.get("operator", "==")
+        expected_value = self.condition.get("value")
+
+        if not field_name:
+            # If no field name is specified, always follow the connection
+            return True
+
+        # Get the field value from the state
+        field_value = getattr(state, field_name, None)
+
+        # Check the condition
+        if operator == "==":
+            return field_value == expected_value
+        elif operator == "!=":
+            return field_value != expected_value
+        elif operator == ">":
+            return field_value > expected_value
+        elif operator == ">=":
+            return field_value >= expected_value
+        elif operator == "<":
+            return field_value < expected_value
+        elif operator == "<=":
+            return field_value <= expected_value
+        elif operator == "in":
+            return field_value in expected_value
+        elif operator == "not in":
+            return field_value not in expected_value
+        else:
+            # Unknown operator, default to always follow
+            logger.warning(
+                f"Unknown condition operator: {operator}. Always following connection."
+            )
+            return True
+
+    def map_state(self, state: BaseModel) -> Dict[str, Any]:
+        """Map state fields according to the mapping configuration.
+
+        Args:
+            state: The current state object
+
+        Returns:
+            Dictionary with mapped state fields
+        """
+        if not self.state_mapping:
+            # If no mapping is specified, use the state as is
+            return {"state": state}
+
+        # Apply mappings
+        mapped_state = {}
+        for source_key, target_key in self.state_mapping.items():
+            if hasattr(state, source_key):
+                mapped_state[target_key] = getattr(state, source_key)
+
+        return mapped_state
+
+
+class WorkflowRunner:
+    """Manages execution across multiple connected workflows.
+
+    This class allows building complex processing pipelines by connecting
+    multiple workflows together, with conditional branching based on state.
+    """
+
+    def __init__(self):
+        """Initialize a workflow runner."""
+        self.workflows = {}
+        self.connections = []
+
+    def add_workflow(self, workflow_id: str, workflow: Workflow) -> None:
+        """Add a workflow to the runner.
+
+        Args:
+            workflow_id: Unique identifier for the workflow
+            workflow: Workflow instance
+
+        Raises:
+            ValueError: If a workflow with the given ID already exists
+        """
+        if workflow_id in self.workflows:
+            raise ValueError(f"Workflow with ID '{workflow_id}' already exists")
+
+        self.workflows[workflow_id] = workflow
+        logger.info(f"Added workflow '{workflow.name}' with ID '{workflow_id}'")
+
+    def connect_workflows(
+        self,
+        source_workflow_id: str,
+        target_workflow_id: str,
+        condition: Optional[Dict[str, Any]] = None,
+        state_mapping: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Connect two workflows.
+
+        Args:
+            source_workflow_id: ID of the source workflow
+            target_workflow_id: ID of the target workflow
+            condition: Optional condition for when this connection should be followed
+            state_mapping: Optional mapping of state fields between workflows
+
+        Raises:
+            ValueError: If any workflow ID is invalid
+        """
+        # Validate workflow IDs
+        if source_workflow_id not in self.workflows:
+            raise ValueError(
+                f"Source workflow with ID '{source_workflow_id}' not found"
+            )
+
+        if target_workflow_id not in self.workflows:
+            raise ValueError(
+                f"Target workflow with ID '{target_workflow_id}' not found"
+            )
+
+        # Create connection
+        connection = WorkflowConnection(
+            source_workflow_id=source_workflow_id,
+            target_workflow_id=target_workflow_id,
+            condition=condition,
+            state_mapping=state_mapping,
+        )
+
+        self.connections.append(connection)
+        logger.info(
+            f"Connected workflow '{source_workflow_id}' to '{target_workflow_id}'"
+        )
+
+    def get_next_workflows(
+        self, current_workflow_id: str, state: BaseModel
+    ) -> List[Tuple[str, Dict[str, Any]]]:
+        """Get the next workflows to execute based on current state.
+
+        Args:
+            current_workflow_id: ID of the current workflow
+            state: Current state object
+
+        Returns:
+            List of (workflow_id, mapped_state) tuples for next workflows
+        """
+        next_workflows = []
+
+        for connection in self.connections:
+            if connection.source_workflow_id == current_workflow_id:
+                if connection.should_follow(state):
+                    mapped_state = connection.map_state(state)
+                    next_workflows.append((connection.target_workflow_id, mapped_state))
+
+        return next_workflows
+
+    def execute(
+        self,
+        entry_workflow_id: str,
+        initial_state: BaseModel,
+        task_manager: Optional[TaskManager] = None,
+        max_steps: int = 10,  # Prevent infinite loops
+    ) -> Tuple[BaseModel, Dict[str, Dict[str, Any]]]:
+        """Execute a sequence of connected workflows.
+
+        Args:
+            entry_workflow_id: ID of the first workflow to execute
+            initial_state: Initial state for workflow execution
+            task_manager: Optional task manager for tracking
+            max_steps: Maximum number of workflow steps to execute
+
+        Returns:
+            Tuple of (final state, all results by workflow)
+
+        Raises:
+            WorkflowExecutionError: If workflow execution fails
+            ValueError: If entry workflow is not found
+        """
+        if entry_workflow_id not in self.workflows:
+            raise ValueError(f"Entry workflow with ID '{entry_workflow_id}' not found")
+
+        # Initialize execution
+        current_workflow_id = entry_workflow_id
+        current_state = initial_state
+        all_results = {}
+        executed_workflows = set()
+        step_count = 0
+
+        # Execute workflows until no more connections to follow
+        while current_workflow_id and step_count < max_steps:
+            step_count += 1
+            logger.info(
+                f"Executing workflow '{current_workflow_id}' (step {step_count}/{max_steps})"
+            )
+
+            # Get the workflow
+            workflow = self.workflows[current_workflow_id]
+
+            # Track executed workflows to detect cycles
+            if current_workflow_id in executed_workflows:
+                logger.warning(
+                    f"Cycle detected in workflow execution: already executed '{current_workflow_id}'"
+                )
+                # Continue to next workflow rather than stopping, to handle intentional cycles
+
+            executed_workflows.add(current_workflow_id)
+
+            try:
+                # Execute the workflow
+                final_state, workflow_results = workflow.execute_with_state(
+                    state_model=current_state, task_manager=task_manager
+                )
+
+                # Store results
+                all_results[current_workflow_id] = workflow_results
+
+                # Update current state
+                current_state = final_state
+
+                # Find next workflows
+                next_workflows = self.get_next_workflows(
+                    current_workflow_id, current_state
+                )
+
+                if not next_workflows:
+                    # No more workflows to execute
+                    logger.info(
+                        f"No more workflows to execute after '{current_workflow_id}'"
+                    )
+                    break
+
+                # Take the first matching workflow as the next one
+                current_workflow_id = next_workflows[0][0]
+
+                # Apply state mapping if needed
+                if next_workflows[0][1]:
+                    # If a complete state object is provided, use it
+                    if "state" in next_workflows[0][1] and isinstance(
+                        next_workflows[0][1]["state"], BaseModel
+                    ):
+                        current_state = next_workflows[0][1]["state"]
+                    # Otherwise, merge the mapped values into the current state
+                    # using StateManager would be ideal here, but keeping it simple for now
+
+            except Exception as e:
+                logger.error(f"Error executing workflow '{current_workflow_id}': {e}")
+                raise WorkflowExecutionError(
+                    f"Failed to execute workflow '{current_workflow_id}': {e}"
+                ) from e
+
+        if step_count >= max_steps:
+            logger.warning(f"Reached maximum steps ({max_steps}) in workflow execution")
+
+        return current_state, all_results

kailash/workflow/state.py

@@ -0,0 +1,238 @@
+"""State management for workflow execution.
+
+This module provides tools for managing immutable state throughout workflow execution,
+making it easier to handle state transitions in a predictable manner.
+"""
+
+import logging
+from copy import deepcopy
+from typing import Any, Generic, List, Tuple, TypeVar
+
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+# Type variable for the state model
+StateT = TypeVar("StateT", bound=BaseModel)
+
+
+class StateManager:
+    """Manages immutable state operations for workflow execution.
+
+    This class provides utilities for updating state objects immutably,
+    focusing on Pydantic models to ensure type safety and validation.
+    """
+
+    @staticmethod
+    def update_in(state_obj: BaseModel, path: List[str], value: Any) -> BaseModel:
+        """Update a nested property in the state and return a new state object.
+
+        Args:
+            state_obj: The Pydantic model state object
+            path: List of attribute names forming a path to the property to update
+            value: The new value to set
+
+        Returns:
+            A new state object with the update applied
+
+        Raises:
+            TypeError: If state_obj is not a Pydantic BaseModel
+            KeyError: If the path is invalid
+        """
+        if not isinstance(state_obj, BaseModel):
+            raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
+
+        # Create deep copy
+        new_state = state_obj.model_copy(deep=True)
+
+        # For simple top-level updates
+        if len(path) == 1:
+            setattr(new_state, path[0], value)
+            return new_state
+
+        # For nested updates
+        current = new_state
+        for i, key in enumerate(path[:-1]):
+            if not hasattr(current, key):
+                raise KeyError(f"Invalid path: {'.'.join(path[:i+1])}")
+
+            # Get the next level object and ensure we're working with a copy
+            next_obj = getattr(current, key)
+            if isinstance(next_obj, BaseModel):
+                next_obj = next_obj.model_copy(deep=True)
+                setattr(current, key, next_obj)
+            elif isinstance(next_obj, dict):
+                next_obj = deepcopy(next_obj)
+                setattr(current, key, next_obj)
+            elif isinstance(next_obj, list):
+                next_obj = deepcopy(next_obj)
+                setattr(current, key, next_obj)
+
+            current = next_obj
+
+        # Set the final value
+        if hasattr(current, path[-1]):
+            setattr(current, path[-1], value)
+        else:
+            raise KeyError(f"Invalid path: {'.'.join(path)}")
+
+        return new_state
+
+    @staticmethod
+    def batch_update(
+        state_obj: BaseModel, updates: List[Tuple[List[str], Any]]
+    ) -> BaseModel:
+        """Apply multiple updates to the state atomically.
+
+        Args:
+            state_obj: The Pydantic model state object
+            updates: List of (path, value) tuples with updates to apply
+
+        Returns:
+            A new state object with all updates applied
+
+        Raises:
+            TypeError: If state_obj is not a Pydantic BaseModel
+            KeyError: If any path is invalid
+        """
+        if not isinstance(state_obj, BaseModel):
+            raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
+
+        # Create deep copy
+        new_state = state_obj.model_copy(deep=True)
+
+        # Apply each update
+        for path, value in updates:
+            new_state = StateManager.update_in(new_state, path, value)
+
+        return new_state
+
+    @staticmethod
+    def get_in(state_obj: BaseModel, path: List[str]) -> Any:
+        """Get the value at a nested path.
+
+        Args:
+            state_obj: The Pydantic model state object
+            path: List of attribute names forming a path to the property to retrieve
+
+        Returns:
+            The value at the specified path
+
+        Raises:
+            TypeError: If state_obj is not a Pydantic BaseModel
+            KeyError: If the path is invalid
+        """
+        if not isinstance(state_obj, BaseModel):
+            raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
+
+        # For simple top-level properties
+        if len(path) == 1:
+            if not hasattr(state_obj, path[0]):
+                raise KeyError(f"Invalid path: {path[0]}")
+            return getattr(state_obj, path[0])
+
+        # For nested properties
+        current = state_obj
+        for i, key in enumerate(path):
+            if not hasattr(current, key):
+                raise KeyError(f"Invalid path: {'.'.join(path[:i+1])}")
+            current = getattr(current, key)
+
+        return current
+
+    @staticmethod
+    def merge(state_obj: BaseModel, **updates) -> BaseModel:
+        """Merge flat updates into state and return a new state.
+
+        Args:
+            state_obj: The Pydantic model state object
+            **updates: Attribute updates to apply to the top level
+
+        Returns:
+            A new state object with the updates applied
+
+        Raises:
+            TypeError: If state_obj is not a Pydantic BaseModel
+        """
+        if not isinstance(state_obj, BaseModel):
+            raise TypeError(f"Expected BaseModel, got {type(state_obj)}")
+
+        return state_obj.model_copy(update=updates)
+
+
+class WorkflowStateWrapper(Generic[StateT]):
+    """Wraps a state object with convenient update methods for use in workflows.
+
+    This wrapper provides a clean interface for immutable state updates
+    within workflow nodes, simplifying state management.
+    """
+
+    def __init__(self, state: StateT):
+        """Initialize the state wrapper.
+
+        Args:
+            state: The Pydantic model state object to wrap
+        """
+        self._state = state
+
+    def update_in(self, path: List[str], value: Any) -> "WorkflowStateWrapper[StateT]":
+        """Update state at path and return new wrapper.
+
+        Args:
+            path: List of attribute names forming a path to the property to update
+            value: The new value to set
+
+        Returns:
+            A new state wrapper with the update applied
+        """
+        new_state = StateManager.update_in(self._state, path, value)
+        return WorkflowStateWrapper(new_state)
+
+    def batch_update(
+        self, updates: List[Tuple[List[str], Any]]
+    ) -> "WorkflowStateWrapper[StateT]":
+        """Apply multiple updates to the state atomically.
+
+        Args:
+            updates: List of (path, value) tuples with updates to apply
+
+        Returns:
+            A new state wrapper with all updates applied
+        """
+        new_state = StateManager.batch_update(self._state, updates)
+        return WorkflowStateWrapper(new_state)
+
+    def get_in(self, path: List[str]) -> Any:
+        """Get the value at a nested path.
+
+        Args:
+            path: List of attribute names forming a path to the property to retrieve
+
+        Returns:
+            The value at the specified path
+        """
+        return StateManager.get_in(self._state, path)
+
+    def merge(self, **updates) -> "WorkflowStateWrapper[StateT]":
+        """Merge flat updates into state and return a new wrapper.
+
+        Args:
+            **updates: Attribute updates to apply to the top level
+
+        Returns:
+            A new state wrapper with the updates applied
+        """
+        new_state = StateManager.merge(self._state, **updates)
+        return WorkflowStateWrapper(new_state)
+
+    def get_state(self) -> StateT:
+        """Get the wrapped state object.
+
+        Returns:
+            The current state object
+        """
+        return self._state
+
+    def __repr__(self) -> str:
+        """Get string representation."""
+        return f"WorkflowStateWrapper({self._state})"