vellum-ai 1.4.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

vellum/workflows/runner/runner.py

@@ -42,6 +42,7 @@ from vellum.workflows.events import (
     WorkflowExecutionStreamingEvent,
 )
 from vellum.workflows.events.node import (
+    NodeEvent,
     NodeExecutionFulfilledBody,
     NodeExecutionInitiatedBody,
     NodeExecutionRejectedBody,
@@ -212,6 +213,10 @@ class WorkflowRunner(Generic[StateType]):
             descriptor for descriptor in self.workflow.Outputs if isinstance(descriptor.instance, StateValueReference)
         ]
 
+        self._background_thread: Optional[Thread] = None
+        self._cancel_thread: Optional[Thread] = None
+        self._stream_thread: Optional[Thread] = None
+
     def _snapshot_state(self, state: StateType, deltas: List[StateDelta]) -> StateType:
         self._workflow_event_inner_queue.put(
             WorkflowExecutionSnapshottedEvent(
@@ -259,17 +264,36 @@ class WorkflowRunner(Generic[StateType]):
         return event
 
     def _run_work_item(self, node: BaseNode[StateType], span_id: UUID) -> None:
+        for event in self.run_node(node, span_id):
+            self._workflow_event_inner_queue.put(event)
+
+    def run_node(
+        self,
+        node: "BaseNode[StateType]",
+        span_id: UUID,
+    ) -> Generator[NodeEvent, None, None]:
+        """
+        Execute a single node and yield workflow events.
+
+        Args:
+            node: The node instance to execute
+            span_id: Unique identifier for this node execution
+
+        Yields:
+            NodeExecutionEvent: Events emitted during node execution (initiated, streaming, fulfilled, rejected)
+        """
         execution = get_execution_context()
-        self._workflow_event_inner_queue.put(
-            NodeExecutionInitiatedEvent(
-                trace_id=execution.trace_id,
-                span_id=span_id,
-                body=NodeExecutionInitiatedBody(
-                    node_definition=node.__class__,
-                    inputs=node._inputs,
-                ),
-                parent=execution.parent_context,
-            )
+
+        node_output_mocks_map = self.workflow.context.node_output_mocks_map
+
+        yield NodeExecutionInitiatedEvent(
+            trace_id=execution.trace_id,
+            span_id=span_id,
+            body=NodeExecutionInitiatedBody(
+                node_definition=node.__class__,
+                inputs=node._inputs,
+            ),
+            parent=execution.parent_context,
         )
 
         logger.debug(f"Started running node: {node.__class__.__name__}")
@@ -282,7 +306,7 @@ class WorkflowRunner(Generic[StateType]):
             )
             node_run_response: NodeRunResponse
             was_mocked: Optional[bool] = None
-            mock_candidates = self.workflow.context.node_output_mocks_map.get(node.Outputs) or []
+            mock_candidates = node_output_mocks_map.get(node.Outputs) or []
             for mock_candidate in mock_candidates:
                 if mock_candidate.when_condition.resolve(node.state):
                     node_run_response = mock_candidate.then_outputs
@@ -312,8 +336,9 @@ class WorkflowRunner(Generic[StateType]):
                 streaming_output_queues: Dict[str, Queue] = {}
                 outputs = node.Outputs()
 
-                def initiate_node_streaming_output(output: BaseOutput) -> None:
-                    execution = get_execution_context()
+                def initiate_node_streaming_output(
+                    output: BaseOutput,
+                ) -> Generator[NodeExecutionStreamingEvent, None, None]:
                     streaming_output_queues[output.name] = Queue()
                     output_descriptor = OutputReference(
                         name=output.name,
@@ -325,57 +350,51 @@ class WorkflowRunner(Generic[StateType]):
                         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._workflow_event_inner_queue.put(
-                        NodeExecutionStreamingEvent(
-                            trace_id=execution.trace_id,
-                            span_id=span_id,
-                            body=NodeExecutionStreamingBody(
-                                node_definition=node.__class__,
-                                output=initiated_output,
-                                invoked_ports=initiated_ports,
-                            ),
-                            parent=execution.parent_context,
+                    yield NodeExecutionStreamingEvent(
+                        trace_id=execution.trace_id,
+                        span_id=span_id,
+                        body=NodeExecutionStreamingBody(
+                            node_definition=node.__class__,
+                            output=initiated_output,
+                            invoked_ports=initiated_ports,
                         ),
+                        parent=execution.parent_context,
                     )
 
                 with execution_context(parent_context=updated_parent_context, trace_id=execution.trace_id):
                     for output in node_run_response:
                         invoked_ports = output > ports
                         if output.is_initiated:
-                            initiate_node_streaming_output(output)
+                            yield from initiate_node_streaming_output(output)
                         elif output.is_streaming:
                             if output.name not in streaming_output_queues:
-                                initiate_node_streaming_output(output)
+                                yield from initiate_node_streaming_output(output)
 
                             streaming_output_queues[output.name].put(output.delta)
-                            self._workflow_event_inner_queue.put(
-                                NodeExecutionStreamingEvent(
-                                    trace_id=execution.trace_id,
-                                    span_id=span_id,
-                                    body=NodeExecutionStreamingBody(
-                                        node_definition=node.__class__,
-                                        output=output,
-                                        invoked_ports=invoked_ports,
-                                    ),
-                                    parent=execution.parent_context,
+                            yield NodeExecutionStreamingEvent(
+                                trace_id=execution.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=output,
+                                    invoked_ports=invoked_ports,
                                 ),
+                                parent=execution.parent_context,
                             )
                         elif output.is_fulfilled:
                             if output.name in streaming_output_queues:
                                 streaming_output_queues[output.name].put(undefined)
 
                             setattr(outputs, output.name, output.value)
-                            self._workflow_event_inner_queue.put(
-                                NodeExecutionStreamingEvent(
-                                    trace_id=execution.trace_id,
-                                    span_id=span_id,
-                                    body=NodeExecutionStreamingBody(
-                                        node_definition=node.__class__,
-                                        output=output,
-                                        invoked_ports=invoked_ports,
-                                    ),
-                                    parent=execution.parent_context,
-                                )
+                            yield NodeExecutionStreamingEvent(
+                                trace_id=execution.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=output,
+                                    invoked_ports=invoked_ports,
+                                ),
+                                parent=execution.parent_context,
                             )
 
             node.state.meta.node_execution_cache.fulfill_node_execution(node.__class__, span_id)
@@ -390,66 +409,57 @@ class WorkflowRunner(Generic[StateType]):
                     node.state.meta.node_outputs[descriptor] = output_value
 
             invoked_ports = ports(outputs, node.state)
-            self._workflow_event_inner_queue.put(
-                NodeExecutionFulfilledEvent(
-                    trace_id=execution.trace_id,
-                    span_id=span_id,
-                    body=NodeExecutionFulfilledBody(
-                        node_definition=node.__class__,
-                        outputs=outputs,
-                        invoked_ports=invoked_ports,
-                        mocked=was_mocked,
-                    ),
-                    parent=execution.parent_context,
-                )
+            yield NodeExecutionFulfilledEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionFulfilledBody(
+                    node_definition=node.__class__,
+                    outputs=outputs,
+                    invoked_ports=invoked_ports,
+                    mocked=was_mocked,
+                ),
+                parent=execution.parent_context,
             )
         except NodeException as e:
             logger.info(e)
             captured_stacktrace = traceback.format_exc()
 
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
-                    trace_id=execution.trace_id,
-                    span_id=span_id,
-                    body=NodeExecutionRejectedBody(
-                        node_definition=node.__class__,
-                        error=e.error,
-                        stacktrace=captured_stacktrace,
-                    ),
-                    parent=execution.parent_context,
-                )
+            yield NodeExecutionRejectedEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionRejectedBody(
+                    node_definition=node.__class__,
+                    error=e.error,
+                    stacktrace=captured_stacktrace,
+                ),
+                parent=execution.parent_context,
             )
         except WorkflowInitializationException as e:
             logger.info(e)
             captured_stacktrace = traceback.format_exc()
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
-                    trace_id=execution.trace_id,
-                    span_id=span_id,
-                    body=NodeExecutionRejectedBody(
-                        node_definition=node.__class__,
-                        error=e.error,
-                        stacktrace=captured_stacktrace,
-                    ),
-                    parent=execution.parent_context,
-                )
+            yield NodeExecutionRejectedEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionRejectedBody(
+                    node_definition=node.__class__,
+                    error=e.error,
+                    stacktrace=captured_stacktrace,
+                ),
+                parent=execution.parent_context,
             )
         except InvalidExpressionException as e:
             logger.info(e)
             captured_stacktrace = traceback.format_exc()
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
-                    trace_id=execution.trace_id,
-                    span_id=span_id,
-                    body=NodeExecutionRejectedBody(
-                        node_definition=node.__class__,
-                        error=e.error,
-                        stacktrace=captured_stacktrace,
-                    ),
-                    parent=execution.parent_context,
-                )
+            yield NodeExecutionRejectedEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionRejectedBody(
+                    node_definition=node.__class__,
+                    error=e.error,
+                    stacktrace=captured_stacktrace,
+                ),
+                parent=execution.parent_context,
             )
-
         except Exception as e:
             error_message = self._parse_error_message(e)
             if error_message is None:
@@ -459,19 +469,17 @@ class WorkflowRunner(Generic[StateType]):
             else:
                 error_code = WorkflowErrorCode.NODE_EXECUTION
 
-            self._workflow_event_inner_queue.put(
-                NodeExecutionRejectedEvent(
-                    trace_id=execution.trace_id,
-                    span_id=span_id,
-                    body=NodeExecutionRejectedBody(
-                        node_definition=node.__class__,
-                        error=WorkflowError(
-                            message=error_message,
-                            code=error_code,
-                        ),
+            yield NodeExecutionRejectedEvent(
+                trace_id=execution.trace_id,
+                span_id=span_id,
+                body=NodeExecutionRejectedBody(
+                    node_definition=node.__class__,
+                    error=WorkflowError(
+                        message=error_message,
+                        code=error_code,
                     ),
-                    parent=execution.parent_context,
                 ),
+                parent=execution.parent_context,
             )
 
         logger.debug(f"Finished running node: {node.__class__.__name__}")
@@ -911,20 +919,20 @@ class WorkflowRunner(Generic[StateType]):
         return False
 
     def _generate_events(self) -> Generator[WorkflowEvent, None, None]:
-        background_thread = Thread(
+        self._background_thread = Thread(
             target=self._run_background_thread,
             name=f"{self.workflow.__class__.__name__}.background_thread",
         )
-        background_thread.start()
+        self._background_thread.start()
 
         cancel_thread_kill_switch = ThreadingEvent()
         if self._cancel_signal:
-            cancel_thread = Thread(
+            self._cancel_thread = Thread(
                 target=self._run_cancel_thread,
                 name=f"{self.workflow.__class__.__name__}.cancel_thread",
                 kwargs={"kill_switch": cancel_thread_kill_switch},
             )
-            cancel_thread.start()
+            self._cancel_thread.start()
 
         event: WorkflowEvent
         if self._is_resuming:
@@ -935,13 +943,13 @@ class WorkflowRunner(Generic[StateType]):
         yield self._emit_event(event)
 
         # The extra level of indirection prevents the runner from waiting on the caller to consume the event stream
-        stream_thread = Thread(
+        self._stream_thread = Thread(
             target=self._stream,
             name=f"{self.workflow.__class__.__name__}.stream_thread",
         )
-        stream_thread.start()
+        self._stream_thread.start()
 
-        while stream_thread.is_alive():
+        while self._stream_thread.is_alive():
             try:
                 event = self._workflow_event_outer_queue.get(timeout=0.1)
             except Empty:
@@ -971,3 +979,17 @@ class WorkflowRunner(Generic[StateType]):
 
     def stream(self) -> WorkflowEventStream:
         return WorkflowEventGenerator(self._generate_events(), self._initial_state.meta.span_id)
+
+    def join(self) -> None:
+        """
+        Wait for all background threads to complete.
+        This ensures all pending work is finished before the runner terminates.
+        """
+        if self._stream_thread and self._stream_thread.is_alive():
+            self._stream_thread.join()
+
+        if self._background_thread and self._background_thread.is_alive():
+            self._background_thread.join()
+
+        if self._cancel_thread and self._cancel_thread.is_alive():
+            self._cancel_thread.join()

vellum/workflows/utils/functions.py

@@ -22,7 +22,7 @@ from vellum.workflows.utils.vellum_variables import vellum_variable_type_to_open
 if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
 
-type_map = {
+type_map: dict[Any, str] = {
     str: "string",
     int: "integer",
     float: "number",
@@ -32,8 +32,13 @@ type_map = {
     None: "null",
     type(None): "null",
     inspect._empty: "null",
+    "None": "null",
 }
 
+for k, v in list(type_map.items()):
+    if isinstance(k, type):
+        type_map[k.__name__] = v
+
 
 def compile_annotation(annotation: Optional[Any], defs: dict[str, Any]) -> dict:
     if annotation is None:

vellum/workflows/utils/tests/test_functions.py

@@ -733,3 +733,43 @@ def test_compile_function_definition__tuples(annotation, expected_schema):
     compiled_function = compile_function_definition(my_function)
     assert isinstance(compiled_function.parameters, dict)
     assert compiled_function.parameters["properties"]["a"] == expected_schema
+
+
+def test_compile_function_definition__string_annotations_with_future_imports():
+    """Test that string annotations work with __future__ import annotations."""
+    # This simulates what happens when using `from __future__ import annotations`
+    # where type annotations become string literals at runtime
+
+    def my_function_with_string_annotations(
+        a: "str",
+        b: "int",
+        c: "float",
+        d: "bool",
+        e: "list",
+        f: "dict",
+        g: "None",
+    ):
+        """Function with string type annotations."""
+        pass
+
+    # WHEN compiling the function
+    compiled_function = compile_function_definition(my_function_with_string_annotations)
+
+    # THEN it should return the compiled function definition with proper types
+    assert compiled_function == FunctionDefinition(
+        name="my_function_with_string_annotations",
+        description="Function with string type annotations.",
+        parameters={
+            "type": "object",
+            "properties": {
+                "a": {"type": "string"},
+                "b": {"type": "integer"},
+                "c": {"type": "number"},
+                "d": {"type": "boolean"},
+                "e": {"type": "array"},
+                "f": {"type": "object"},
+                "g": {"type": "null"},
+            },
+            "required": ["a", "b", "c", "d", "e", "f", "g"],
+        },
+    )

vellum/workflows/workflows/base.py

@@ -31,6 +31,7 @@ from vellum.workflows.edges import Edge
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
 from vellum.workflows.events.node import (
+    NodeEvent,
     NodeExecutionFulfilledBody,
     NodeExecutionFulfilledEvent,
     NodeExecutionInitiatedBody,
@@ -252,6 +253,7 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
         self.resolvers = resolvers or (self.resolvers if hasattr(self, "resolvers") else [])
         self._store = store or Store()
         self._execution_context = self._context.execution_context
+        self._current_runner: Optional[WorkflowRunner] = None
 
         # Register context with all emitters
         for emitter in self.emitters:
@@ -412,7 +414,7 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
             subworkflows or nodes that utilizes threads.
         """
 
-        events = WorkflowRunner(
+        runner = WorkflowRunner(
             self,
             inputs=inputs,
             state=state,
@@ -423,7 +425,9 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
             node_output_mocks=node_output_mocks,
             max_concurrency=max_concurrency,
             init_execution_context=self._execution_context,
-        ).stream()
+        )
+        self._current_runner = runner
+        events = runner.stream()
         first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
         last_event = None
         for event in events:
@@ -531,7 +535,7 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
         """
 
         should_yield = event_filter or workflow_event_filter
-        runner_stream = WorkflowRunner(
+        runner = WorkflowRunner(
             self,
             inputs=inputs,
             state=state,
@@ -542,7 +546,9 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
             node_output_mocks=node_output_mocks,
             max_concurrency=max_concurrency,
             init_execution_context=self._execution_context,
-        ).stream()
+        )
+        self._current_runner = runner
+        runner_stream = runner.stream()
 
         def _generate_filtered_events() -> Generator[BaseWorkflow.WorkflowEvent, None, None]:
             for event in runner_stream:
@@ -559,6 +565,11 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
         # https://app.shortcut.com/vellum/story/4327
         pass
 
+    def run_node(self, node: Type[BaseNode]) -> Generator[NodeEvent, None, None]:
+        runner = WorkflowRunner(self)
+        span_id = uuid4()
+        return runner.run_node(node=node(state=self.get_default_state(), context=self._context), span_id=span_id)
+
     @classmethod
     @lru_cache
     def _get_parameterized_classes(
@@ -689,9 +700,12 @@ class BaseWorkflow(Generic[InputsType, StateType], BaseExecutable, metaclass=_Ba
 
     def join(self) -> None:
         """
-        Wait for all emitters to complete their background work.
+        Wait for all emitters and runner to complete their background work.
         This ensures all pending events are processed before the workflow terminates.
         """
+        if self._current_runner:
+            self._current_runner.join()
+
         for emitter in self.emitters:
             emitter.join()
 

vellum/workflows/workflows/tests/test_base_workflow.py

@@ -3,6 +3,7 @@ import logging
 from uuid import UUID, uuid4
 
 from vellum.workflows.edges.edge import Edge
+from vellum.workflows.events.node import NodeExecutionFulfilledEvent, NodeExecutionInitiatedEvent
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
@@ -684,3 +685,56 @@ def test_base_workflow__deserialize_state_with_invalid_workflow_definition(raw_w
 
     # AND the workflow definition should be BaseWorkflow
     assert state.meta.workflow_definition == BaseWorkflow
+
+
+def test_base_workflow__join_calls_runner_join():
+    """
+    Test that BaseWorkflow.join() calls runner.join() when runner exists.
+    """
+
+    # GIVEN a test workflow
+    class TestWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        pass
+
+    workflow = TestWorkflow()
+
+    # WHEN we run the workflow to create a runner
+    workflow.run()
+
+    workflow.join()
+
+    # THEN the runner should have been joined (verified by no hanging threads)
+    assert workflow._current_runner is not None
+
+
+def test_base_workflow__run_node_emits_correct_events():
+    """Test that WorkflowRunner.run_node method emits the expected events."""
+
+    class TestInputs(BaseInputs):
+        pass
+
+    class TestState(BaseState):
+        pass
+
+    class TestNode(BaseNode[TestState]):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+        def run(self) -> "TestNode.Outputs":
+            return self.Outputs(result="test_output")
+
+    class TestWorkflow(BaseWorkflow[TestInputs, TestState]):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            result: str
+
+    workflow = TestWorkflow()
+
+    events = list(workflow.run_node(node=TestNode))
+
+    assert len(events) == 2
+    assert isinstance(events[0], NodeExecutionInitiatedEvent)
+    assert isinstance(events[1], NodeExecutionFulfilledEvent)
+    assert events[0].span_id == events[1].span_id
+    assert events[1].body.outputs.result == "test_output"

{vellum_ai-1.4.2.dist-info → vellum_ai-1.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.4.2
+Version: 1.5.0
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0