diff --git a/.changeset/neat-chicken-bake.md b/.changeset/neat-chicken-bake.md new file mode 100644 index 00000000..13810701 --- /dev/null +++ b/.changeset/neat-chicken-bake.md @@ -0,0 +1,5 @@ +--- +"llama-index-workflows": minor +--- + +Make WorkflowTick serializable, and support switching workflow name and runtime before launch diff --git a/packages/llama-index-workflows/src/workflows/context/state_store.py b/packages/llama-index-workflows/src/workflows/context/state_store.py index 45af8473..fbc23a1c 100644 --- a/packages/llama-index-workflows/src/workflows/context/state_store.py +++ b/packages/llama-index-workflows/src/workflows/context/state_store.py @@ -192,6 +192,121 @@ def assign_path_step(obj: Any, segment: str, value: Any) -> None: setattr(obj, segment, value) +def get_by_path(state: Any, path: str, default: Any = Ellipsis) -> Any: + """Get a nested value from state using a dot-separated path. + + Args: + state: The root state object. + path: Dot-separated path, e.g. "user.profile.name". + default: If provided, return this when the path does not exist; + otherwise, raise ValueError. + + Returns: + The resolved value. + + Raises: + ValueError: If the path is invalid and no default is provided, + or if path depth exceeds MAX_DEPTH. + """ + segments = path.split(".") if path else [] + if len(segments) > MAX_DEPTH: + raise ValueError(f"Path length exceeds {MAX_DEPTH} segments") + + try: + value: Any = state + for segment in segments: + value = traverse_path_step(value, segment) + except Exception: + if default is not Ellipsis: + return default + raise ValueError(f"Path '{path}' not found in state") + return value + + +def set_by_path(state: Any, path: str, value: Any) -> None: + """Set a nested value on state using a dot-separated path. + + Intermediate dicts are created as needed. + + Args: + state: The root state object (mutated in place). + path: Dot-separated path to write. + value: Value to assign. + + Raises: + ValueError: If the path is empty or exceeds MAX_DEPTH. + """ + if not path: + raise ValueError("Path cannot be empty") + + segments = path.split(".") + if len(segments) > MAX_DEPTH: + raise ValueError(f"Path length exceeds {MAX_DEPTH} segments") + + current = state + for segment in segments[:-1]: + try: + current = traverse_path_step(current, segment) + except (KeyError, AttributeError, IndexError, TypeError): + intermediate: Any = {} + assign_path_step(current, segment, intermediate) + current = intermediate + + assign_path_step(current, segments[-1], value) + + +def merge_state(current_state: MODEL_T, incoming: BaseModel) -> MODEL_T: + """Replace or merge incoming state onto current state. + + If incoming is the same type (or subclass) of current, it replaces directly. + If current's type is a subclass of incoming's type (parent provided), + fields are merged preserving child-specific fields. + + Args: + current_state: The existing state. + incoming: The new state to apply. + + Returns: + The resulting state after merge/replace. + + Raises: + ValueError: If the types are not compatible. + """ + current_type = type(current_state) + new_type = type(incoming) + + if isinstance(incoming, current_type): + return incoming # type: ignore[return-value] + elif issubclass(current_type, new_type): + parent_data = incoming.model_dump() + return current_type.model_validate( + {**current_state.model_dump(), **parent_data} + ) + else: + raise ValueError( + f"State must be of type {current_type.__name__} or a parent type, " + f"got {new_type.__name__}" + ) + + +def create_cleared_state(state_type: Type[MODEL_T]) -> MODEL_T: + """Create a default instance of the state type, wrapping ValidationError. + + Args: + state_type: The state model class to instantiate. + + Returns: + A new default instance. + + Raises: + ValueError: If the model cannot be instantiated from defaults. + """ + try: + return state_type() + except ValidationError: + raise ValueError("State must have defaults for all fields") + + # Only warn once about unserializable keys class UnserializableKeyWarning(Warning): pass @@ -377,27 +492,8 @@ class InMemoryStateStore(Generic[MODEL_T]): Raises: ValueError: If the types are not compatible (neither same nor parent). """ - current_type = type(self._state) - new_type = type(state) - - if isinstance(state, current_type): - # Exact match or subclass - direct replacement - async with self._lock: - self._state = state - elif issubclass(current_type, new_type): - # Parent type provided - merge fields onto current state - # This preserves child-specific fields while updating parent fields - async with self._lock: - # Get the fields from the parent type and update them on the current state - parent_data = state.model_dump() - self._state = current_type.model_validate( - {**self._state.model_dump(), **parent_data} - ) - else: - raise ValueError( - f"State must be of type {current_type.__name__} or a parent type, " - f"got {new_type.__name__}" - ) + async with self._lock: + self._state = merge_state(self._state, state) def to_dict(self, serializer: "BaseSerializer") -> dict[str, Any]: """Serialize the state and model metadata for persistence. @@ -465,9 +561,6 @@ class InMemoryStateStore(Generic[MODEL_T]): async def get(self, path: str, default: Any = Ellipsis) -> Any: """Get a nested value using dot-separated paths. - Supports dict keys, list indices, and attribute access transparently at - each segment. - Args: path (str): Dot-separated path, e.g. "user.profile.name". default (Any): If provided, return this when the path does not @@ -480,30 +573,12 @@ class InMemoryStateStore(Generic[MODEL_T]): ValueError: If the path is invalid and no default is provided or if the path depth exceeds limits. """ - segments = path.split(".") if path else [] - if len(segments) > MAX_DEPTH: - raise ValueError(f"Path length exceeds {MAX_DEPTH} segments") - async with self._lock: - try: - value: Any = self._state - for segment in segments: - value = traverse_path_step(value, segment) - except Exception: - if default is not Ellipsis: - return default - - msg = f"Path '{path}' not found in state" - raise ValueError(msg) - - return value + return get_by_path(self._state, path, default) async def set(self, path: str, value: Any) -> None: """Set a nested value using dot-separated paths. - Intermediate containers are created as needed. Dicts, lists, tuples, and - Pydantic models are supported where appropriate. - Args: path (str): Dot-separated path to write. value (Any): Value to assign. @@ -511,28 +586,8 @@ class InMemoryStateStore(Generic[MODEL_T]): Raises: ValueError: If the path is empty or exceeds the maximum depth. """ - if not path: - raise ValueError("Path cannot be empty") - - segments = path.split(".") - if len(segments) > MAX_DEPTH: - raise ValueError(f"Path length exceeds {MAX_DEPTH} segments") - async with self._lock: - current = self._state - - # Navigate/create intermediate segments - for segment in segments[:-1]: - try: - current = traverse_path_step(current, segment) - except (KeyError, AttributeError, IndexError, TypeError): - # Create intermediate object and assign it - intermediate: Any = {} - assign_path_step(current, segment, intermediate) - current = intermediate - - # Assign the final value - assign_path_step(current, segments[-1], value) + set_by_path(self._state, path, value) async def clear(self) -> None: """Reset the state to its type defaults. @@ -541,10 +596,33 @@ class InMemoryStateStore(Generic[MODEL_T]): ValueError: If the model type cannot be instantiated from defaults (i.e., fields missing default values). """ + await self.set_state(create_cleared_state(self._state.__class__)) + + +def deserialize_dict_state_data( + data: dict[str, Any], + serializer: BaseSerializer, +) -> DictState: + """Deserialize DictState from {"_data": {...}} format. + + Args: + data: Dict with {"_data": {...}} structure containing serialized values. + serializer: Strategy for decoding values. + + Returns: + DictState with deserialized values. + + Raises: + ValueError: If deserialization fails for any key. + """ + _data_serialized = data.get("_data", {}) + deserialized_data = {} + for key, value in _data_serialized.items(): try: - await self.set_state(self._state.__class__()) - except ValidationError: - raise ValueError("State must have defaults for all fields") + deserialized_data[key] = serializer.deserialize(value) + except Exception as e: + raise ValueError(f"Failed to deserialize state value for key {key}: {e}") + return DictState(_data=deserialized_data) def deserialize_state_from_dict( @@ -570,16 +648,7 @@ def deserialize_state_from_dict( state_type_name = serialized_state.get("state_type", "DictState") if state_type_name == "DictState": - _data_serialized = state_data.get("_data", {}) - deserialized_data = {} - for key, value in _data_serialized.items(): - try: - deserialized_data[key] = serializer.deserialize(value) - except Exception as e: - raise ValueError( - f"Failed to deserialize state value for key {key}: {e}" - ) - return DictState(_data=deserialized_data) + return deserialize_dict_state_data(state_data, serializer) else: return serializer.deserialize(state_data) diff --git a/packages/llama-index-workflows/src/workflows/plugins/basic.py b/packages/llama-index-workflows/src/workflows/plugins/basic.py index f416c95d..5bdcab2d 100644 --- a/packages/llama-index-workflows/src/workflows/plugins/basic.py +++ b/packages/llama-index-workflows/src/workflows/plugins/basic.py @@ -127,7 +127,7 @@ class InternalAsyncioAdapter(InternalRunAdapter, SnapshottableAdapter): except asyncio.TimeoutError: return WaitResultTimeout() - def on_tick(self, tick: WorkflowTick) -> None: + async def on_tick(self, tick: WorkflowTick) -> None: self._queues.ticks.append(tick) def replay(self) -> list[WorkflowTick]: @@ -147,7 +147,8 @@ class ExternalAsyncioAdapter( and stream events published by the workflow. """ - def __init__(self, queues: AsyncioAdapterQueues) -> None: + def __init__(self, outer: BasicRuntime, queues: AsyncioAdapterQueues) -> None: + self._outer = outer self._queues = queues @property @@ -170,9 +171,6 @@ class ExternalAsyncioAdapter( if isinstance(item, StopEvent): break - def on_tick(self, tick: WorkflowTick) -> None: - self._queues.ticks.append(tick) - def replay(self) -> list[WorkflowTick]: return self._queues.ticks @@ -195,6 +193,7 @@ class ExternalAsyncioAdapter( """Abort by cancelling the control loop task.""" if not self._queues.complete.done(): self._queues.complete.cancel() + self._outer._queues.pop(self.run_id, None) @property def init_state(self) -> BrokerState: @@ -205,6 +204,7 @@ class BasicRuntime(Runtime): """Default asyncio-based runtime with no durability.""" def __init__(self) -> None: + super().__init__() # WeakValueDictionary allows queues to be GC'd when no adapters reference them. # The task closure in run_workflow() captures a strong reference, keeping # queues alive for fire-and-forget workflows even if the external adapter is dropped. @@ -315,7 +315,7 @@ class BasicRuntime(Runtime): def get_external_adapter(self, run_id: str) -> ExternalRunAdapter: if run_id not in self._queues: raise RuntimeError(f"No active workflow with run_id '{run_id}'. ") - return ExternalAsyncioAdapter(self._queues[run_id]) + return ExternalAsyncioAdapter(self, self._queues[run_id]) _current_run_id: ContextVar[str | None] = ContextVar("current_run_id", default=None) diff --git a/packages/llama-index-workflows/src/workflows/runtime/control_loop.py b/packages/llama-index-workflows/src/workflows/runtime/control_loop.py index b7a957d4..3ada4dbb 100644 --- a/packages/llama-index-workflows/src/workflows/runtime/control_loop.py +++ b/packages/llama-index-workflows/src/workflows/runtime/control_loop.py @@ -11,7 +11,6 @@ import traceback from dataclasses import replace from typing import TYPE_CHECKING -from workflows.decorators import R from workflows.errors import ( WorkflowCancelledByUser, WorkflowRuntimeError, @@ -50,7 +49,6 @@ from workflows.runtime.types.named_task import NamedTask from workflows.runtime.types.plugin import ( InternalRunAdapter, WaitResultTick, - as_snapshottable_adapter, get_current_run, ) from workflows.runtime.types.results import ( @@ -128,7 +126,6 @@ class _ControlLoopRunner: # avoiding TypeError from comparing WorkflowTick objects that don't implement __lt__ self.scheduled_wakeups: list[tuple[float, int, WorkflowTick]] = [] self._wakeup_sequence = 0 - self.snapshot_adapter = as_snapshottable_adapter(adapter) # Pull task sequence counter for deterministic journaling self._pull_sequence = 0 # Map from worker task to (step_name, worker_id) key @@ -431,8 +428,7 @@ class _ControlLoopRunner: ) raise - if self.snapshot_adapter is not None: - self.snapshot_adapter.on_tick(tick) + await self.adapter.on_tick(tick) for command in commands: try: @@ -557,7 +553,7 @@ def _check_idle_state(state: BrokerState) -> bool: def _process_step_result_tick( - tick: TickStepResult[R], init: BrokerState, now_seconds: float + tick: TickStepResult, init: BrokerState, now_seconds: float ) -> tuple[BrokerState, list[WorkflowCommand]]: """ processes the results from a step function execution diff --git a/packages/llama-index-workflows/src/workflows/runtime/types/plugin.py b/packages/llama-index-workflows/src/workflows/runtime/types/plugin.py index debfc8f8..f940deb0 100644 --- a/packages/llama-index-workflows/src/workflows/runtime/types/plugin.py +++ b/packages/llama-index-workflows/src/workflows/runtime/types/plugin.py @@ -7,6 +7,7 @@ A runtime interface to switch out a broker runtime (external library or service from __future__ import annotations import asyncio +import weakref from abc import ABC, abstractmethod from contextlib import contextmanager from contextvars import ContextVar, Token @@ -63,7 +64,7 @@ WaitResult = Union[WaitResultTick, WaitResultTimeout] class RegisteredWorkflow: workflow: Workflow workflow_run_fn: WorkflowRunFunction - steps: dict[str, StepWorkerFunction[Any]] + steps: dict[str, StepWorkerFunction] class InternalRunAdapter(ABC): @@ -180,6 +181,17 @@ class InternalRunAdapter(ABC): """ pass + async def on_tick(self, tick: WorkflowTick) -> None: + """ + Called whenever a tick event is processed by the control loop. + + This method is invoked for both external ticks (sent via send_event) + and internal ticks (generated by step completions, timeouts, etc.). + Adapters can override to record ticks, persist them, etc. + Default is no-op. + """ + pass + async def wait_for_next_task( self, task_set: list[NamedTask], @@ -310,7 +322,7 @@ class RunContext: workflow: Workflow run_adapter: InternalRunAdapter context: Context - steps: dict[str, StepWorkerFunction[Any]] + steps: dict[str, StepWorkerFunction] _current_run: ContextVar[RunContext | None] = ContextVar("current_run", default=None) @@ -334,6 +346,48 @@ def get_current_run() -> RunContext: return ctx +class WorkflowSet: + """Identity-based weak set for tracking Workflow instances. + + Uses id() as the key and weakref.ref with cleanup callbacks to + avoid hashability requirements and memory leaks. + """ + + def __init__(self) -> None: + self._refs: dict[int, weakref.ref[Workflow]] = {} + + def add(self, workflow: Workflow) -> None: + obj_id = id(workflow) + if obj_id in self._refs: + return + + def _cleanup(ref: weakref.ref[Workflow], _id: int = obj_id) -> None: + self._refs.pop(_id, None) + + self._refs[obj_id] = weakref.ref(workflow, _cleanup) + + def discard(self, workflow: Workflow) -> None: + self._refs.pop(id(workflow), None) + + def __contains__(self, workflow: Workflow) -> bool: + ref = self._refs.get(id(workflow)) + if ref is None: + return False + return ref() is not None + + def __iter__(self) -> Generator[Workflow, None, None]: + for ref in list(self._refs.values()): + obj = ref() + if obj is not None: + yield obj + + def __len__(self) -> int: + return sum(1 for _ in self) + + def __bool__(self) -> bool: + return any(ref() is not None for ref in self._refs.values()) + + class Runtime(ABC): """ Abstract base class for workflow execution runtimes. @@ -351,9 +405,13 @@ class Runtime(ABC): Use registering() context manager for implicit workflow registration. """ + def __init__(self) -> None: + self._pending: WorkflowSet = WorkflowSet() + self._launched: bool = False + _token: Token[Runtime | None] - def get_or_register(self, workflow: "Workflow") -> RegisteredWorkflow: + def get_or_register(self, workflow: Workflow) -> RegisteredWorkflow: """Get the registered workflow if available, otherwise register it.""" registered = self.get_registered(workflow) if registered is None: @@ -361,7 +419,7 @@ class Runtime(ABC): return registered @abstractmethod - def register(self, workflow: "Workflow") -> RegisteredWorkflow: + def register(self, workflow: Workflow) -> RegisteredWorkflow: """ Register a workflow with the runtime. @@ -399,7 +457,7 @@ class Runtime(ABC): ... @abstractmethod - def get_internal_adapter(self, workflow: "Workflow") -> InternalRunAdapter: + def get_internal_adapter(self, workflow: Workflow) -> InternalRunAdapter: """ Get the internal adapter for a workflow run. @@ -429,7 +487,11 @@ class Runtime(ABC): For many runtime's, this must be called before running workflows. """ - pass + self._launched = True + for wf in self._pending: + self.register(wf) + wf._runtime_locked = True + self._pending = WorkflowSet() def destroy(self) -> None: """ @@ -439,17 +501,19 @@ class Runtime(ABC): """ pass - def track_workflow(self, workflow: "Workflow") -> None: + def track_workflow(self, workflow: Workflow) -> None: """ Track a workflow instance for registration at launch time. Called by Workflow.__init__ to register with the runtime. - Override in runtimes that need to track workflows. - Default implementation is a no-op. """ - pass + self._pending.add(workflow) - def get_registered(self, workflow: "Workflow") -> RegisteredWorkflow | None: + def untrack_workflow(self, workflow: Workflow) -> None: + """Remove a workflow from this runtime's tracking set.""" + self._pending.discard(workflow) + + def get_registered(self, workflow: Workflow) -> RegisteredWorkflow | None: """ Get the registered workflow if available. @@ -462,8 +526,7 @@ class Runtime(ABC): """ Context manager for implicit workflow registration. - Workflows created inside this block will automatically register - with this runtime. Does NOT call launch() on exit. + Workflows created inside this block will automatically set this runtime as their runtime. """ token = _current_runtime.set(self) try: @@ -478,35 +541,19 @@ class SnapshottableAdapter(ABC): This is a standalone mixin (not inheriting from InternalRunAdapter or ExternalRunAdapter) that can be combined with adapter implementations - to add tick recording for debugging or replay purposes. - - Adapters that implement this interface can record ticks as they occur - and replay them later. `on_tick` is called whenever a tick event is - received externally OR as a result from an internal command (e.g., a - step function completing, a timeout occurring, etc.) + to add init_state and replay capabilities for state reconstruction. Use `as_snapshottable_adapter()` to check if an adapter supports snapshotting. """ @property @abstractmethod - def init_state(self) -> "BrokerState": + def init_state(self) -> BrokerState: """ Get the initial state of the adapter. """ ... - @abstractmethod - def on_tick(self, tick: WorkflowTick) -> None: - """ - Called whenever a tick event is received. - - This method is invoked for both external ticks (sent via send_event) - and internal ticks (generated by step completions, timeouts, etc.). - Implementations should record the tick for later replay. - """ - ... - @abstractmethod def replay(self) -> list[WorkflowTick]: """ diff --git a/packages/llama-index-workflows/src/workflows/runtime/types/results.py b/packages/llama-index-workflows/src/workflows/runtime/types/results.py index 4f388a56..f09c653f 100644 --- a/packages/llama-index-workflows/src/workflows/runtime/types/results.py +++ b/packages/llama-index-workflows/src/workflows/runtime/types/results.py @@ -7,19 +7,21 @@ import dataclasses from contextvars import ContextVar from dataclasses import dataclass from typing import ( - TYPE_CHECKING, Any, Generic, + Literal, TypeVar, Union, ) -from workflows.decorators import R +from pydantic import BaseModel, ConfigDict, model_serializer, model_validator from workflows.events import Event - -if TYPE_CHECKING: - pass - +from workflows.runtime.types.serialization_helpers import ( + SerializableEvent, + SerializableEventType, + SerializableException, + SerializableOptionalEvent, +) EventType = TypeVar("EventType", bound=Event) @@ -29,7 +31,7 @@ EventType = TypeVar("EventType", bound=Event) @dataclass(frozen=True) -class StepWorkerContext(Generic[R]): +class StepWorkerContext: """ Base state passed to step functions and returned by step functions. """ @@ -37,7 +39,7 @@ class StepWorkerContext(Generic[R]): # immutable state of the step events at start of the step function execution state: StepWorkerState # add commands here to mutate the internal worker state after step execution - returns: Returns[R] + returns: Returns @dataclass(frozen=True) @@ -79,13 +81,13 @@ class StepWorkerWaiter(Generic[EventType]): @dataclass() -class Returns(Generic[R]): +class Returns: """ Mutate to add return values to the step function. These are only executed after the step function has completed (including errors!) """ - return_values: list[StepFunctionResult[R]] + return_values: list[StepFunctionResult] class WaitingForEvent(Exception, Generic[EventType]): @@ -110,71 +112,83 @@ StepWorkerStateContextVar = ContextVar[StepWorkerContext]("step_worker") ################################### -@dataclass -class StepWorkerResult(Generic[R]): - """ - Returned after a step function has been successfully executed. - """ +class StepWorkerResult(BaseModel): + """Returned after a step function has been successfully executed.""" - result: R + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["result"] = "result" + result: SerializableOptionalEvent = None -@dataclass -class StepWorkerFailed(Generic[R]): - """ - Returned after a step function has failed - """ +class StepWorkerFailed(BaseModel): + """Returned after a step function has failed.""" - exception: Exception + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["failed"] = "failed" + exception: SerializableException failed_at: float -@dataclass -class DeleteWaiter: - """ - Returned after a waiter condition has been successfully resolved. - """ +class DeleteWaiter(BaseModel): + """Returned after a waiter condition has been successfully resolved.""" + model_config = ConfigDict(frozen=True) + type: Literal["delete_waiter"] = "delete_waiter" waiter_id: str -@dataclass -class DeleteCollectedEvent: - """ - Returned after a collected event has been successfully resolved. - """ +class DeleteCollectedEvent(BaseModel): + """Returned after a collected event has been successfully resolved.""" + model_config = ConfigDict(frozen=True) + type: Literal["delete_collected"] = "delete_collected" event_id: str -@dataclass -class AddCollectedEvent: - """ - Returned after a collected event has been added, and is not yet resolved. - """ +class AddCollectedEvent(BaseModel): + """Returned after a collected event has been added, and is not yet resolved.""" + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["add_collected"] = "add_collected" event_id: str - event: Event + event: SerializableEvent -@dataclass -class AddWaiter(Generic[EventType]): - """ - Returned after a waiter has been added, and is not yet resolved. - """ +class AddWaiter(BaseModel, Generic[EventType]): + """Returned after a waiter has been added, and is not yet resolved.""" + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["add_waiter"] = "add_waiter" waiter_id: str - waiter_event: Event | None - requirements: dict[str, Any] - timeout: float | None - event_type: type[EventType] + waiter_event: SerializableOptionalEvent = None + requirements: dict[str, Any] = {} + timeout: float | None = None + event_type: SerializableEventType + has_requirements: bool = False + + @model_serializer(mode="wrap") + def _serialize(self, handler: Any) -> dict[str, Any]: + data = handler(self) + # Always serialize requirements as {} and record whether they existed + data["has_requirements"] = bool(self.requirements) + data["requirements"] = {} + return data + + @model_validator(mode="wrap") + @classmethod + def _validate(cls, data: Any, handler: Any) -> AddWaiter: + if isinstance(data, dict): + # Strip has_requirements before validation (it's computed) + data = dict(data) + data.pop("has_requirements", None) + return handler(data) # A step function result "command" communicates back to the workflow how the step function was resolved # e.g. are we collecting events, waiting for an event, or just returning a result? StepFunctionResult = Union[ - StepWorkerResult[R], - StepWorkerFailed[R], + StepWorkerResult, + StepWorkerFailed, AddCollectedEvent, DeleteCollectedEvent, AddWaiter[Event], diff --git a/packages/llama-index-workflows/src/workflows/runtime/types/serialization_helpers.py b/packages/llama-index-workflows/src/workflows/runtime/types/serialization_helpers.py new file mode 100644 index 00000000..6597d53d --- /dev/null +++ b/packages/llama-index-workflows/src/workflows/runtime/types/serialization_helpers.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Annotated types for Pydantic serialization of tricky tick/result fields. + +Provides custom serializers/validators for: +- Events (polymorphic Pydantic models) +- Exceptions (not natively serializable) +- Event types (type[Event] as qualified name strings) +""" + +from __future__ import annotations + +from typing import Annotated, Any, Optional + +from pydantic import PlainSerializer, PlainValidator +from workflows.context.serializers import JsonSerializer +from workflows.context.utils import ( + import_module_from_qualified_name, +) +from workflows.events import Event + +_json_serializer = JsonSerializer() + + +def _serialize_event(event: Event) -> Any: + return _json_serializer.serialize_value(event) + + +def _deserialize_event(data: Any) -> Event: + return _json_serializer.deserialize_value(data) + + +SerializableEvent = Annotated[ + Event, + PlainSerializer(_serialize_event, return_type=Any), + PlainValidator(_deserialize_event), +] + + +def _serialize_optional_event(event: Event | None) -> Any: + if event is None: + return None + return _json_serializer.serialize_value(event) + + +def _deserialize_optional_event(data: Any) -> Event | None: + if data is None: + return None + return _json_serializer.deserialize_value(data) + + +SerializableOptionalEvent = Annotated[ + Optional[Event], + PlainSerializer(_serialize_optional_event, return_type=Any), + PlainValidator(_deserialize_optional_event), +] + + +def _serialize_exception(exc: Exception) -> dict[str, Any]: + exc_type = type(exc) + qualified_name = f"{exc_type.__module__}.{exc_type.__qualname__}" + return { + "exception_type": qualified_name, + "exception_message": str(exc), + } + + +def _deserialize_exception(data: Any) -> Exception: + if isinstance(data, Exception): + return data + exc_message = data["exception_message"] + try: + exc_cls = import_module_from_qualified_name(data["exception_type"]) + return exc_cls(exc_message) + except (ImportError, AttributeError, ValueError): + return Exception(exc_message) + + +SerializableException = Annotated[ + Exception, + PlainSerializer(_serialize_exception, return_type=dict[str, Any]), + PlainValidator(_deserialize_exception), +] + + +def _serialize_event_type(event_type: type[Event]) -> str: + return f"{event_type.__module__}.{event_type.__qualname__}" + + +def _deserialize_event_type(data: Any) -> type[Event]: + if isinstance(data, type): + return data + return import_module_from_qualified_name(data) + + +SerializableEventType = Annotated[ + type[Event], + PlainSerializer(_serialize_event_type, return_type=str), + PlainValidator(_deserialize_event_type), +] diff --git a/packages/llama-index-workflows/src/workflows/runtime/types/step_function.py b/packages/llama-index-workflows/src/workflows/runtime/types/step_function.py index 3a0a18a2..3443d2dc 100644 --- a/packages/llama-index-workflows/src/workflows/runtime/types/step_function.py +++ b/packages/llama-index-workflows/src/workflows/runtime/types/step_function.py @@ -7,10 +7,10 @@ import asyncio import functools import time from contextvars import copy_context -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Protocol +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Protocol, TypeVar from llama_index_instrumentation.dispatcher import instrument_tags -from workflows.decorators import P, R, StepConfig +from workflows.decorators import P, StepConfig from workflows.errors import WorkflowRuntimeError from workflows.events import ( Event, @@ -40,24 +40,26 @@ from workflows.workflow import Workflow if TYPE_CHECKING: from workflows.context.context import Context +StepReturnT = TypeVar("StepReturnT", bound=Optional[Event]) -class StepWorkerFunction(Protocol, Generic[R]): + +class StepWorkerFunction(Protocol): def __call__( self, state: StepWorkerState, step_name: str, event: Event, workflow: Workflow, - ) -> Awaitable[list[StepFunctionResult[R]]]: ... + ) -> Awaitable[list[StepFunctionResult]]: ... async def partial( - func: Callable[..., R], + func: Callable[..., Any], step_config: StepConfig, event: Event, context: Context, workflow: Workflow, -) -> Callable[[], R]: +) -> Callable[[], Any]: kwargs: dict[str, Any] = {} kwargs[step_config.event_name] = event if step_config.context_parameter: @@ -75,14 +77,16 @@ async def partial( def as_step_worker_functions(workflow: Workflow) -> dict[str, StepWorkerFunction]: step_funcs = workflow._get_steps() - step_workers: dict[str, StepWorkerFunction[Any]] = { + step_workers: dict[str, StepWorkerFunction] = { name: as_step_worker_function(getattr(func, "__func__", func)) for name, func in step_funcs.items() } return step_workers -def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFunction[R]: +def as_step_worker_function( + func: Callable[P, Awaitable[StepReturnT]], +) -> StepWorkerFunction: """ Wrap a step function, setting context variables and handling exceptions to instead return the appropriate StepFunctionResult. @@ -90,7 +94,7 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti # Keep original function reference for free-function steps; for methods we # will resolve the currently-bound method from the provided workflow at call time. - original_func: Callable[..., Awaitable[R]] = func + original_func: Callable[..., Awaitable[StepReturnT]] = func # Avoid functools.wraps here because it would set __wrapped__ to the bound # method (when present), which would strongly reference the workflow @@ -100,11 +104,11 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti step_name: str, event: Event, workflow: Workflow, - ) -> list[StepFunctionResult[R]]: + ) -> list[StepFunctionResult]: from workflows.context.context import Context internal_context = Context._create_internal(workflow=workflow) - returns = Returns[R](return_values=[]) + returns = Returns(return_values=[]) token = StepWorkerStateContextVar.set( StepWorkerContext(state=state, returns=returns) @@ -134,15 +138,17 @@ def as_step_worker_function(func: Callable[P, Awaitable[R]]) -> StepWorkerFuncti # run_in_executor doesn't accept **kwargs, so we need to use partial copy = copy_context() - result: R = await asyncio.get_event_loop().run_in_executor( - None, - lambda: copy.run(partial_func), # type: ignore + result: StepReturnT = ( + await asyncio.get_event_loop().run_in_executor( + None, + lambda: copy.run(partial_func), # type: ignore + ) ) else: result = await partial_func() - if result is not None and not isinstance(result, Event): - msg = f"Step function {step_name} returned {type(result).__name__} instead of an Event instance." - raise WorkflowRuntimeError(msg) + if result is not None and not isinstance(result, Event): + msg = f"Step function {step_name} returned {type(result).__name__} instead of an Event instance." + raise WorkflowRuntimeError(msg) returns.return_values.append(StepWorkerResult(result=result)) except WaitingForEvent as e: await asyncio.sleep(0) diff --git a/packages/llama-index-workflows/src/workflows/runtime/types/ticks.py b/packages/llama-index-workflows/src/workflows/runtime/types/ticks.py index ed197af2..7a34c8e6 100644 --- a/packages/llama-index-workflows/src/workflows/runtime/types/ticks.py +++ b/packages/llama-index-workflows/src/workflows/runtime/types/ticks.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: MIT # Copyright (c) 2026 LlamaIndex Inc. - """ Ticks (events) that drive the control loop. @@ -16,55 +15,61 @@ events that can occur during workflow execution: from __future__ import annotations -from dataclasses import dataclass -from typing import Generic, Union +from typing import Annotated, Literal, Union -from workflows.decorators import R -from workflows.events import Event +from pydantic import BaseModel, ConfigDict, Discriminator, TypeAdapter from workflows.runtime.types.results import StepFunctionResult +from workflows.runtime.types.serialization_helpers import SerializableEvent -@dataclass(frozen=True) -class TickStepResult(Generic[R]): +class TickStepResult(BaseModel): """When processed, executes a step function and publishes the result""" + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["step_result"] = "step_result" step_name: str worker_id: int - event: Event - result: list[StepFunctionResult[R]] + event: SerializableEvent + result: list[Annotated[StepFunctionResult, Discriminator("type")]] -@dataclass(frozen=True) -class TickAddEvent: +class TickAddEvent(BaseModel): """When sent, adds an event to the workflow's event queue""" - event: Event + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["add_event"] = "add_event" + event: SerializableEvent step_name: str | None = None attempts: int | None = None first_attempt_at: float | None = None -@dataclass(frozen=True) -class TickCancelRun: +class TickCancelRun(BaseModel): """When processed, cancels the workflow run""" - pass + model_config = ConfigDict(frozen=True) + type: Literal["cancel_run"] = "cancel_run" -@dataclass(frozen=True) -class TickPublishEvent: +class TickPublishEvent(BaseModel): """When sent, publishes an event to workflow consumers, e.g. a UI or a callback""" - event: Event + model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True) + type: Literal["publish_event"] = "publish_event" + event: SerializableEvent -@dataclass(frozen=True) -class TickTimeout: +class TickTimeout(BaseModel): """When processed, times the workflow out, cancelling it""" + model_config = ConfigDict(frozen=True) + type: Literal["timeout"] = "timeout" timeout: float -WorkflowTick = Union[ - TickStepResult[R], TickAddEvent, TickCancelRun, TickPublishEvent, TickTimeout +WorkflowTick = Annotated[ + Union[TickStepResult, TickAddEvent, TickCancelRun, TickPublishEvent, TickTimeout], + Discriminator("type"), ] + +WorkflowTickAdapter: TypeAdapter[WorkflowTick] = TypeAdapter(WorkflowTick) diff --git a/packages/llama-index-workflows/src/workflows/workflow.py b/packages/llama-index-workflows/src/workflows/workflow.py index dc381b2f..de7ff7dc 100644 --- a/packages/llama-index-workflows/src/workflows/workflow.py +++ b/packages/llama-index-workflows/src/workflows/workflow.py @@ -148,6 +148,7 @@ class Workflow(metaclass=WorkflowMeta): self._resource_manager = resource_manager or ResourceManager() # Instrumentation self._dispatcher = dispatcher + self._runtime_locked = False # Runtime registration: explicit > context-scoped > basic_runtime from workflows.plugins._context import get_current_runtime @@ -178,6 +179,18 @@ class Workflow(metaclass=WorkflowMeta): """The runtime this workflow is registered with.""" return self._runtime + def _switch_runtime(self, new_runtime: Runtime) -> None: + if new_runtime is self._runtime: + return + if self._runtime_locked: + raise RuntimeError( + "Cannot reassign runtime after workflow has been launched" + ) + old = self._runtime + old.untrack_workflow(self) + self._runtime = new_runtime + new_runtime.track_workflow(self) + @property def workflow_name(self) -> str: """ @@ -198,6 +211,13 @@ class Workflow(metaclass=WorkflowMeta): cls = self.__class__ return f"{cls.__module__}.{cls.__qualname__}" + def _switch_workflow_name(self, name: str) -> None: + if self._runtime_locked and name != self._workflow_name: + raise RuntimeError( + "Cannot change workflow_name after workflow has been launched" + ) + self._workflow_name = name + def _ensure_start_event_class(self) -> type[StartEvent]: """ Returns the StartEvent type used in this workflow. @@ -410,6 +430,10 @@ class Workflow(metaclass=WorkflowMeta): """ from workflows.context import Context + if not self._runtime_locked: + # don't allow switching runtime after a workflow has been launched + self._runtime_locked = True + # Validate the workflow self._validate() diff --git a/packages/llama-index-workflows/tests/conftest.py b/packages/llama-index-workflows/tests/conftest.py index 86b4db38..e966e60d 100644 --- a/packages/llama-index-workflows/tests/conftest.py +++ b/packages/llama-index-workflows/tests/conftest.py @@ -6,10 +6,12 @@ from typing import Any, AsyncGenerator import pytest from pydantic import Field from workflows.context import Context -from workflows.context.state_store import DictState, InMemoryStateStore +from workflows.context.external_context import ExternalContext from workflows.decorators import step from workflows.events import Event, StartEvent, StopEvent -from workflows.plugins.basic import AsyncioAdapterQueues, ExternalAsyncioAdapter +from workflows.plugins.basic import ( + BasicRuntime, +) from workflows.runtime.types.internal_state import BrokerState from workflows.workflow import Workflow @@ -52,16 +54,15 @@ def events() -> list: @pytest.fixture() async def ctx(workflow: Workflow) -> AsyncGenerator[Context[Any], None]: - from workflows.context.external_context import ExternalContext + runtime = BasicRuntime() - queues = AsyncioAdapterQueues( + _ = runtime._get_or_create_queues( run_id="test-run", init_state=BrokerState.from_workflow(workflow), - state_store=InMemoryStateStore(DictState()), ) ctx = Context._create_external( workflow=workflow, - external_adapter=ExternalAsyncioAdapter(queues=queues), + external_adapter=runtime.get_external_adapter("test-run"), ) assert isinstance(ctx._face, ExternalContext) try: diff --git a/packages/llama-index-workflows/tests/runtime/conftest.py b/packages/llama-index-workflows/tests/runtime/conftest.py index 1069104a..843f93d7 100644 --- a/packages/llama-index-workflows/tests/runtime/conftest.py +++ b/packages/llama-index-workflows/tests/runtime/conftest.py @@ -124,7 +124,7 @@ class MockRunAdapter( workers={}, ) - def on_tick(self, tick: WorkflowTick) -> None: + async def on_tick(self, tick: WorkflowTick) -> None: """Record a tick for replay.""" self._ticks.append(tick) diff --git a/packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py b/packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py index a8f3a498..6936724e 100644 --- a/packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py +++ b/packages/llama-index-workflows/tests/runtime/test_control_loop_transformations.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: MIT # Copyright (c) 2026 LlamaIndex Inc. - """ Unit tests for control loop transformation functions. @@ -10,7 +9,7 @@ testing them in isolation without running the full async control loop. from __future__ import annotations -from typing import Any, cast +from typing import cast import pytest from workflows.decorators import StepConfig @@ -253,7 +252,7 @@ def test_step_worker_failed_with_retry(base_state: BrokerState) -> None: event = MyTestEvent(value=42) add_worker(base_state, event) - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, @@ -278,7 +277,7 @@ def test_step_worker_failed_without_retry(base_state: BrokerState) -> None: event = MyTestEvent(value=42) add_worker(base_state, event) - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, @@ -297,7 +296,7 @@ def test_collected_events(base_state: BrokerState) -> None: add_worker(base_state, event) # Add event - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, @@ -334,12 +333,12 @@ def test_waiters(base_state: BrokerState) -> None: event_type=OtherEvent, ) # Add waiter - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, result=[ - cast(StepFunctionResult[Any], result), + cast(StepFunctionResult, result), ], ) new_state, _ = _process_step_result_tick(tick, base_state, now_seconds=110.0) @@ -634,7 +633,7 @@ def test_step_worker_failed_retry_preserves_delay(base_state: BrokerState) -> No event = MyTestEvent(value=42) add_worker(base_state, event) - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, @@ -676,7 +675,7 @@ def test_step_worker_failed_retry_preserves_first_attempt_at( ) ) - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, @@ -761,11 +760,11 @@ def test_idle_event_emitted_on_transition_to_idle(base_state: BrokerState) -> No event_type=OtherEvent, ) - tick: TickStepResult[Any] = TickStepResult[Any]( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, - result=[cast(StepFunctionResult[Any], result)], + result=[cast(StepFunctionResult, result)], ) new_state, commands = _process_step_result_tick(tick, base_state, now_seconds=110.0) @@ -853,7 +852,7 @@ def test_no_idle_event_when_work_remains(base_state: BrokerState) -> None: EventAttempt(event=MyTestEvent(value=99)) ) - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, @@ -887,7 +886,7 @@ def test_no_idle_event_when_workflow_completes(base_state: BrokerState) -> None: base_state.workers["test_step"].collected_waiters.append(waiter) # Complete the workflow with StopEvent - tick: TickStepResult[Any] = TickStepResult( + tick: TickStepResult = TickStepResult( step_name="test_step", worker_id=0, event=event, diff --git a/packages/llama-index-workflows/tests/runtime/test_tick_serialization.py b/packages/llama-index-workflows/tests/runtime/test_tick_serialization.py new file mode 100644 index 00000000..c7acc964 --- /dev/null +++ b/packages/llama-index-workflows/tests/runtime/test_tick_serialization.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. + +from __future__ import annotations + +import json +import time + +import pytest +from pydantic import TypeAdapter +from workflows.events import Event, StartEvent, StopEvent +from workflows.runtime.types.results import ( + AddCollectedEvent, + AddWaiter, + DeleteCollectedEvent, + DeleteWaiter, + StepWorkerFailed, + StepWorkerResult, +) +from workflows.runtime.types.serialization_helpers import ( + _deserialize_event, + _deserialize_event_type, + _deserialize_exception, + _serialize_event, + _serialize_event_type, + _serialize_exception, +) +from workflows.runtime.types.ticks import ( + TickAddEvent, + TickCancelRun, + TickPublishEvent, + TickStepResult, + TickTimeout, + WorkflowTick, +) + + +class MyEvent(Event): + value: str = "hello" + + +# -- Serialization helper roundtrip tests -- + + +def test_event_roundtrip() -> None: + event = MyEvent(value="world") + serialized = _serialize_event(event) + result = _deserialize_event(serialized) + assert isinstance(result, MyEvent) + assert result.value == "world" + + +def test_exception_roundtrip() -> None: + exc = ValueError("something went wrong") + serialized = _serialize_exception(exc) + result = _deserialize_exception(serialized) + assert isinstance(result, ValueError) + assert str(result) == "something went wrong" + + +def test_exception_roundtrip_unimportable() -> None: + CustomError = type("CustomError", (Exception,), {}) + exc = CustomError("oops") + serialized = _serialize_exception(exc) + result = _deserialize_exception(serialized) + assert type(result) is Exception + assert str(result) == "oops" + + +def test_event_type_roundtrip() -> None: + serialized = _serialize_event_type(MyEvent) + result = _deserialize_event_type(serialized) + assert result is MyEvent + + +# -- Tick roundtrip tests -- + + +@pytest.mark.parametrize( + "tick", + [ + pytest.param( + TickAddEvent( + event=StartEvent(), + step_name="my_step", + attempts=3, + first_attempt_at=1234567890.0, + ), + id="add_event", + ), + pytest.param( + TickPublishEvent(event=MyEvent(value="world")), + id="publish_event", + ), + pytest.param( + TickCancelRun(), + id="cancel_run", + ), + pytest.param( + TickTimeout(timeout=30.5), + id="timeout", + ), + pytest.param( + TickStepResult( + step_name="process", + worker_id=42, + event=MyEvent(value="trigger"), + result=[StepWorkerResult(result=StopEvent(result="done"))], + ), + id="step_result_with_event", + ), + pytest.param( + TickStepResult( + step_name="process", + worker_id=1, + event=StartEvent(), + result=[StepWorkerResult(result=None)], + ), + id="step_result_with_none", + ), + pytest.param( + TickStepResult( + step_name="collector", + worker_id=2, + event=StartEvent(), + result=[ + AddCollectedEvent( + event_id="evt-1", event=MyEvent(value="collected") + ) + ], + ), + id="step_result_add_collected_event", + ), + pytest.param( + TickStepResult( + step_name="collector", + worker_id=3, + event=StartEvent(), + result=[DeleteCollectedEvent(event_id="evt-2")], + ), + id="step_result_delete_collected_event", + ), + pytest.param( + TickStepResult( + step_name="cleanup", + worker_id=5, + event=StartEvent(), + result=[DeleteWaiter(waiter_id="w-2")], + ), + id="step_result_delete_waiter", + ), + ], +) +def test_tick_roundtrip(tick: WorkflowTick) -> None: + serialized = tick.model_dump(mode="json") + roundtripped = json.loads(json.dumps(serialized)) + result = type(tick).model_validate(roundtripped) + assert result == tick + + +# -- Tick roundtrip tests with lossy serialization -- + + +def test_tick_step_result_with_failed_value_error() -> None: + failed_at = time.time() + tick = TickStepResult( + step_name="broken_step", + worker_id=7, + event=StartEvent(), + result=[ + StepWorkerFailed( + exception=ValueError("something went wrong"), failed_at=failed_at + ) + ], + ) + serialized = tick.model_dump(mode="json") + roundtripped = json.loads(json.dumps(serialized)) + result = TickStepResult.model_validate(roundtripped) + + assert isinstance(result, TickStepResult) + r = result.result[0] + assert isinstance(r, StepWorkerFailed) + assert isinstance(r.exception, ValueError) + assert str(r.exception) == "something went wrong" + assert r.failed_at == failed_at + + +def test_tick_step_result_with_failed_unimportable_exception() -> None: + CustomError = type("CustomError", (Exception,), {}) + failed_at = time.time() + tick = TickStepResult( + step_name="broken_step", + worker_id=8, + event=StartEvent(), + result=[StepWorkerFailed(exception=CustomError("oops"), failed_at=failed_at)], + ) + serialized = tick.model_dump(mode="json") + roundtripped = json.loads(json.dumps(serialized)) + result = TickStepResult.model_validate(roundtripped) + + assert isinstance(result, TickStepResult) + r = result.result[0] + assert isinstance(r, StepWorkerFailed) + assert type(r.exception) is Exception + assert str(r.exception) == "oops" + assert r.failed_at == failed_at + + +def test_tick_step_result_with_add_waiter() -> None: + tick = TickStepResult( + step_name="waiter_step", + worker_id=4, + event=StartEvent(), + result=[ + AddWaiter( + waiter_id="w-1", + waiter_event=MyEvent(value="waiting"), + requirements={"key": "value"}, + timeout=60.0, + event_type=MyEvent, + ) + ], + ) + serialized = tick.model_dump(mode="json") + + # Verify the serialized form captures has_requirements correctly + waiter_data = serialized["result"][0] + assert waiter_data["has_requirements"] is True + assert waiter_data["requirements"] == {} + + roundtripped = json.loads(json.dumps(serialized)) + result = TickStepResult.model_validate(roundtripped) + + assert isinstance(result, TickStepResult) + r = result.result[0] + assert isinstance(r, AddWaiter) + assert r.waiter_id == "w-1" + assert isinstance(r.waiter_event, MyEvent) + assert r.waiter_event.value == "waiting" + # Requirements are always {} after deserialization + assert r.requirements == {} + assert r.timeout == 60.0 + assert r.event_type is MyEvent + + +# -- WorkflowTick discriminated union tests -- + + +def test_workflow_tick_discriminated_union_roundtrip() -> None: + """Verify that WorkflowTick TypeAdapter can roundtrip all tick types.""" + adapter = TypeAdapter(WorkflowTick) + + ticks = [ + TickAddEvent(event=StartEvent(), step_name="s"), + TickPublishEvent(event=MyEvent(value="x")), + TickCancelRun(), + TickTimeout(timeout=10.0), + TickStepResult( + step_name="s", + worker_id=0, + event=StartEvent(), + result=[StepWorkerResult(result=None)], + ), + ] + for tick in ticks: + dumped = adapter.dump_python(tick, mode="json") + roundtripped = json.loads(json.dumps(dumped)) + restored = adapter.validate_python(roundtripped) + assert type(restored) is type(tick) diff --git a/packages/llama-index-workflows/tests/runtime/test_workflow_set_and_tracking.py b/packages/llama-index-workflows/tests/runtime/test_workflow_set_and_tracking.py new file mode 100644 index 00000000..464d1535 --- /dev/null +++ b/packages/llama-index-workflows/tests/runtime/test_workflow_set_and_tracking.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Tests for WorkflowSet, runtime tracking, and workflow mutation.""" + +from __future__ import annotations + +import gc +from typing import Any + +import pytest +from workflows import Workflow, step +from workflows.events import StartEvent, StopEvent +from workflows.plugins import BasicRuntime +from workflows.runtime.types.plugin import WorkflowSet + + +class SimpleWorkflow(Workflow): + @step + async def start(self, ev: StartEvent) -> StopEvent: + return StopEvent(result="done") + + +class UnhashableWorkflow(Workflow): + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.data = [1, 2, 3] # Makes it unhashable + + @step + async def start(self, ev: StartEvent) -> StopEvent: + return StopEvent(result="done") + + +@pytest.fixture +def basic_runtime() -> BasicRuntime: + return BasicRuntime() + + +@pytest.fixture +def workflow_set() -> WorkflowSet: + return WorkflowSet() + + +# --------------------------------------------------------------------------- +# WorkflowSet tests +# --------------------------------------------------------------------------- + + +def test_workflow_set_add_and_contains( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + workflow_set.add(wf) + assert wf in workflow_set + + +def test_workflow_set_add_unhashable_workflow( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + wf = UnhashableWorkflow(runtime=basic_runtime) + workflow_set.add(wf) + assert wf in workflow_set + items = list(workflow_set) + assert wf in items + + +def test_workflow_set_discard( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + workflow_set.add(wf) + assert wf in workflow_set + workflow_set.discard(wf) + assert wf not in workflow_set + + +def test_workflow_set_len_and_bool( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + assert not workflow_set + assert len(workflow_set) == 0 + + wf = SimpleWorkflow(runtime=basic_runtime) + workflow_set.add(wf) + assert workflow_set + assert len(workflow_set) == 1 + + +def test_workflow_set_iter( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + wf1 = SimpleWorkflow(runtime=basic_runtime) + wf2 = SimpleWorkflow(runtime=basic_runtime) + wf3 = SimpleWorkflow(runtime=basic_runtime) + workflow_set.add(wf1) + workflow_set.add(wf2) + workflow_set.add(wf3) + items = set(id(w) for w in workflow_set) + assert items == {id(wf1), id(wf2), id(wf3)} + + +def test_workflow_set_gc_cleanup( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + workflow_set.add(wf) + assert len(workflow_set) == 1 + del wf + gc.collect() + assert len(workflow_set) == 0 + + +def test_workflow_set_add_idempotent( + workflow_set: WorkflowSet, basic_runtime: BasicRuntime +) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + workflow_set.add(wf) + workflow_set.add(wf) + assert len(workflow_set) == 1 + + +# --------------------------------------------------------------------------- +# Runtime tracking tests +# --------------------------------------------------------------------------- + + +def test_track_workflow_adds_to_set(basic_runtime: BasicRuntime) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + assert wf in basic_runtime._pending + + +def test_untrack_workflow_removes_from_set(basic_runtime: BasicRuntime) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + assert wf in basic_runtime._pending + basic_runtime.untrack_workflow(wf) + assert wf not in basic_runtime._pending + + +def test_launch_locks_tracked_workflows(basic_runtime: BasicRuntime) -> None: + wf1 = SimpleWorkflow(runtime=basic_runtime) + wf2 = SimpleWorkflow(runtime=basic_runtime) + basic_runtime.launch() + assert wf1._runtime_locked is True + assert wf2._runtime_locked is True + + +def test_relaunch_locks_new_workflows(basic_runtime: BasicRuntime) -> None: + wf1 = SimpleWorkflow(runtime=basic_runtime) + basic_runtime.launch() + assert wf1._runtime_locked is True + + wf2 = SimpleWorkflow(runtime=basic_runtime) + assert wf2._runtime_locked is False + basic_runtime.launch() + assert wf1._runtime_locked is True + assert wf2._runtime_locked is True + + +def test_weak_reference_cleanup(basic_runtime: BasicRuntime) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + assert len(basic_runtime._pending) == 1 + del wf + gc.collect() + assert len(basic_runtime._pending) == 0 + + +def test_basic_runtime_launch_sets_launched_flag(basic_runtime: BasicRuntime) -> None: + assert basic_runtime._launched is False + basic_runtime.launch() + assert basic_runtime._launched is True + + +# --------------------------------------------------------------------------- +# Workflow mutation tests +# --------------------------------------------------------------------------- + + +def test_workflow_name_setter(basic_runtime: BasicRuntime) -> None: + wf = SimpleWorkflow(runtime=basic_runtime) + wf._switch_workflow_name("custom-name") + assert wf.workflow_name == "custom-name" + + +def test_workflow_name_setter_raises_after_launch() -> None: + rt = BasicRuntime() + wf = SimpleWorkflow(runtime=rt) + wf._switch_workflow_name("before-launch") + assert wf.workflow_name == "before-launch" + + rt.launch() + with pytest.raises(RuntimeError, match="Cannot change workflow_name"): + wf._switch_workflow_name("after-launch") + + +def test_runtime_setter_swaps_tracking() -> None: + rt1 = BasicRuntime() + rt2 = BasicRuntime() + wf = SimpleWorkflow(runtime=rt1) + assert wf in rt1._pending + assert wf not in rt2._pending + + wf._switch_runtime(rt2) + assert wf not in rt1._pending + assert wf in rt2._pending + + +def test_runtime_setter_post_launch_raises() -> None: + rt1 = BasicRuntime() + rt2 = BasicRuntime() + wf = SimpleWorkflow(runtime=rt1) + rt1.launch() + + with pytest.raises(RuntimeError, match="Cannot reassign runtime"): + wf._switch_runtime(rt2) + + +def test_runtime_setter_same_runtime_after_launch_is_noop() -> None: + rt = BasicRuntime() + wf = SimpleWorkflow(runtime=rt) + rt.launch() + # Assigning the same runtime should not raise + wf._switch_runtime(rt) + assert wf.runtime is rt + + +@pytest.mark.asyncio +async def test_run_locks_runtime() -> None: + rt1 = BasicRuntime() + rt2 = BasicRuntime() + wf = SimpleWorkflow(runtime=rt1) + assert wf._runtime_locked is False + + handler = wf.run() + assert wf._runtime_locked is True + + with pytest.raises(RuntimeError, match="Cannot reassign runtime"): + wf._switch_runtime(rt2) + + await handler + + +def test_runtime_setter_before_launch_then_launch_locks() -> None: + rt1 = BasicRuntime() + rt2 = BasicRuntime() + wf = SimpleWorkflow(runtime=rt1) + wf._switch_runtime(rt2) + assert wf._runtime_locked is False + rt2.launch() + assert wf._runtime_locked is True