rappel 0.10.0__py3-none-win_amd64.whl

rappel/workflow.py

@@ -0,0 +1,338 @@
+"""
+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 types import UnionType
+from typing import (
+    Any,
+    Awaitable,
+    ClassVar,
+    Optional,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+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, _coerce_value
+
+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, _blocking: bool = True, _priority: Optional[int] = None, **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,
+        priority: Optional[int] = 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)
+
+        if priority is not None:
+            message.priority = priority
+
+        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,
+        _blocking: bool = True,
+        _priority: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Any:
+        if _running_under_pytest():
+            logger.debug(
+                "pytest run: workflow=%s in_memory=%s",
+                cls.short_name(),
+                os.environ.get("RAPPEL_BRIDGE_IN_MEMORY"),
+            )
+            _enable_in_memory_broker()
+            initial_context = cls._build_initial_context(args, kwargs)
+            payload = cls._build_registration_payload(initial_context, priority=_priority)
+            result_bytes = await bridge.execute_workflow(payload.SerializeToString())
+            return _deserialize_workflow_result(cls, result_bytes)
+
+        initial_context = cls._build_initial_context(args, kwargs)
+
+        payload = cls._build_registration_payload(initial_context, priority=_priority)
+        run_result = await bridge.run_instance(payload.SerializeToString())
+        cls._workflow_version_id = run_result.workflow_version_id
+
+        if not _blocking:
+            return run_result.workflow_instance_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"
+            )
+        return _deserialize_workflow_result(cls, result_bytes)
+
+    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"}
+
+
+def _enable_in_memory_broker() -> None:
+    os.environ.setdefault("RAPPEL_BRIDGE_IN_MEMORY", "1")
+
+
+def _deserialize_workflow_result(
+    workflow_cls: type[Workflow],
+    result_bytes: bytes,
+) -> Any:
+    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 = workflow_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
+
+    value = result.result
+    target_type = _resolve_return_type(workflow_cls)
+    if target_type is None:
+        return value
+    return _coerce_result_value(value, target_type)
+
+
+def _resolve_return_type(workflow_cls: type[Workflow]) -> Optional[type]:
+    try:
+        run_impl = workflow_cls.__workflow_run_impl__  # type: ignore[attr-defined]
+    except AttributeError:
+        run_impl = workflow_cls.run
+    try:
+        type_hints = get_type_hints(run_impl)
+    except Exception:
+        return None
+    target_type = type_hints.get("return")
+    if target_type is None or target_type is Any:
+        return None
+    return target_type
+
+
+def _coerce_result_value(value: Any, target_type: type) -> Any:
+    origin = get_origin(target_type)
+    if origin is UnionType or origin is Union:
+        for arg in get_args(target_type):
+            if arg is type(None) and value is None:
+                return None
+            try:
+                coerced = _coerce_value(value, arg)
+            except Exception:
+                continue
+            if coerced is not value:
+                return coerced
+            if isinstance(arg, type) and isinstance(value, arg):
+                return value
+        return value
+    try:
+        return _coerce_value(value, target_type)
+    except Exception:
+        return value

rappel/workflow_runtime.py

@@ -0,0 +1,292 @@
+"""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
+
+    module = None
+    if module_name:
+        import importlib
+
+        module = 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 and module is not None:
+        import importlib
+
+        module = importlib.reload(module)
+        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,
+        )