refactor _done out

Author: adrianlyjak

src/workflows/context/context.py

@@ -22,13 +22,11 @@
 from workflows.decorators import StepConfig
 from workflows.errors import (
     ContextSerdeError,
-    WorkflowDone,
     WorkflowRuntimeError,
 )
 from workflows.events import (
     Event,
     StartEvent,
-    StopEvent,
 )
 from workflows.runtime.state import WorkflowBrokerState
 from workflows.runtime.broker import WorkflowBroker
@@ -151,7 +149,7 @@ def __init__(
 
         state_types: set[Type[BaseModel]] = set()
         for _, step_func in workflow._get_steps().items():
-            step_config: StepConfig = getattr(step_func, "__step_config")
+            step_config: StepConfig = step_func._step_config
             if (
                 step_config.context_state_type is not None
                 and step_config.context_state_type != DictState
@@ -245,21 +243,6 @@ async def after_complete() -> None:
             after_complete=after_complete,
         )
 
-    def _internal_finalize_run(
-        self,
-        event: StopEvent,
-        result: RunResultT,
-    ) -> None:
-        """
-        Called internally from the workflow on a context's completion
-        """
-        if self._broker_run is None:
-            raise WorkflowRuntimeError("Workflow run is not yet running")
-        self._broker_run.finalize_run(event, result)
-        self.write_event_to_stream(event)
-        # Signal we want to stop the workflow
-        raise WorkflowDone
-
     def _internal_cancel_run(self) -> None:
         """
         Called internally from the handler to cancel a context's run

src/workflows/decorators.py

@@ -55,9 +55,9 @@ class StepConfig(BaseModel):
 
 
 class StepFunction(Protocol, Generic[P, R]):
-    """A decorated function, that has some __step_config metadata from the @step decorator"""
+    """A decorated function, that has some step_config metadata from the @step decorator"""
 
-    __step_config: StepConfig
+    _step_config: StepConfig
 
     __name__: str
     __qualname__: str
@@ -135,24 +135,7 @@ def decorator(func: Callable[P, R]) -> StepFunction[P, R]:
                 "num_workers must be an integer greater than 0"
             )
 
-        # This will raise providing a message with the specific validation failure
-        spec = inspect_signature(func)
-        validate_step_signature(spec)
-
-        event_name, accepted_events = next(iter(spec.accepted_events.items()))
-
-        func = cast(StepFunction[P, R], func)
-        # store the configuration in the function object
-        func.__step_config = StepConfig(
-            accepted_events=accepted_events,
-            event_name=event_name,
-            return_types=spec.return_types,
-            context_parameter=spec.context_parameter,
-            context_state_type=spec.context_state_type,
-            num_workers=num_workers,
-            retry_policy=retry_policy,
-            resources=spec.resources,
-        )
+        func = make_step_function(func, num_workers, retry_policy)
 
         # If this is a free function, call add_step() explicitly.
         if is_free_function(func.__qualname__):
@@ -167,3 +150,27 @@ def decorator(func: Callable[P, R]) -> StepFunction[P, R]:
         # The decorator was used without parentheses, like `@step`
         return decorator(func)
     return decorator
+
+
+def make_step_function(
+    func: Callable[P, R], num_workers: int = 4, retry_policy: RetryPolicy | None = None
+) -> StepFunction[P, R]:
+    # This will raise providing a message with the specific validation failure
+    spec = inspect_signature(func)
+    validate_step_signature(spec)
+
+    event_name, accepted_events = next(iter(spec.accepted_events.items()))
+
+    casted = cast(StepFunction[P, R], func)
+    casted._step_config = StepConfig(
+        accepted_events=accepted_events,
+        event_name=event_name,
+        return_types=spec.return_types,
+        context_parameter=spec.context_parameter,
+        context_state_type=spec.context_state_type,
+        num_workers=num_workers,
+        retry_policy=retry_policy,
+        resources=spec.resources,
+    )
+
+    return casted

src/workflows/runtime/broker.py

@@ -23,7 +23,7 @@
 
 from llama_index_instrumentation.dispatcher import Dispatcher
 
-from workflows.decorators import StepConfig
+from workflows.decorators import StepConfig, StepFunction, step
 from workflows.errors import (
     WorkflowCancelledByUser,
     WorkflowDone,
@@ -111,6 +111,86 @@ class WorkflowBroker(Generic[MODEL_T]):
     # run state
     _handler: WorkflowHandler | None
 
+    def __init__(
+        self,
+        workflow: Workflow,
+        context: Context[MODEL_T],
+        state: WorkflowBrokerState,
+        run_id: str,
+    ) -> None:
+        self._context = context
+        self._handler = None
+
+        # Store the step configs of this workflow, to be used in send_event
+        self._step_configs = {}
+
+        steps: dict[str, StepFunction] = {
+            "_done": self._done,
+            **workflow._get_steps(),
+        }
+        for step_name, step_func in steps.items():
+            self._step_configs[step_name] = step_func._step_config
+
+        # Transient runtime fields (always reinitialized)
+        self._tasks = set()
+        self._cancel_flag = asyncio.Event()
+        self._step_flags = {}
+        self._step_events_holding = None
+        self._step_lock = asyncio.Lock()
+        self._retval = None
+
+        self._lock = asyncio.Lock()
+
+        self._dispatcher = workflow._dispatcher
+
+        self._step_condition = asyncio.Condition(lock=self._step_lock)
+        self._step_event_written = asyncio.Condition(lock=self._step_lock)
+        # Keep track of the steps currently running (transient)
+        self._currently_running_steps = defaultdict(int)
+        # Default initial values for persistent fields
+
+        self._state = state
+
+        # initialize running state from workflow
+        for name, step_func in steps.items():
+            if name not in self._state.queues:
+                self._state.queues[name] = asyncio.Queue()
+
+            if name not in self._step_flags:
+                self._step_flags[name] = asyncio.Event()
+
+            step_config: StepConfig = step_func._step_config
+
+            # Make the system step "_done" accept custom stop events
+            if (
+                name == "_done"
+                and workflow._stop_event_class not in step_config.accepted_events
+            ):
+                step_config.accepted_events.append(workflow._stop_event_class)
+
+            for _ in range(step_config.num_workers):
+                self._add_step_worker(
+                    name=name,
+                    step=step_func,
+                    config=step_config,
+                    verbose=workflow._verbose,
+                    run_id=run_id,
+                    worker_id=str(uuid.uuid4()),
+                    resource_manager=workflow._resource_manager,
+                )
+
+        # add dedicated cancel task
+        self._add_cancel_worker()
+
+    @step(num_workers=1)
+    async def _done(self, ev: StopEvent) -> None:
+        """Tears down the whole workflow and stop execution."""
+        result = ev.result if type(ev) is StopEvent else ev
+        self.finalize_run(ev, result)
+        self.write_event_to_stream(ev)
+        # Signal we want to stop the workflow
+        raise WorkflowDone
+
     @property
     def is_running(self) -> bool:
         return self._state.is_running
@@ -359,73 +439,6 @@ async def shutdown(self) -> None:
         await asyncio.gather(*self._tasks, return_exceptions=True)
         self._tasks.clear()
 
-    def __init__(
-        self,
-        workflow: Workflow,
-        context: Context[MODEL_T],
-        state: WorkflowBrokerState,
-        run_id: str,
-    ) -> None:
-        self._context = context
-        self._handler = None
-
-        # Store the step configs of this workflow, to be used in send_event
-        self._step_configs = {}
-        for step_name, step_func in workflow._get_steps().items():
-            self._step_configs[step_name] = getattr(step_func, "__step_config", None)
-
-        # Transient runtime fields (always reinitialized)
-        self._tasks = set()
-        self._cancel_flag = asyncio.Event()
-        self._step_flags = {}
-        self._step_events_holding = None
-        self._step_lock = asyncio.Lock()
-        self._retval = None
-
-        self._lock = asyncio.Lock()
-
-        self._dispatcher = workflow._dispatcher
-
-        self._step_condition = asyncio.Condition(lock=self._step_lock)
-        self._step_event_written = asyncio.Condition(lock=self._step_lock)
-        # Keep track of the steps currently running (transient)
-        self._currently_running_steps = defaultdict(int)
-        # Default initial values for persistent fields
-
-        self._state = state
-
-        # initialize running state from workflow
-        for name, step_func in workflow._get_steps().items():
-            if name not in self._state.queues:
-                self._state.queues[name] = asyncio.Queue()
-
-            if name not in self._step_flags:
-                self._step_flags[name] = asyncio.Event()
-
-            # At this point, step_func is guaranteed to have the `__step_config` attribute
-            step_config: StepConfig = getattr(step_func, "__step_config")
-
-            # Make the system step "_done" accept custom stop events
-            if (
-                name == "_done"
-                and workflow._stop_event_class not in step_config.accepted_events
-            ):
-                step_config.accepted_events.append(workflow._stop_event_class)
-
-            for _ in range(step_config.num_workers):
-                self._add_step_worker(
-                    name=name,
-                    step=step_func,
-                    config=step_config,
-                    verbose=workflow._verbose,
-                    run_id=run_id,
-                    worker_id=str(uuid.uuid4()),
-                    resource_manager=workflow._resource_manager,
-                )
-
-        # add dedicated cancel task
-        self._add_cancel_worker()
-
     async def _mark_in_progress(
         self, name: str, ev: Event, worker_id: str = ""
     ) -> None:

src/workflows/server/representation_utils.py

@@ -85,9 +85,7 @@ def _extract_workflow_structure(
     # Assuming that `Workflow` is validated before drawing, it's enough to find the first one.
     current_stop_event = None
     for step_name, step_func in steps.items():
-        step_config = getattr(step_func, "__step_config", None)
-        if step_config is None:
-            continue
+        step_config = step_func._step_config
 
         for return_type in step_config.return_types:
             if issubclass(return_type, StopEvent):
@@ -99,10 +97,8 @@ def _extract_workflow_structure(
 
     # First pass: Add all nodes
     for step_name, step_func in steps.items():
-        step_config = getattr(step_func, "__step_config", None)
-        if step_config is None:
-            continue
-
+        step_config = step_func._step_config
+        
         # Add step node
         step_label = (
             _truncate_label(step_name, max_label_length)
@@ -198,10 +194,8 @@ def _extract_workflow_structure(
 
     # Second pass: Add edges
     for step_name, step_func in steps.items():
-        step_config = getattr(step_func, "__step_config", None)
-        if step_config is None:
-            continue
-
+        step_config = step_func._step_config
+        
         # Edges from steps to return types
         for return_type in step_config.return_types:
             if return_type is not type(None):

src/workflows/utils.py

@@ -169,7 +169,7 @@ def get_steps_from_class(_class: object) -> dict[str, StepFunction]:
     all_methods = inspect.getmembers(_class, predicate=inspect.isfunction)
 
     for name, method in all_methods:
-        if hasattr(method, "__step_config"):
+        if hasattr(method, "_step_config"):
             step_methods[name] = cast(StepFunction, method)
 
     return step_methods
@@ -192,7 +192,7 @@ def get_steps_from_instance(workflow: object) -> dict[str, StepFunction]:
     all_methods = inspect.getmembers(workflow, predicate=inspect.ismethod)
 
     for name, method in all_methods:
-        if hasattr(method, "__step_config"):
+        if hasattr(method, "_step_config"):
             step_methods[name] = cast(StepFunction, method)
 
     return step_methods

src/workflows/workflow.py

@@ -14,7 +14,7 @@
 from pydantic import ValidationError
 
 from .context import Context
-from .decorators import StepConfig, StepFunction, step
+from .decorators import StepConfig, StepFunction
 from .errors import (
     WorkflowConfigurationError,
     WorkflowRuntimeError,
@@ -135,7 +135,7 @@ def _ensure_start_event_class(self) -> type[StartEvent]:
         """
         start_events_found: set[type[StartEvent]] = set()
         for step_func in self._get_steps().values():
-            step_config: StepConfig = getattr(step_func, "__step_config")
+            step_config: StepConfig = step_func._step_config
             for event_type in step_config.accepted_events:
                 if issubclass(event_type, StartEvent):
                     start_events_found.add(event_type)
@@ -196,7 +196,7 @@ def _ensure_stop_event_class(self) -> type[RunResultT]:
         """
         stop_events_found: set[type[StopEvent]] = set()
         for step_func in self._get_steps().values():
-            step_config: StepConfig = getattr(step_func, "__step_config")
+            step_config: StepConfig = step_func._step_config
             for event_type in step_config.return_types:
                 if issubclass(event_type, StopEvent):
                     stop_events_found.add(event_type)
@@ -226,7 +226,7 @@ def add_step(cls, func: StepFunction) -> None:
 
         It raises an exception if a step with the same name was already added to the workflow.
         """
-        step_config: StepConfig | None = getattr(func, "__step_config", None)
+        step_config: StepConfig | None = getattr(func, "_step_config", None)
         if not step_config:
             msg = f"Step function {func.__name__} is missing the `@step` decorator."
             raise WorkflowValidationError(msg)
@@ -336,13 +336,6 @@ def run(
             workflow=self, start_event=start_event_instance, semaphore=self._sem
         )
 
-    @step(num_workers=1)
-    async def _done(self, ctx: Context, ev: StopEvent) -> None:
-        """Tears down the whole workflow and stop execution."""
-        ctx._internal_finalize_run(
-            ev, ev.result if self._stop_event_class is StopEvent else ev
-        )
-
     def _validate(self) -> bool:
         """
         Validate the workflow to ensure it's well-formed.
@@ -359,9 +352,7 @@ def _validate(self) -> bool:
         steps_accepting_stop_event: list[str] = []
 
         for name, step_func in self._get_steps().items():
-            step_config: StepConfig | None = getattr(step_func, "__step_config")
-            # At this point we know step config is not None, let's make the checker happy
-            assert step_config is not None
+            step_config: StepConfig  = step_func._step_config
 
             # Check that no user-defined step accepts StopEvent (only _done step should)
             if name != "_done":