"""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
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: dict[str, Any] | None = None,
state_mapping: dict[str, str] | None = 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 == ">":
if field_value is None or expected_value is None:
return False
return field_value > expected_value
elif operator == ">=":
if field_value is None or expected_value is None:
return False
return field_value >= expected_value
elif operator == "<":
if field_value is None or expected_value is None:
return False
return field_value < expected_value
elif operator == "<=":
if field_value is None or expected_value is None:
return False
return field_value <= expected_value
elif operator == "in":
if field_value is None or expected_value is None:
return False
return field_value in expected_value
elif operator == "not in":
if field_value is None or expected_value is None:
return True
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
[docs]
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.
"""
[docs]
def __init__(self):
"""Initialize a workflow runner."""
self.workflows = {}
self.connections = []
[docs]
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}'")
[docs]
def connect_workflows(
self,
source_workflow_id: str,
target_workflow_id: str,
condition: dict[str, Any] | None = None,
state_mapping: dict[str, str] | None = 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}'"
)
[docs]
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
[docs]
def execute(
self,
entry_workflow_id: str,
initial_state: BaseModel,
task_manager: TaskManager | None = 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