vellum-ai 0.10.8__py3-none-any.whl → 0.11.0__py3-none-any.whl

vellum/workflows/runner/runner.py

@@ -3,7 +3,7 @@ from copy import deepcopy
 import logging
 from queue import Empty, Queue
 from threading import Event as ThreadingEvent, Thread
-from uuid import UUID, uuid4
+from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Type, Union
 
 from vellum.workflows.constants import UNDEF
@@ -29,7 +29,6 @@ from vellum.workflows.events.node import (
     NodeExecutionStreamingBody,
 )
 from vellum.workflows.events.types import BaseEvent, ParentContext, WorkflowParentContext
-from vellum.workflows.events.utils import is_terminal_event
 from vellum.workflows.events.workflow import (
     WorkflowExecutionFulfilledBody,
     WorkflowExecutionInitiatedBody,
@@ -38,6 +37,8 @@ from vellum.workflows.events.workflow import (
     WorkflowExecutionRejectedBody,
     WorkflowExecutionResumedBody,
     WorkflowExecutionResumedEvent,
+    WorkflowExecutionSnapshottedBody,
+    WorkflowExecutionSnapshottedEvent,
     WorkflowExecutionStreamingBody,
 )
 from vellum.workflows.exceptions import NodeException
@@ -46,7 +47,6 @@ from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references import ExternalInputReference, OutputReference
-from vellum.workflows.runner.types import WorkItemEvent
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.generics import OutputsType, StateType, WorkflowInputsType
 
@@ -77,6 +77,7 @@ class WorkflowRunner(Generic[StateType]):
             raise ValueError("Can only run a Workflow providing one of state or external inputs, not both")
 
         self.workflow = workflow
+        self._is_resuming = False
         if entrypoint_nodes:
             if len(list(entrypoint_nodes)) > 1:
                 raise ValueError("Cannot resume from multiple nodes")
@@ -99,6 +100,7 @@ class WorkflowRunner(Generic[StateType]):
                 for ei in external_inputs
                 if issubclass(ei.inputs_class.__parent_class__, BaseNode)
             ]
+            self._is_resuming = True
         else:
             normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
             if state:
@@ -108,9 +110,16 @@ class WorkflowRunner(Generic[StateType]):
                 self._initial_state = self.workflow.get_default_state(normalized_inputs)
             self._entrypoints = self.workflow.get_entrypoints()
 
-        self._work_item_event_queue: Queue[WorkItemEvent[StateType]] = Queue()
-        self._workflow_event_queue: Queue[WorkflowEvent] = Queue()
+        # This queue is responsible for sending events from WorkflowRunner to the outside world
+        self._workflow_event_outer_queue: Queue[WorkflowEvent] = Queue()
+
+        # This queue is responsible for sending events from the inner worker threads to WorkflowRunner
+        self._workflow_event_inner_queue: Queue[WorkflowEvent] = Queue()
+
+        # This queue is responsible for sending events from WorkflowRunner to the background thread
+        # for user defined emitters
         self._background_thread_queue: Queue[BackgroundThreadItem] = Queue()
+
         self._dependencies: Dict[Type[BaseNode], Set[Type[BaseNode]]] = defaultdict(set)
         self._state_forks: Set[StateType] = {self._initial_state}
 
@@ -123,8 +132,20 @@ class WorkflowRunner(Generic[StateType]):
             "__snapshot_callback__",
             lambda s: self._snapshot_state(s),
         )
+        self.workflow.context._register_event_queue(self._workflow_event_inner_queue)
 
     def _snapshot_state(self, state: StateType) -> StateType:
+        self._workflow_event_inner_queue.put(
+            WorkflowExecutionSnapshottedEvent(
+                trace_id=state.meta.trace_id,
+                span_id=state.meta.span_id,
+                body=WorkflowExecutionSnapshottedBody(
+                    workflow_definition=self.workflow.__class__,
+                    state=state,
+                ),
+                parent=self._parent_context,
+            )
+        )
         self.workflow._store.append_state_snapshot(state)
         self._background_thread_queue.put(state)
         return state
@@ -135,21 +156,19 @@ class WorkflowRunner(Generic[StateType]):
         return event
 
     def _run_work_item(self, node: BaseNode[StateType], span_id: UUID) -> None:
-        self._work_item_event_queue.put(
-            WorkItemEvent(
-                node=node,
-                event=NodeExecutionInitiatedEvent(
-                    trace_id=node.state.meta.trace_id,
+        self._workflow_event_inner_queue.put(
+            NodeExecutionInitiatedEvent(
+                trace_id=node.state.meta.trace_id,
+                span_id=span_id,
+                body=NodeExecutionInitiatedBody(
+                    node_definition=node.__class__,
+                    inputs=node._inputs,
+                ),
+                parent=WorkflowParentContext(
                     span_id=span_id,
-                    body=NodeExecutionInitiatedBody(
-                        node_definition=node.__class__,
-                        inputs=node._inputs,
-                    ),
-                    parent=WorkflowParentContext(
-                        span_id=span_id,
-                        workflow_definition=self.workflow.__class__,
-                        parent=self._parent_context
-                    )
+                    workflow_definition=self.workflow.__class__,
+                    parent=self._parent_context,
+                    type="WORKFLOW",
                 ),
             )
         )
@@ -186,24 +205,23 @@ class WorkflowRunner(Generic[StateType]):
                         outputs_class=node.Outputs,
                     )
                     node.state.meta.node_outputs[output_descriptor] = streaming_output_queues[output.name]
-                    self._work_item_event_queue.put(
-                        WorkItemEvent(
-                            node=node,
-                            event=NodeExecutionStreamingEvent(
-                                trace_id=node.state.meta.trace_id,
+                    initiated_output: BaseOutput = BaseOutput(name=output.name)
+                    initiated_ports = initiated_output > ports
+                    self._workflow_event_inner_queue.put(
+                        NodeExecutionStreamingEvent(
+                            trace_id=node.state.meta.trace_id,
+                            span_id=span_id,
+                            body=NodeExecutionStreamingBody(
+                                node_definition=node.__class__,
+                                output=initiated_output,
+                                invoked_ports=initiated_ports,
+                            ),
+                            parent=WorkflowParentContext(
                                 span_id=span_id,
-                                body=NodeExecutionStreamingBody(
-                                    node_definition=node.__class__,
-                                    output=BaseOutput(name=output.name),
-                                    invoked_ports=invoked_ports,
-                                ),
-                                parent=WorkflowParentContext(
-                                    span_id=span_id,
-                                    workflow_definition=self.workflow.__class__,
-                                    parent=self._parent_context
-                                )
+                                workflow_definition=self.workflow.__class__,
+                                parent=self._parent_context,
                             ),
-                        )
+                        ),
                     )
 
                 for output in node_run_response:
@@ -215,50 +233,47 @@ class WorkflowRunner(Generic[StateType]):
                             initiate_node_streaming_output(output)
 
                         streaming_output_queues[output.name].put(output.delta)
-                        self._work_item_event_queue.put(
-                            WorkItemEvent(
-                                node=node,
-                                event=NodeExecutionStreamingEvent(
-                                    trace_id=node.state.meta.trace_id,
+                        self._workflow_event_inner_queue.put(
+                            NodeExecutionStreamingEvent(
+                                trace_id=node.state.meta.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=output,
+                                    invoked_ports=invoked_ports,
+                                ),
+                                parent=WorkflowParentContext(
                                     span_id=span_id,
-                                    body=NodeExecutionStreamingBody(
-                                        node_definition=node.__class__,
-                                        output=output,
-                                        invoked_ports=invoked_ports,
-                                    ),
-                                    parent=WorkflowParentContext(
-                                        span_id=span_id,
-                                        workflow_definition=self.workflow.__class__,
-                                        parent=self._parent_context,
-                                    )
+                                    workflow_definition=self.workflow.__class__,
+                                    parent=self._parent_context,
                                 ),
-                            )
+                            ),
                         )
                     elif output.is_fulfilled:
                         if output.name in streaming_output_queues:
                             streaming_output_queues[output.name].put(UNDEF)
 
                         setattr(outputs, output.name, output.value)
-                        self._work_item_event_queue.put(
-                            WorkItemEvent(
-                                node=node,
-                                event=NodeExecutionStreamingEvent(
-                                    trace_id=node.state.meta.trace_id,
+                        self._workflow_event_inner_queue.put(
+                            NodeExecutionStreamingEvent(
+                                trace_id=node.state.meta.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=output,
+                                    invoked_ports=invoked_ports,
+                                ),
+                                parent=WorkflowParentContext(
                                     span_id=span_id,
-                                    body=NodeExecutionStreamingBody(
-                                        node_definition=node.__class__,
-                                        output=output,
-                                        invoked_ports=invoked_ports,
-                                    ),
-                                    parent=WorkflowParentContext(
-                                        span_id=span_id,
-                                        workflow_definition=self.workflow.__class__,
-                                        parent=self._parent_context,
-                                    )
+                                    workflow_definition=self.workflow.__class__,
+                                    parent=self._parent_context,
                                 ),
                             )
                         )
 
+            invoked_ports = ports(outputs, node.state)
+            node.state.meta.node_execution_cache.fulfill_node_execution(node.__class__, span_id)
+
             for descriptor, output_value in outputs:
                 if output_value is UNDEF:
                     if descriptor in node.state.meta.node_outputs:
@@ -267,70 +282,58 @@ class WorkflowRunner(Generic[StateType]):
 
                 node.state.meta.node_outputs[descriptor] = output_value
 
-            invoked_ports = ports(outputs, node.state)
-            node.state.meta.node_execution_cache.fulfill_node_execution(node.__class__, span_id)
-
-            self._work_item_event_queue.put(
-                WorkItemEvent(
-                    node=node,
-                    event=NodeExecutionFulfilledEvent(
-                        trace_id=node.state.meta.trace_id,
+            self._workflow_event_inner_queue.put(
+                NodeExecutionFulfilledEvent(
+                    trace_id=node.state.meta.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionFulfilledBody(
+                        node_definition=node.__class__,
+                        outputs=outputs,
+                        invoked_ports=invoked_ports,
+                    ),
+                    parent=WorkflowParentContext(
                         span_id=span_id,
-                        body=NodeExecutionFulfilledBody(
-                            node_definition=node.__class__,
-                            outputs=outputs,
-                            invoked_ports=invoked_ports,
-                        ),
-                        parent=WorkflowParentContext(
-                            span_id=span_id,
-                            workflow_definition=self.workflow.__class__,
-                            parent=self._parent_context,
-                        )
+                        workflow_definition=self.workflow.__class__,
+                        parent=self._parent_context,
                     ),
                 )
             )
         except NodeException as e:
-            self._work_item_event_queue.put(
-                WorkItemEvent(
-                    node=node,
-                    event=NodeExecutionRejectedEvent(
-                        trace_id=node.state.meta.trace_id,
+            self._workflow_event_inner_queue.put(
+                NodeExecutionRejectedEvent(
+                    trace_id=node.state.meta.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionRejectedBody(
+                        node_definition=node.__class__,
+                        error=e.error,
+                    ),
+                    parent=WorkflowParentContext(
                         span_id=span_id,
-                        body=NodeExecutionRejectedBody(
-                            node_definition=node.__class__,
-                            error=e.error,
-                        ),
-                        parent=WorkflowParentContext(
-                            span_id=span_id,
-                            workflow_definition=self.workflow.__class__,
-                            parent=self._parent_context,
-                        )
+                        workflow_definition=self.workflow.__class__,
+                        parent=self._parent_context,
                     ),
                 )
             )
         except Exception as e:
             logger.exception(f"An unexpected error occurred while running node {node.__class__.__name__}")
 
-            self._work_item_event_queue.put(
-                WorkItemEvent(
-                    node=node,
-                    event=NodeExecutionRejectedEvent(
-                        trace_id=node.state.meta.trace_id,
-                        span_id=span_id,
-                        body=NodeExecutionRejectedBody(
-                            node_definition=node.__class__,
-                            error=VellumError(
-                                message=str(e),
-                                code=VellumErrorCode.INTERNAL_ERROR,
-                            ),
+            self._workflow_event_inner_queue.put(
+                NodeExecutionRejectedEvent(
+                    trace_id=node.state.meta.trace_id,
+                    span_id=span_id,
+                    body=NodeExecutionRejectedBody(
+                        node_definition=node.__class__,
+                        error=VellumError(
+                            message=str(e),
+                            code=VellumErrorCode.INTERNAL_ERROR,
                         ),
-                        parent=WorkflowParentContext(
-                            span_id=span_id,
-                            workflow_definition=self.workflow.__class__,
-                            parent=self._parent_context
-                        )
                     ),
-                )
+                    parent=WorkflowParentContext(
+                        span_id=span_id,
+                        workflow_definition=self.workflow.__class__,
+                        parent=self._parent_context,
+                    ),
+                ),
             )
 
         logger.debug(f"Finished running node: {node.__class__.__name__}")
@@ -350,7 +353,10 @@ class WorkflowRunner(Generic[StateType]):
                 self._run_node_if_ready(next_state, edge.to_node, edge)
 
     def _run_node_if_ready(
-        self, state: StateType, node_class: Type[BaseNode], invoked_by: Optional[Edge] = None
+        self,
+        state: StateType,
+        node_class: Type[BaseNode],
+        invoked_by: Optional[Edge] = None,
     ) -> None:
         with state.__lock__:
             for descriptor in node_class.ExternalInputs:
@@ -362,22 +368,23 @@ class WorkflowRunner(Generic[StateType]):
                     return
 
             all_deps = self._dependencies[node_class]
-            if not node_class.Trigger.should_initiate(state, all_deps, invoked_by):
+            node_span_id = state.meta.node_execution_cache.queue_node_execution(node_class, all_deps, invoked_by)
+            if not node_class.Trigger.should_initiate(state, all_deps, node_span_id):
                 return
 
             node = node_class(state=state, context=self.workflow.context)
-            node_span_id = uuid4()
             state.meta.node_execution_cache.initiate_node_execution(node_class, node_span_id)
             self._active_nodes_by_execution_id[node_span_id] = node
 
-            worker_thread = Thread(target=self._run_work_item, kwargs={"node": node, "span_id": node_span_id})
+            worker_thread = Thread(
+                target=self._run_work_item,
+                kwargs={"node": node, "span_id": node_span_id},
+            )
             worker_thread.start()
 
-    def _handle_work_item_event(self, work_item_event: WorkItemEvent[StateType]) -> Optional[VellumError]:
-        node = work_item_event.node
-        event = work_item_event.event
-
-        if event.name == "node.execution.initiated":
+    def _handle_work_item_event(self, event: WorkflowEvent) -> Optional[VellumError]:
+        node = self._active_nodes_by_execution_id.get(event.span_id)
+        if not node:
             return None
 
         if event.name == "node.execution.rejected":
@@ -394,7 +401,7 @@ class WorkflowRunner(Generic[StateType]):
                 if node_output_descriptor.name != event.output.name:
                     continue
 
-                self._workflow_event_queue.put(
+                self._workflow_event_outer_queue.put(
                     self._stream_workflow_event(
                         BaseOutput(
                             name=workflow_output_descriptor.name,
@@ -414,7 +421,7 @@ class WorkflowRunner(Generic[StateType]):
 
             return None
 
-        raise ValueError(f"Invalid event name: {event.name}")
+        return None
 
     def _initiate_workflow_event(self) -> WorkflowExecutionInitiatedEvent:
         return WorkflowExecutionInitiatedEvent(
@@ -435,7 +442,7 @@ class WorkflowRunner(Generic[StateType]):
                 workflow_definition=self.workflow.__class__,
                 output=output,
             ),
-            parent=self._parent_context
+            parent=self._parent_context,
         )
 
     def _fulfill_workflow_event(self, outputs: OutputsType) -> WorkflowExecutionFulfilledEvent:
@@ -484,9 +491,12 @@ class WorkflowRunner(Generic[StateType]):
         # TODO: We should likely handle this during initialization
         # https://app.shortcut.com/vellum/story/4327
         if not self._entrypoints:
-            self._workflow_event_queue.put(
+            self._workflow_event_outer_queue.put(
                 self._reject_workflow_event(
-                    VellumError(message="No entrypoints defined", code=VellumErrorCode.INVALID_WORKFLOW)
+                    VellumError(
+                        message="No entrypoints defined",
+                        code=VellumErrorCode.INVALID_WORKFLOW,
+                    )
                 )
             )
             return
@@ -498,12 +508,12 @@ class WorkflowRunner(Generic[StateType]):
             try:
                 self._run_node_if_ready(self._initial_state, node_cls)
             except NodeException as e:
-                self._workflow_event_queue.put(self._reject_workflow_event(e.error))
+                self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
                 return
             except Exception:
                 err_message = f"An unexpected error occurred while initializing node {node_cls.__name__}"
                 logger.exception(err_message)
-                self._workflow_event_queue.put(
+                self._workflow_event_outer_queue.put(
                     self._reject_workflow_event(
                         VellumError(code=VellumErrorCode.INTERNAL_ERROR, message=err_message),
                     )
@@ -516,21 +526,20 @@ class WorkflowRunner(Generic[StateType]):
             if not self._active_nodes_by_execution_id:
                 break
 
-            work_item_event = self._work_item_event_queue.get()
-            event = work_item_event.event
+            event = self._workflow_event_inner_queue.get()
 
-            self._workflow_event_queue.put(event)
+            self._workflow_event_outer_queue.put(event)
 
-            rejection_error = self._handle_work_item_event(work_item_event)
+            rejection_error = self._handle_work_item_event(event)
             if rejection_error:
                 break
 
         # Handle any remaining events
         try:
-            while work_item_event := self._work_item_event_queue.get_nowait():
-                self._workflow_event_queue.put(work_item_event.event)
+            while event := self._workflow_event_inner_queue.get_nowait():
+                self._workflow_event_outer_queue.put(event)
 
-                rejection_error = self._handle_work_item_event(work_item_event)
+                rejection_error = self._handle_work_item_event(event)
                 if rejection_error:
                     break
         except Empty:
@@ -546,14 +555,13 @@ class WorkflowRunner(Generic[StateType]):
             if node_input_value is UNDEF
         }
         if unresolved_external_inputs:
-            self._workflow_event_queue.put(
+            self._workflow_event_outer_queue.put(
                 self._pause_workflow_event(unresolved_external_inputs),
             )
             return
 
-        final_state.meta.is_terminated = True
         if rejection_error:
-            self._workflow_event_queue.put(self._reject_workflow_event(rejection_error))
+            self._workflow_event_outer_queue.put(self._reject_workflow_event(rejection_error))
             return
 
         fulfilled_outputs = self.workflow.Outputs()
@@ -561,9 +569,13 @@ class WorkflowRunner(Generic[StateType]):
             if isinstance(value, BaseDescriptor):
                 setattr(fulfilled_outputs, descriptor.name, value.resolve(final_state))
             elif isinstance(descriptor.instance, BaseDescriptor):
-                setattr(fulfilled_outputs, descriptor.name, descriptor.instance.resolve(final_state))
+                setattr(
+                    fulfilled_outputs,
+                    descriptor.name,
+                    descriptor.instance.resolve(final_state),
+                )
 
-        self._workflow_event_queue.put(self._fulfill_workflow_event(fulfilled_outputs))
+        self._workflow_event_outer_queue.put(self._fulfill_workflow_event(fulfilled_outputs))
 
     def _run_background_thread(self) -> None:
         state_class = self.workflow.get_state_class()
@@ -584,55 +596,71 @@ class WorkflowRunner(Generic[StateType]):
             return
 
         self._cancel_signal.wait()
-        self._workflow_event_queue.put(
+        self._workflow_event_outer_queue.put(
             self._reject_workflow_event(
-                VellumError(code=VellumErrorCode.WORKFLOW_CANCELLED, message="Workflow run cancelled")
+                VellumError(
+                    code=VellumErrorCode.WORKFLOW_CANCELLED,
+                    message="Workflow run cancelled",
+                )
             )
         )
 
+    def _is_terminal_event(self, event: WorkflowEvent) -> bool:
+        if (
+            event.name == "workflow.execution.fulfilled"
+            or event.name == "workflow.execution.rejected"
+            or event.name == "workflow.execution.paused"
+        ):
+            return event.workflow_definition == self.workflow.__class__
+        return False
+
     def stream(self) -> WorkflowEventStream:
         background_thread = Thread(
-            target=self._run_background_thread, name=f"{self.workflow.__class__.__name__}.background_thread"
+            target=self._run_background_thread,
+            name=f"{self.workflow.__class__.__name__}.background_thread",
         )
         background_thread.start()
 
         if self._cancel_signal:
             cancel_thread = Thread(
-                target=self._run_cancel_thread, name=f"{self.workflow.__class__.__name__}.cancel_thread"
+                target=self._run_cancel_thread,
+                name=f"{self.workflow.__class__.__name__}.cancel_thread",
             )
             cancel_thread.start()
 
         event: WorkflowEvent
-        if self._initial_state.meta.is_terminated or self._initial_state.meta.is_terminated is None:
-            event = self._initiate_workflow_event()
-        else:
+        if self._is_resuming:
             event = self._resume_workflow_event()
+        else:
+            event = self._initiate_workflow_event()
 
         yield self._emit_event(event)
-        self._initial_state.meta.is_terminated = False
 
         # The extra level of indirection prevents the runner from waiting on the caller to consume the event stream
-        stream_thread = Thread(target=self._stream, name=f"{self.workflow.__class__.__name__}.stream_thread")
+        stream_thread = Thread(
+            target=self._stream,
+            name=f"{self.workflow.__class__.__name__}.stream_thread",
+        )
         stream_thread.start()
 
         while stream_thread.is_alive():
             try:
-                event = self._workflow_event_queue.get(timeout=0.1)
+                event = self._workflow_event_outer_queue.get(timeout=0.1)
             except Empty:
                 continue
 
             yield self._emit_event(event)
 
-            if is_terminal_event(event):
+            if self._is_terminal_event(event):
                 break
 
         try:
-            while event := self._workflow_event_queue.get_nowait():
+            while event := self._workflow_event_outer_queue.get_nowait():
                 yield self._emit_event(event)
         except Empty:
             pass
 
-        if not is_terminal_event(event):
+        if not self._is_terminal_event(event):
             yield self._reject_workflow_event(
                 VellumError(
                     code=VellumErrorCode.INTERNAL_ERROR,