rappel 0.7.2__py3-none-macosx_15_0_arm64.whl

rappel/workflow.py

@@ -0,0 +1,251 @@
+"""
+Workflow base class and registration decorator.
+
+This module provides the foundation for defining workflows that can be
+compiled to IR and executed by the Rappel runtime.
+"""
+
+import hashlib
+import inspect
+import os
+from dataclasses import dataclass
+from datetime import timedelta
+from functools import wraps
+from threading import RLock
+from typing import Any, Awaitable, ClassVar, Optional, TypeVar
+
+from proto import ast_pb2 as ir
+from proto import messages_pb2 as pb2
+
+from . import bridge
+from .actions import deserialize_result_payload
+from .ir_builder import build_workflow_ir
+from .logger import configure as configure_logger
+from .serialization import build_arguments_from_kwargs
+from .workflow_runtime import WorkflowNodeResult
+
+logger = configure_logger("rappel.workflow")
+
+TWorkflow = TypeVar("TWorkflow", bound="Workflow")
+TResult = TypeVar("TResult")
+
+
+@dataclass(frozen=True)
+class RetryPolicy:
+    """Retry policy for action execution.
+
+    Maps to IR RetryPolicy: [ExceptionType -> retry: N, backoff: Xs]
+
+    Args:
+        attempts: Maximum number of retry attempts.
+        exception_types: List of exception type names to retry on. Empty = catch all.
+        backoff_seconds: Constant backoff duration between retries in seconds.
+    """
+
+    attempts: Optional[int] = None
+    exception_types: Optional[list[str]] = None
+    backoff_seconds: Optional[float] = None
+
+
+class Workflow:
+    """Base class for workflow definitions."""
+
+    name: ClassVar[Optional[str]] = None
+    """Human-friendly identifier. Override to pin the registry key; defaults to lowercase class name."""
+
+    concurrent: ClassVar[bool] = False
+    """When True, downstream engines may respect DAG-parallel execution; False preserves sequential semantics."""
+
+    _workflow_ir: ClassVar[Optional[ir.Program]] = None
+    _ir_lock: ClassVar[RLock] = RLock()
+    _workflow_version_id: ClassVar[Optional[str]] = None
+
+    async def run(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError
+
+    @classmethod
+    def _normalize_run_inputs(cls, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+        try:
+            run_impl = cls.__workflow_run_impl__  # type: ignore[attr-defined]
+        except AttributeError:
+            run_impl = cls.run
+        sig = inspect.signature(run_impl)
+        params = list(sig.parameters.keys())[1:]  # Skip 'self'
+
+        normalized = dict(kwargs)
+        for i, arg in enumerate(args):
+            if i < len(params):
+                normalized[params[i]] = arg
+
+        bound = sig.bind_partial(None, **normalized)
+        bound.apply_defaults()
+        return {key: value for key, value in bound.arguments.items() if key != "self"}
+
+    @classmethod
+    def _build_initial_context(
+        cls, args: tuple[Any, ...], kwargs: dict[str, Any]
+    ) -> pb2.WorkflowArguments:
+        initial_kwargs = cls._normalize_run_inputs(args, kwargs)
+        return build_arguments_from_kwargs(initial_kwargs)
+
+    async def run_action(
+        self,
+        awaitable: Awaitable[TResult],
+        *,
+        retry: Optional[RetryPolicy] = None,
+        timeout: Optional[float | int | timedelta] = None,
+    ) -> TResult:
+        """Helper that simply awaits the provided action coroutine.
+
+        The retry and timeout arguments are consumed by the workflow compiler
+        (IR builder) rather than the runtime execution path.
+
+        Args:
+            awaitable: The action coroutine to execute.
+            retry: Retry policy including max attempts, exception types, and backoff.
+            timeout: Timeout duration in seconds (or timedelta).
+        """
+        # Parameters are intentionally unused at runtime; the workflow compiler
+        # inspects the AST to record them.
+        del retry, timeout
+        return await awaitable
+
+    @classmethod
+    def short_name(cls) -> str:
+        if cls.name:
+            return cls.name
+        return cls.__name__.lower()
+
+    @classmethod
+    def workflow_ir(cls) -> ir.Program:
+        """Build and cache the IR program for this workflow."""
+        if cls._workflow_ir is None:
+            with cls._ir_lock:
+                if cls._workflow_ir is None:
+                    cls._workflow_ir = build_workflow_ir(cls)
+        return cls._workflow_ir
+
+    @classmethod
+    def _build_registration_payload(
+        cls, initial_context: Optional[pb2.WorkflowArguments] = None
+    ) -> pb2.WorkflowRegistration:
+        """Build a registration payload with the serialized IR."""
+        program = cls.workflow_ir()
+
+        # Serialize IR to bytes
+        ir_bytes = program.SerializeToString()
+        ir_hash = hashlib.sha256(ir_bytes).hexdigest()
+
+        message = pb2.WorkflowRegistration(
+            workflow_name=cls.short_name(),
+            ir=ir_bytes,
+            ir_hash=ir_hash,
+            concurrent=cls.concurrent,
+        )
+
+        if initial_context:
+            message.initial_context.CopyFrom(initial_context)
+
+        return message
+
+
+class WorkflowRegistry:
+    """Registry of workflow definitions keyed by workflow name."""
+
+    def __init__(self) -> None:
+        self._workflows: dict[str, type[Workflow]] = {}
+        self._lock = RLock()
+
+    def register(self, name: str, workflow_cls: type[Workflow]) -> None:
+        with self._lock:
+            if name in self._workflows:
+                raise ValueError(f"workflow '{name}' already registered")
+            self._workflows[name] = workflow_cls
+
+    def get(self, name: str) -> Optional[type[Workflow]]:
+        with self._lock:
+            return self._workflows.get(name)
+
+    def names(self) -> list[str]:
+        with self._lock:
+            return sorted(self._workflows.keys())
+
+    def reset(self) -> None:
+        with self._lock:
+            self._workflows.clear()
+
+
+workflow_registry = WorkflowRegistry()
+
+
+def workflow(cls: type[TWorkflow]) -> type[TWorkflow]:
+    """Decorator that registers workflow classes and caches their IR."""
+
+    if not issubclass(cls, Workflow):
+        raise TypeError("workflow decorator requires Workflow subclasses")
+    run_impl = cls.run
+    if not inspect.iscoroutinefunction(run_impl):
+        raise TypeError("workflow run() must be defined with 'async def'")
+
+    @wraps(run_impl)
+    async def run_public(self: Workflow, *args: Any, **kwargs: Any) -> Any:
+        if _running_under_pytest():
+            cls.workflow_ir()
+            return await run_impl(self, *args, **kwargs)
+
+        initial_context = cls._build_initial_context(args, kwargs)
+
+        payload = cls._build_registration_payload(initial_context)
+        run_result = await bridge.run_instance(payload.SerializeToString())
+        cls._workflow_version_id = run_result.workflow_version_id
+        if _skip_wait_for_instance():
+            logger.info(
+                "Skipping wait_for_instance for workflow %s due to RAPPEL_SKIP_WAIT_FOR_INSTANCE",
+                cls.short_name(),
+            )
+            return None
+        result_bytes = await bridge.wait_for_instance(
+            instance_id=run_result.workflow_instance_id,
+            poll_interval_secs=1.0,
+        )
+        if result_bytes is None:
+            raise TimeoutError(
+                f"workflow instance {run_result.workflow_instance_id} did not complete"
+            )
+        arguments = pb2.WorkflowArguments()
+        arguments.ParseFromString(result_bytes)
+        result = deserialize_result_payload(arguments)
+        if result.error:
+            raise RuntimeError(f"workflow failed: {result.error}")
+
+        # Unwrap WorkflowNodeResult if present (internal worker representation)
+        if isinstance(result.result, WorkflowNodeResult):
+            # Extract the actual result from the variables dict
+            variables = result.result.variables
+            program = cls.workflow_ir()
+            # Get the return variable from the IR if available
+            if program.functions:
+                outputs = list(program.functions[0].io.outputs)
+                if outputs:
+                    return_var = outputs[0]
+                    if return_var in variables:
+                        return variables[return_var]
+            return None
+
+        return result.result
+
+    cls.__workflow_run_impl__ = run_impl
+    cls.run = run_public  # type: ignore[assignment]
+    workflow_registry.register(cls.short_name(), cls)
+    return cls
+
+
+def _running_under_pytest() -> bool:
+    return bool(os.environ.get("PYTEST_CURRENT_TEST"))
+
+
+def _skip_wait_for_instance() -> bool:
+    value = os.environ.get("RAPPEL_SKIP_WAIT_FOR_INSTANCE")
+    if not value:
+        return False
+    return value.strip().lower() not in {"0", "false", "no"}

rappel/workflow_runtime.py

@@ -0,0 +1,287 @@
+"""Runtime helpers for executing actions inside the worker.
+
+This module provides the execution layer for Python workers that receive
+action dispatch commands from the Rust scheduler.
+"""
+
+import asyncio
+import dataclasses
+from base64 import b64decode
+from dataclasses import dataclass
+from datetime import date, datetime, time, timedelta
+from decimal import Decimal
+from pathlib import Path, PurePath
+from typing import Any, Dict, get_args, get_origin, get_type_hints
+from uuid import UUID
+
+from pydantic import BaseModel
+
+from proto import messages_pb2 as pb2
+
+from .dependencies import provide_dependencies
+from .registry import registry
+from .serialization import arguments_to_kwargs
+
+
+class WorkflowNodeResult(BaseModel):
+    """Result from a workflow node execution containing variable bindings."""
+
+    variables: Dict[str, Any]
+
+
+@dataclass
+class ActionExecutionResult:
+    """Result of an action execution."""
+
+    result: Any
+    exception: BaseException | None = None
+
+
+def _is_pydantic_model(cls: type) -> bool:
+    """Check if a class is a Pydantic BaseModel subclass."""
+    try:
+        return isinstance(cls, type) and issubclass(cls, BaseModel)
+    except TypeError:
+        return False
+
+
+def _is_dataclass_type(cls: type) -> bool:
+    """Check if a class is a dataclass."""
+    return dataclasses.is_dataclass(cls) and isinstance(cls, type)
+
+
+def _coerce_primitive(value: Any, target_type: type) -> Any:
+    """Coerce a value to a primitive type based on target_type.
+
+    Handles conversion of serialized values (strings, floats) back to their
+    native Python types (UUID, datetime, etc.).
+    """
+    # Handle None
+    if value is None:
+        return None
+
+    # UUID from string
+    if target_type is UUID:
+        if isinstance(value, UUID):
+            return value
+        if isinstance(value, str):
+            return UUID(value)
+        return value
+
+    # datetime from ISO string
+    if target_type is datetime:
+        if isinstance(value, datetime):
+            return value
+        if isinstance(value, str):
+            return datetime.fromisoformat(value)
+        return value
+
+    # date from ISO string
+    if target_type is date:
+        if isinstance(value, date):
+            return value
+        if isinstance(value, str):
+            return date.fromisoformat(value)
+        return value
+
+    # time from ISO string
+    if target_type is time:
+        if isinstance(value, time):
+            return value
+        if isinstance(value, str):
+            return time.fromisoformat(value)
+        return value
+
+    # timedelta from total seconds
+    if target_type is timedelta:
+        if isinstance(value, timedelta):
+            return value
+        if isinstance(value, (int, float)):
+            return timedelta(seconds=value)
+        return value
+
+    # Decimal from string
+    if target_type is Decimal:
+        if isinstance(value, Decimal):
+            return value
+        if isinstance(value, (str, int, float)):
+            return Decimal(str(value))
+        return value
+
+    # bytes from base64 string
+    if target_type is bytes:
+        if isinstance(value, bytes):
+            return value
+        if isinstance(value, str):
+            return b64decode(value)
+        return value
+
+    # Path from string
+    if target_type is Path or target_type is PurePath:
+        if isinstance(value, PurePath):
+            return value
+        if isinstance(value, str):
+            return Path(value)
+        return value
+
+    return value
+
+
+# Types that can be coerced from serialized form
+COERCIBLE_TYPES = (UUID, datetime, date, time, timedelta, Decimal, bytes, Path, PurePath)
+
+
+def _coerce_dict_to_model(value: Any, target_type: type) -> Any:
+    """Convert a dict to a Pydantic model or dataclass if needed.
+
+    If value is a dict and target_type is a Pydantic model or dataclass,
+    instantiate the model with the dict values. Otherwise, return value unchanged.
+    """
+    if not isinstance(value, dict):
+        return value
+
+    if _is_pydantic_model(target_type):
+        # Use model_validate for Pydantic v2, fall back to direct instantiation
+        model_validate = getattr(target_type, "model_validate", None)
+        if model_validate is not None:
+            return model_validate(value)
+        return target_type(**value)
+
+    if _is_dataclass_type(target_type):
+        return target_type(**value)
+
+    return value
+
+
+def _coerce_value(value: Any, target_type: type) -> Any:
+    """Coerce a value to the target type.
+
+    Handles:
+    - Primitive types (UUID, datetime, etc.)
+    - Pydantic models and dataclasses (from dicts)
+    - Generic collections like list[UUID], set[datetime]
+    """
+    # Handle None
+    if value is None:
+        return None
+
+    # Check for coercible primitive types
+    if isinstance(target_type, type) and issubclass(target_type, COERCIBLE_TYPES):
+        return _coerce_primitive(value, target_type)
+
+    # Check for Pydantic models or dataclasses
+    if isinstance(value, dict):
+        coerced = _coerce_dict_to_model(value, target_type)
+        if coerced is not value:
+            return coerced
+
+    # Handle generic types like list[UUID], set[datetime]
+    origin = get_origin(target_type)
+    if origin is not None:
+        args = get_args(target_type)
+
+        # Handle list[T]
+        if origin is list and isinstance(value, list) and args:
+            item_type = args[0]
+            return [_coerce_value(item, item_type) for item in value]
+
+        # Handle set[T] (serialized as list)
+        if origin is set and isinstance(value, list) and args:
+            item_type = args[0]
+            return {_coerce_value(item, item_type) for item in value}
+
+        # Handle frozenset[T] (serialized as list)
+        if origin is frozenset and isinstance(value, list) and args:
+            item_type = args[0]
+            return frozenset(_coerce_value(item, item_type) for item in value)
+
+        # Handle tuple[T, ...] (serialized as list)
+        if origin is tuple and isinstance(value, (list, tuple)) and args:
+            # Variable length tuple like tuple[int, ...]
+            if len(args) == 2 and args[1] is ...:
+                item_type = args[0]
+                return tuple(_coerce_value(item, item_type) for item in value)
+            # Fixed length tuple like tuple[int, str, UUID]
+            return tuple(
+                _coerce_value(item, item_type) for item, item_type in zip(value, args, strict=False)
+            )
+
+        # Handle dict[K, V]
+        if origin is dict and isinstance(value, dict) and len(args) == 2:
+            key_type, val_type = args
+            return {
+                _coerce_value(k, key_type): _coerce_value(v, val_type) for k, v in value.items()
+            }
+
+    return value
+
+
+def _coerce_kwargs_to_type_hints(handler: Any, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """Coerce kwargs to expected types based on handler's type hints.
+
+    Handles:
+    - Pydantic models and dataclasses (from dicts)
+    - Primitive types like UUID, datetime, Decimal, etc.
+    - Generic collections like list[UUID], dict[str, datetime]
+    """
+    try:
+        type_hints = get_type_hints(handler)
+    except Exception:
+        # If we can't get type hints (e.g., forward references), return as-is
+        return kwargs
+
+    coerced = {}
+    for key, value in kwargs.items():
+        if key in type_hints:
+            target_type = type_hints[key]
+            coerced[key] = _coerce_value(value, target_type)
+        else:
+            coerced[key] = value
+
+    return coerced
+
+
+async def execute_action(dispatch: pb2.ActionDispatch) -> ActionExecutionResult:
+    """Execute an action based on the dispatch command.
+
+    Args:
+        dispatch: The action dispatch command from the Rust scheduler.
+
+    Returns:
+        The result of executing the action.
+    """
+    action_name = dispatch.action_name
+    module_name = dispatch.module_name
+
+    # Import the module if specified (this registers actions via @action decorator)
+    if module_name:
+        import importlib
+
+        importlib.import_module(module_name)
+
+    # Get the action handler using both module and name
+    handler = registry.get(module_name, action_name)
+    if handler is None:
+        return ActionExecutionResult(
+            result=None,
+            exception=KeyError(f"action '{module_name}:{action_name}' not registered"),
+        )
+
+    # Deserialize kwargs
+    kwargs = arguments_to_kwargs(dispatch.kwargs)
+
+    # Coerce dict arguments to Pydantic models or dataclasses based on type hints
+    # This is needed because the IR converts model constructor calls to dicts
+    kwargs = _coerce_kwargs_to_type_hints(handler, kwargs)
+
+    try:
+        async with provide_dependencies(handler, kwargs) as call_kwargs:
+            value = handler(**call_kwargs)
+            if asyncio.iscoroutine(value):
+                value = await value
+        return ActionExecutionResult(result=value)
+    except Exception as e:
+        return ActionExecutionResult(
+            result=None,
+            exception=e,
+        )