vellum-ai 0.13.28__py3-none-any.whl → 0.14.0__py3-none-any.whl

vellum/workflows/outputs/base.py

@@ -4,7 +4,7 @@ from typing_extensions import dataclass_transform
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 
-from vellum.workflows.constants import UNDEF
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
@@ -19,17 +19,17 @@ _Accumulated = TypeVar("_Accumulated")
 
 
 class BaseOutput(Generic[_Delta, _Accumulated]):
-    _value: Union[_Accumulated, Type[UNDEF]]
-    _delta: Union[_Delta, Type[UNDEF]]
+    _value: Union[_Accumulated, Type[undefined]]
+    _delta: Union[_Delta, Type[undefined]]
     _name: str
 
     def __init__(
         self,
         name: str,
-        value: Union[_Accumulated, Type[UNDEF]] = UNDEF,
-        delta: Union[_Delta, Type[UNDEF]] = UNDEF,
+        value: Union[_Accumulated, Type[undefined]] = undefined,
+        delta: Union[_Delta, Type[undefined]] = undefined,
     ) -> None:
-        if value is not UNDEF and delta is not UNDEF:
+        if value is not undefined and delta is not undefined:
             raise ValueError("Cannot set both value and delta")
 
         self._name = name
@@ -37,24 +37,24 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
         self._delta = delta
 
     @property
-    def delta(self) -> Union[_Delta, Type[UNDEF]]:
+    def delta(self) -> Union[_Delta, Type[undefined]]:
         return self._delta
 
     @property
-    def value(self) -> Union[_Accumulated, Type[UNDEF]]:
+    def value(self) -> Union[_Accumulated, Type[undefined]]:
         return self._value
 
     @property
     def is_initiated(self) -> bool:
-        return self._delta is UNDEF and self._value is UNDEF
+        return self._delta is undefined and self._value is undefined
 
     @property
     def is_streaming(self) -> bool:
-        return self._delta is not UNDEF and self._value is UNDEF
+        return self._delta is not undefined and self._value is undefined
 
     @property
     def is_fulfilled(self) -> bool:
-        return self._delta is UNDEF and self._value is not UNDEF
+        return self._delta is undefined and self._value is not undefined
 
     @property
     def name(self) -> str:
@@ -71,18 +71,18 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
             "name": self.name,
         }
 
-        if self.value is not UNDEF:
+        if self.value is not undefined:
             data["value"] = self.value
 
-        if self.delta is not UNDEF:
+        if self.delta is not undefined:
             data["delta"] = self.delta
 
         return data
 
     def __repr__(self) -> str:
-        if self.value is not UNDEF:
+        if self.value is not undefined:
             return f"{self.__class__.__name__}({self.name}={self.value})"
-        elif self.delta is not UNDEF:
+        elif self.delta is not undefined:
             return f"{self.__class__.__name__}({self.name}={self.delta})"
         else:
             return f"{self.__class__.__name__}(name='{self.name}')"
@@ -144,7 +144,7 @@ class _BaseOutputsMeta(type):
             # We first try to resolve the instance that this class attribute name is mapped to. If it's not found,
             # we iterate through its inheritance hierarchy to find the first base class that has this attribute
             # and use its mapping.
-            instance = vars(cls).get(name, UNDEF)
+            instance = vars(cls).get(name, undefined)
             if not instance:
                 for base in cls.__mro__[1:]:
                     if hasattr(base, name):
@@ -204,7 +204,9 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
         if not isinstance(other, dict):
             return super().__eq__(other)
 
-        outputs = {name: value for name, value in vars(self).items() if not name.startswith("_") and value is not UNDEF}
+        outputs = {
+            name: value for name, value in vars(self).items() if not name.startswith("_") and value is not undefined
+        }
         return outputs == other
 
     def __repr__(self) -> str:
@@ -213,9 +215,9 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
 
     def __iter__(self) -> Iterator[Tuple[OutputReference, Any]]:
         for output_descriptor in self.__class__:
-            output_value = getattr(self, output_descriptor.name, UNDEF)
+            output_value = getattr(self, output_descriptor.name, undefined)
             if isinstance(output_value, BaseDescriptor):
-                output_value = UNDEF
+                output_value = undefined
 
             yield (output_descriptor, output_value)
 

vellum/workflows/references/external_input.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Generic, Optional, Tuple, Type, TypeVar,
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 
-from vellum.workflows.constants import UNDEF
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
@@ -34,7 +34,7 @@ class ExternalInputReference(BaseDescriptor[_InputType], Generic[_InputType]):
 
     def resolve(self, state: "BaseState") -> _InputType:
         external_input = state.meta.external_inputs.get(self)
-        if external_input is not UNDEF:
+        if external_input is not undefined:
             return cast(_InputType, external_input)
 
         if state.meta.parent:

vellum/workflows/references/lazy.py

@@ -2,7 +2,7 @@ import ast
 import inspect
 from typing import TYPE_CHECKING, Callable, Generic, TypeVar, Union, get_args
 
-from vellum.workflows.constants import UNDEF
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 
 if TYPE_CHECKING:
@@ -37,7 +37,7 @@ class LazyReference(BaseDescriptor[_T], Generic[_T]):
 
             # Fix typing surrounding the return value of node outputs/output descriptors
             # https://app.shortcut.com/vellum/story/4783
-            return UNDEF  # type: ignore[return-value]
+            return undefined  # type: ignore[return-value]
 
         return resolve_value(self._get(), state)
 

vellum/workflows/references/output.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Generator, Generic, Optional, Tuple, Type
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 
-from vellum.workflows.constants import UNDEF
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 
 if TYPE_CHECKING:
@@ -32,13 +32,13 @@ class OutputReference(BaseDescriptor[_OutputType], Generic[_OutputType]):
         return self._outputs_class
 
     def resolve(self, state: "BaseState") -> _OutputType:
-        node_output = state.meta.node_outputs.get(self, UNDEF)
+        node_output = state.meta.node_outputs.get(self, undefined)
         if isinstance(node_output, Queue):
             # Fix typing surrounding the return value of node outputs
             # https://app.shortcut.com/vellum/story/4783
             return self._as_generator(node_output)  # type: ignore[return-value]
 
-        if node_output is not UNDEF:
+        if node_output is not undefined:
             return cast(_OutputType, node_output)
 
         if state.meta.parent:
@@ -46,13 +46,13 @@ class OutputReference(BaseDescriptor[_OutputType], Generic[_OutputType]):
 
         # Fix typing surrounding the return value of node outputs
         # https://app.shortcut.com/vellum/story/4783
-        return cast(Type[UNDEF], node_output)  # type: ignore[return-value]
+        return cast(Type[undefined], node_output)  # type: ignore[return-value]
 
-    def _as_generator(self, node_output: Queue) -> Generator[_OutputType, None, Type[UNDEF]]:
+    def _as_generator(self, node_output: Queue) -> Generator[_OutputType, None, Type[undefined]]:
         while True:
             item = node_output.get()
-            if item is UNDEF:
-                return UNDEF
+            if item is undefined:
+                return undefined
             yield cast(_OutputType, item)
 
     def __eq__(self, other: object) -> bool:

vellum/workflows/runner/runner.py

@@ -6,7 +6,7 @@ from threading import Event as ThreadingEvent, Thread
 from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Tuple, Type, Union
 
-from vellum.workflows.constants import UNDEF
+from vellum.workflows.constants import undefined
 from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges.edge import Edge
@@ -267,7 +267,7 @@ class WorkflowRunner(Generic[StateType]):
                             )
                         elif output.is_fulfilled:
                             if output.name in streaming_output_queues:
-                                streaming_output_queues[output.name].put(UNDEF)
+                                streaming_output_queues[output.name].put(undefined)
 
                             setattr(outputs, output.name, output.value)
                             self._workflow_event_inner_queue.put(
@@ -286,7 +286,7 @@ class WorkflowRunner(Generic[StateType]):
             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 output_value is undefined:
                     if descriptor in node.state.meta.node_outputs:
                         del node.state.meta.node_outputs[descriptor]
                     continue
@@ -386,8 +386,8 @@ class WorkflowRunner(Generic[StateType]):
                 if not isinstance(descriptor, ExternalInputReference):
                     continue
 
-                if state.meta.external_inputs.get(descriptor, UNDEF) is UNDEF:
-                    state.meta.external_inputs[descriptor] = UNDEF
+                if state.meta.external_inputs.get(descriptor, undefined) is undefined:
+                    state.meta.external_inputs[descriptor] = undefined
                     return
 
             all_deps = self._dependencies[node_class]
@@ -577,7 +577,7 @@ class WorkflowRunner(Generic[StateType]):
         unresolved_external_inputs = {
             descriptor
             for descriptor, node_input_value in final_state.meta.external_inputs.items()
-            if node_input_value is UNDEF
+            if node_input_value is undefined
         }
         if unresolved_external_inputs:
             self._workflow_event_outer_queue.put(
@@ -616,19 +616,21 @@ class WorkflowRunner(Generic[StateType]):
                 for emitter in self.workflow.emitters:
                     emitter.emit_event(item)
 
-    def _run_cancel_thread(self) -> None:
+    def _run_cancel_thread(self, kill_switch: ThreadingEvent) -> None:
         if not self._cancel_signal:
             return
 
-        self._cancel_signal.wait()
-        self._workflow_event_outer_queue.put(
-            self._reject_workflow_event(
-                WorkflowError(
-                    code=WorkflowErrorCode.WORKFLOW_CANCELLED,
-                    message="Workflow run cancelled",
+        while not kill_switch.wait(timeout=0.1):
+            if self._cancel_signal.is_set():
+                self._workflow_event_outer_queue.put(
+                    self._reject_workflow_event(
+                        WorkflowError(
+                            code=WorkflowErrorCode.WORKFLOW_CANCELLED,
+                            message="Workflow run cancelled",
+                        )
+                    )
                 )
-            )
-        )
+                return
 
     def _is_terminal_event(self, event: WorkflowEvent) -> bool:
         if (
@@ -646,10 +648,12 @@ class WorkflowRunner(Generic[StateType]):
         )
         background_thread.start()
 
+        cancel_thread_kill_switch = ThreadingEvent()
         if self._cancel_signal:
             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()
 
@@ -694,3 +698,4 @@ class WorkflowRunner(Generic[StateType]):
             )
 
         self._background_thread_queue.put(None)
+        cancel_thread_kill_switch.set()

vellum/workflows/state/base.py

@@ -12,7 +12,7 @@ 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.constants import undefined
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
@@ -347,7 +347,7 @@ class BaseState(metaclass=_BaseStateMeta):
             if key.startswith("_"):
                 continue
 
-            if getattr(latest_state, key, UNDEF) == UNDEF:
+            if getattr(latest_state, key, undefined) == undefined:
                 setattr(latest_state, key, value)
 
         return cast(StateType, latest_state)

vellum/workflows/state/tests/test_state.py

@@ -23,6 +23,9 @@ class MockState(BaseState):
 
 
 class MockNode(BaseNode):
+    class ExternalInputs(BaseNode.ExternalInputs):
+        message: str
+
     class Outputs(BaseOutputs):
         baz: str
 
@@ -70,11 +73,7 @@ def test_state_snapshot__external_input_edit():
     assert snapshot_count[id(state)] == 0
 
     # WHEN we add an external input to state
-    class MockExternalInputs(BaseNode.ExternalInputs):
-        message: str
-
-    # WHEN we edit external inputs dictionary
-    state.meta.external_inputs[MockExternalInputs.message] = "hello"
+    state.meta.external_inputs[MockNode.ExternalInputs.message] = "hello"
 
     # THEN the snapshot is emitted
     assert snapshot_count[id(state)] == 1
@@ -137,19 +136,16 @@ def test_state_deepcopy__with_external_input_updates():
     state = MockState(foo="bar")
 
     # AND we add an external input to state
-    class MockExternalInputs(BaseNode.ExternalInputs):
-        message: str
-
-    state.meta.external_inputs[MockExternalInputs.message] = "hello"
+    state.meta.external_inputs[MockNode.ExternalInputs.message] = "hello"
 
     # AND we deepcopy the state
     deepcopied_state = deepcopy(state)
 
     # AND we update the original state
-    state.meta.external_inputs[MockExternalInputs.message] = "world"
+    state.meta.external_inputs[MockNode.ExternalInputs.message] = "world"
 
     # THEN the copied state is not updated
-    assert deepcopied_state.meta.external_inputs[MockExternalInputs.message] == "hello"
+    assert deepcopied_state.meta.external_inputs[MockNode.ExternalInputs.message] == "hello"
 
     # AND the original state has had the correct number of snapshots
     assert snapshot_count[id(state)] == 2

vellum/workflows/workflows/base.py

@@ -80,6 +80,26 @@ class _BaseWorkflowMeta(type):
         if "graph" not in dct:
             dct["graph"] = set()
 
+        if "Outputs" in dct:
+            outputs_class = dct["Outputs"]
+
+            if not any(issubclass(base, BaseOutputs) for base in outputs_class.__bases__):
+                parent_outputs_class = next(
+                    (base.Outputs for base in bases if hasattr(base, "Outputs")),
+                    BaseOutputs,  # Default to BaseOutputs only if no parent has Outputs
+                )
+
+                filtered_bases = tuple(base for base in outputs_class.__bases__ if base is not object)
+
+                new_dct = {key: value for key, value in outputs_class.__dict__.items() if not key.startswith("__")}
+                new_dct["__module__"] = dct["__module__"]
+
+                dct["Outputs"] = type(
+                    f"{name}.Outputs",
+                    (parent_outputs_class,) + filtered_bases,
+                    new_dct,
+                )
+
         cls = super().__new__(mcs, name, bases, dct)
         workflow_class = cast(Type["BaseWorkflow"], cls)
         workflow_class.__id__ = uuid4_from_hash(workflow_class.__qualname__)

vellum/workflows/workflows/tests/test_base_workflow.py

@@ -0,0 +1,80 @@
+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
+from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.workflows.base import BaseWorkflow
+
+
+def test_base_workflow__inherit_base_outputs():
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+        def run(self):
+            return self.Outputs(foo="bar")
+
+    class MyWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = MyNode
+
+        class Outputs:
+            output = MyNode.Outputs.foo
+
+    # TEST that the Outputs class is a subclass of BaseOutputs
+    assert issubclass(MyWorkflow.Outputs, BaseOutputs)
+
+    # TEST that the Outputs class does not inherit from object
+    assert object not in MyWorkflow.Outputs.__bases__
+
+    workflow = MyWorkflow()
+    terminal_event = workflow.run()
+
+    # TEST that the Outputs class has the correct attributes
+    assert hasattr(MyWorkflow.Outputs, "output")
+
+    # TEST that the outputs should be correct
+    assert terminal_event.name == "workflow.execution.fulfilled", terminal_event
+    assert terminal_event.outputs == {"output": "bar"}
+
+
+def test_subworkflow__inherit_base_outputs():
+    class StartNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+        def run(self):
+            return self.Outputs(foo="bar")
+
+    class SubWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = StartNode
+
+        class Outputs:
+            output = StartNode.Outputs.foo
+
+    class SubworkflowNode(InlineSubworkflowNode):
+        subworkflow = SubWorkflow
+
+    class MainWorkflow(BaseWorkflow[BaseInputs, BaseState]):
+        graph = SubworkflowNode
+
+        class Outputs:
+            output = SubworkflowNode.Outputs.output
+
+    # TEST that the Outputs classes are subclasses of BaseOutputs
+    assert issubclass(MainWorkflow.Outputs, BaseOutputs)
+    assert issubclass(SubWorkflow.Outputs, BaseOutputs)
+
+    # TEST that the Outputs classes do not inherit from object
+    assert object not in MainWorkflow.Outputs.__bases__
+    assert object not in SubWorkflow.Outputs.__bases__
+
+    # TEST execution
+    workflow = MainWorkflow()
+    terminal_event = workflow.run()
+
+    # TEST that the Outputs class has the correct attributes
+    assert hasattr(MainWorkflow.Outputs, "output")
+
+    # TEST that the outputs are correct
+    assert terminal_event.name == "workflow.execution.fulfilled", terminal_event
+    assert terminal_event.outputs == {"output": "bar"}

{vellum_ai-0.13.28.dist-info → vellum_ai-0.14.0.dist-info}/METADATA

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