vellum-ai 0.10.8__py3-none-any.whl → 0.10.9__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,
@@ -74,7 +73,9 @@ class WorkflowRunner(Generic[StateType]):
         parent_context: Optional[ParentContext] = None,
     ):
         if state and external_inputs:
-            raise ValueError("Can only run a Workflow providing one of state or external inputs, not both")
+            raise ValueError(
+                "Can only run a Workflow providing one of state or external inputs, not both"
+            )
 
         self.workflow = workflow
         if entrypoint_nodes:
@@ -100,7 +101,9 @@ class WorkflowRunner(Generic[StateType]):
                 if issubclass(ei.inputs_class.__parent_class__, BaseNode)
             ]
         else:
-            normalized_inputs = deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
+            normalized_inputs = (
+                deepcopy(inputs) if inputs else self.workflow.get_default_inputs()
+            )
             if state:
                 self._initial_state = deepcopy(state)
                 self._initial_state.meta.workflow_inputs = normalized_inputs
@@ -123,6 +126,7 @@ class WorkflowRunner(Generic[StateType]):
             "__snapshot_callback__",
             lambda s: self._snapshot_state(s),
         )
+        self.workflow.context._register_event_queue(self._workflow_event_queue)
 
     def _snapshot_state(self, state: StateType) -> StateType:
         self.workflow._store.append_state_snapshot(state)
@@ -148,8 +152,9 @@ class WorkflowRunner(Generic[StateType]):
                     parent=WorkflowParentContext(
                         span_id=span_id,
                         workflow_definition=self.workflow.__class__,
-                        parent=self._parent_context
-                    )
+                        parent=self._parent_context,
+                        type="WORKFLOW",
+                    ),
                 ),
             )
         )
@@ -185,7 +190,11 @@ class WorkflowRunner(Generic[StateType]):
                         instance=None,
                         outputs_class=node.Outputs,
                     )
-                    node.state.meta.node_outputs[output_descriptor] = streaming_output_queues[output.name]
+                    node.state.meta.node_outputs[output_descriptor] = (
+                        streaming_output_queues[output.name]
+                    )
+                    initiated_output: BaseOutput = BaseOutput(name=output.name)
+                    initiated_ports = initiated_output > ports
                     self._work_item_event_queue.put(
                         WorkItemEvent(
                             node=node,
@@ -194,14 +203,14 @@ class WorkflowRunner(Generic[StateType]):
                                 span_id=span_id,
                                 body=NodeExecutionStreamingBody(
                                     node_definition=node.__class__,
-                                    output=BaseOutput(name=output.name),
-                                    invoked_ports=invoked_ports,
+                                    output=initiated_output,
+                                    invoked_ports=initiated_ports,
                                 ),
                                 parent=WorkflowParentContext(
                                     span_id=span_id,
                                     workflow_definition=self.workflow.__class__,
-                                    parent=self._parent_context
-                                )
+                                    parent=self._parent_context,
+                                ),
                             ),
                         )
                     )
@@ -230,7 +239,7 @@ class WorkflowRunner(Generic[StateType]):
                                         span_id=span_id,
                                         workflow_definition=self.workflow.__class__,
                                         parent=self._parent_context,
-                                    )
+                                    ),
                                 ),
                             )
                         )
@@ -254,7 +263,7 @@ class WorkflowRunner(Generic[StateType]):
                                         span_id=span_id,
                                         workflow_definition=self.workflow.__class__,
                                         parent=self._parent_context,
-                                    )
+                                    ),
                                 ),
                             )
                         )
@@ -268,7 +277,9 @@ 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)
+            node.state.meta.node_execution_cache.fulfill_node_execution(
+                node.__class__, span_id
+            )
 
             self._work_item_event_queue.put(
                 WorkItemEvent(
@@ -285,7 +296,7 @@ class WorkflowRunner(Generic[StateType]):
                             span_id=span_id,
                             workflow_definition=self.workflow.__class__,
                             parent=self._parent_context,
-                        )
+                        ),
                     ),
                 )
             )
@@ -304,12 +315,14 @@ class WorkflowRunner(Generic[StateType]):
                             span_id=span_id,
                             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__}")
+            logger.exception(
+                f"An unexpected error occurred while running node {node.__class__.__name__}"
+            )
 
             self._work_item_event_queue.put(
                 WorkItemEvent(
@@ -327,15 +340,17 @@ class WorkflowRunner(Generic[StateType]):
                         parent=WorkflowParentContext(
                             span_id=span_id,
                             workflow_definition=self.workflow.__class__,
-                            parent=self._parent_context
-                        )
+                            parent=self._parent_context,
+                        ),
                     ),
                 )
             )
 
         logger.debug(f"Finished running node: {node.__class__.__name__}")
 
-    def _handle_invoked_ports(self, state: StateType, ports: Optional[Iterable[Port]]) -> None:
+    def _handle_invoked_ports(
+        self, state: StateType, ports: Optional[Iterable[Port]]
+    ) -> None:
         if not ports:
             return
 
@@ -350,7 +365,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,18 +380,27 @@ 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)
+            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]:
+    def _handle_work_item_event(
+        self, work_item_event: WorkItemEvent[StateType]
+    ) -> Optional[VellumError]:
         node = work_item_event.node
         event = work_item_event.event
 
@@ -389,7 +416,10 @@ class WorkflowRunner(Generic[StateType]):
                 node_output_descriptor = workflow_output_descriptor.instance
                 if not isinstance(node_output_descriptor, OutputReference):
                     continue
-                if node_output_descriptor.outputs_class != event.node_definition.Outputs:
+                if (
+                    node_output_descriptor.outputs_class
+                    != event.node_definition.Outputs
+                ):
                     continue
                 if node_output_descriptor.name != event.output.name:
                     continue
@@ -427,7 +457,9 @@ class WorkflowRunner(Generic[StateType]):
             parent=self._parent_context,
         )
 
-    def _stream_workflow_event(self, output: BaseOutput) -> WorkflowExecutionStreamingEvent:
+    def _stream_workflow_event(
+        self, output: BaseOutput
+    ) -> WorkflowExecutionStreamingEvent:
         return WorkflowExecutionStreamingEvent(
             trace_id=self._initial_state.meta.trace_id,
             span_id=self._initial_state.meta.span_id,
@@ -435,10 +467,12 @@ 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:
+    def _fulfill_workflow_event(
+        self, outputs: OutputsType
+    ) -> WorkflowExecutionFulfilledEvent:
         return WorkflowExecutionFulfilledEvent(
             trace_id=self._initial_state.meta.trace_id,
             span_id=self._initial_state.meta.span_id,
@@ -449,7 +483,9 @@ class WorkflowRunner(Generic[StateType]):
             parent=self._parent_context,
         )
 
-    def _reject_workflow_event(self, error: VellumError) -> WorkflowExecutionRejectedEvent:
+    def _reject_workflow_event(
+        self, error: VellumError
+    ) -> WorkflowExecutionRejectedEvent:
         return WorkflowExecutionRejectedEvent(
             trace_id=self._initial_state.meta.trace_id,
             span_id=self._initial_state.meta.span_id,
@@ -469,7 +505,9 @@ class WorkflowRunner(Generic[StateType]):
             ),
         )
 
-    def _pause_workflow_event(self, external_inputs: Iterable[ExternalInputReference]) -> WorkflowExecutionPausedEvent:
+    def _pause_workflow_event(
+        self, external_inputs: Iterable[ExternalInputReference]
+    ) -> WorkflowExecutionPausedEvent:
         return WorkflowExecutionPausedEvent(
             trace_id=self._initial_state.meta.trace_id,
             span_id=self._initial_state.meta.span_id,
@@ -486,7 +524,10 @@ class WorkflowRunner(Generic[StateType]):
         if not self._entrypoints:
             self._workflow_event_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
@@ -505,7 +546,9 @@ class WorkflowRunner(Generic[StateType]):
                 logger.exception(err_message)
                 self._workflow_event_queue.put(
                     self._reject_workflow_event(
-                        VellumError(code=VellumErrorCode.INTERNAL_ERROR, message=err_message),
+                        VellumError(
+                            code=VellumErrorCode.INTERNAL_ERROR, message=err_message
+                        ),
                     )
                 )
                 return
@@ -561,7 +604,11 @@ 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))
 
@@ -586,24 +633,41 @@ class WorkflowRunner(Generic[StateType]):
         self._cancel_signal.wait()
         self._workflow_event_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:
+        if (
+            self._initial_state.meta.is_terminated
+            or self._initial_state.meta.is_terminated is None
+        ):
             event = self._initiate_workflow_event()
         else:
             event = self._resume_workflow_event()
@@ -612,7 +676,10 @@ class WorkflowRunner(Generic[StateType]):
         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():
@@ -623,7 +690,7 @@ class WorkflowRunner(Generic[StateType]):
 
             yield self._emit_event(event)
 
-            if is_terminal_event(event):
+            if self._is_terminal_event(event):
                 break
 
         try:
@@ -632,7 +699,7 @@ class WorkflowRunner(Generic[StateType]):
         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,

vellum/workflows/state/base.py

@@ -5,15 +5,15 @@ from datetime import datetime
 from queue import Queue
 from threading import Lock
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Optional, Sequence, Set, Tuple, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Type, cast
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, field_serializer
 from pydantic_core import core_schema
 
 from vellum.core.pydantic_utilities import UniversalBaseModel
-
 from vellum.workflows.constants import UNDEF
+from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
 from vellum.workflows.types.generics import StateType
@@ -71,58 +71,92 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
 
 
 class NodeExecutionCache:
-    _node_execution_ids: Dict[Type["BaseNode"], Stack[UUID]]
+    _node_executions_fulfilled: Dict[Type["BaseNode"], Stack[UUID]]
     _node_executions_initiated: Dict[Type["BaseNode"], Set[UUID]]
-    _dependencies_invoked: Dict[str, Set[str]]
+    _node_executions_queued: Dict[Type["BaseNode"], List[UUID]]
+    _dependencies_invoked: Dict[UUID, Set[Type["BaseNode"]]]
 
     def __init__(
         self,
         dependencies_invoked: Optional[Dict[str, Sequence[str]]] = None,
-        node_execution_ids: Optional[Dict[str, Sequence[str]]] = None,
+        node_executions_fulfilled: Optional[Dict[str, Sequence[str]]] = None,
         node_executions_initiated: Optional[Dict[str, Sequence[str]]] = None,
+        node_executions_queued: Optional[Dict[str, Sequence[str]]] = None,
     ) -> None:
         self._dependencies_invoked = defaultdict(set)
-        self._node_execution_ids = defaultdict(Stack[UUID])
+        self._node_executions_fulfilled = defaultdict(Stack[UUID])
         self._node_executions_initiated = defaultdict(set)
+        self._node_executions_queued = defaultdict(list)
 
-        for node, dependencies in (dependencies_invoked or {}).items():
-            self._dependencies_invoked[node].update(dependencies)
+        for execution_id, dependencies in (dependencies_invoked or {}).items():
+            self._dependencies_invoked[UUID(execution_id)] = {get_class_by_qualname(dep) for dep in dependencies}
 
-        for node, execution_ids in (node_execution_ids or {}).items():
+        for node, execution_ids in (node_executions_fulfilled or {}).items():
             node_class = get_class_by_qualname(node)
-            self._node_execution_ids[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
+            self._node_executions_fulfilled[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
 
         for node, execution_ids in (node_executions_initiated or {}).items():
             node_class = get_class_by_qualname(node)
             self._node_executions_initiated[node_class].update({UUID(execution_id) for execution_id in execution_ids})
 
-    @property
-    def dependencies_invoked(self) -> Dict[str, Set[str]]:
-        return self._dependencies_invoked
+        for node, execution_ids in (node_executions_queued or {}).items():
+            node_class = get_class_by_qualname(node)
+            self._node_executions_queued[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
 
-    def is_node_initiated(self, node: Type["BaseNode"]) -> bool:
-        return node in self._node_executions_initiated and len(self._node_executions_initiated[node]) > 0
+    def _invoke_dependency(
+        self,
+        execution_id: UUID,
+        node: Type["BaseNode"],
+        dependency: Type["BaseNode"],
+        dependencies: Set["Type[BaseNode]"],
+    ) -> None:
+        self._dependencies_invoked[execution_id].add(dependency)
+        if all(dep in self._dependencies_invoked[execution_id] for dep in dependencies):
+            self._node_executions_queued[node].remove(execution_id)
+
+    def queue_node_execution(
+        self, node: Type["BaseNode"], dependencies: Set["Type[BaseNode]"], invoked_by: Optional[Edge] = None
+    ) -> UUID:
+        execution_id = uuid4()
+        if not invoked_by:
+            return execution_id
+
+        source_node = invoked_by.from_port.node_class
+        for queued_node_execution_id in self._node_executions_queued[node]:
+            if source_node not in self._dependencies_invoked[queued_node_execution_id]:
+                self._invoke_dependency(queued_node_execution_id, node, source_node, dependencies)
+                return queued_node_execution_id
+
+        self._node_executions_queued[node].append(execution_id)
+        self._invoke_dependency(execution_id, node, source_node, dependencies)
+        return execution_id
+
+    def is_node_execution_initiated(self, node: Type["BaseNode"], execution_id: UUID) -> bool:
+        return execution_id in self._node_executions_initiated[node]
 
     def initiate_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
         self._node_executions_initiated[node].add(execution_id)
 
     def fulfill_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
-        self._node_executions_initiated[node].remove(execution_id)
-        self._node_execution_ids[node].push(execution_id)
+        self._node_executions_fulfilled[node].push(execution_id)
 
     def get_execution_count(self, node: Type["BaseNode"]) -> int:
-        return self._node_execution_ids[node].size()
+        return self._node_executions_fulfilled[node].size()
 
     def dump(self) -> Dict[str, Any]:
         return {
             "dependencies_invoked": {
-                node: list(dependencies) for node, dependencies in self._dependencies_invoked.items()
+                str(execution_id): [str(dep) for dep in dependencies]
+                for execution_id, dependencies in self._dependencies_invoked.items()
             },
             "node_executions_initiated": {
                 str(node): list(execution_ids) for node, execution_ids in self._node_executions_initiated.items()
             },
-            "node_execution_ids": {
-                str(node): execution_ids.dump() for node, execution_ids in self._node_execution_ids.items()
+            "node_executions_fulfilled": {
+                str(node): execution_ids.dump() for node, execution_ids in self._node_executions_fulfilled.items()
+            },
+            "node_executions_queued": {
+                str(node): execution_ids for node, execution_ids in self._node_executions_queued.items()
             },
         }
 

vellum/workflows/state/context.py

@@ -1,14 +1,24 @@
 from functools import cached_property
-from typing import Optional
+from queue import Queue
+from typing import TYPE_CHECKING, Optional
 
 from vellum import Vellum
-
+from vellum.workflows.events.types import ParentContext
 from vellum.workflows.vellum_client import create_vellum_client
 
+if TYPE_CHECKING:
+    from vellum.workflows.events.workflow import WorkflowEvent
+
 
 class WorkflowContext:
-    def __init__(self, _vellum_client: Optional[Vellum] = None):
+    def __init__(
+        self,
+        _vellum_client: Optional[Vellum] = None,
+        _parent_context: Optional[ParentContext] = None,
+    ):
         self._vellum_client = _vellum_client
+        self._parent_context = _parent_context
+        self._event_queue: Optional[Queue["WorkflowEvent"]] = None
 
     @cached_property
     def vellum_client(self) -> Vellum:
@@ -16,3 +26,16 @@ class WorkflowContext:
             return self._vellum_client
 
         return create_vellum_client()
+
+    @cached_property
+    def parent_context(self) -> Optional[ParentContext]:
+        if self._parent_context:
+            return self._parent_context
+        return None
+
+    def _emit_subworkflow_event(self, event: "WorkflowEvent") -> None:
+        if self._event_queue:
+            self._event_queue.put(event)
+
+    def _register_event_queue(self, event_queue: Queue["WorkflowEvent"]) -> None:
+        self._event_queue = event_queue

vellum/workflows/types/core.py

@@ -83,7 +83,7 @@ EntityInputsInterface = Dict[
 class MergeBehavior(Enum):
     AWAIT_ALL = "AWAIT_ALL"
     AWAIT_ANY = "AWAIT_ANY"
-
+    AWAIT_ATTRIBUTES = "AWAIT_ATTRIBUTES"
 
 class ConditionType(Enum):
     IF = "IF"