vellum-ai 1.11.2__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum/workflows/runner/tests/test_runner.py

@@ -0,0 +1,170 @@
+from vellum.client.core.api_error import ApiError
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.events.node import NodeExecutionInitiatedEvent, NodeExecutionRejectedEvent
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.workflows.base import BaseWorkflow
+
+
+def test_workflow_runner__handles_400_api_error_with_integration_details():
+    """
+    Tests that WorkflowRunner handles 400 API errors with integration details
+    as INVALID_INPUTS errors.
+    """
+
+    # GIVEN a node that raises an ApiError with status_code 400 and integration details
+    class TestInputs(BaseInputs):
+        pass
+
+    class TestState(BaseState):
+        pass
+
+    class TestNode(BaseNode[TestState]):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+        def run(self) -> "TestNode.Outputs":
+            raise ApiError(
+                status_code=400,
+                body={
+                    "message": "Invalid request to integration",
+                    "integration": {
+                        "name": "test_integration",
+                        "provider": "test_provider",
+                    },
+                },
+            )
+
+    class TestWorkflow(BaseWorkflow[TestInputs, TestState]):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            result: str
+
+    workflow = TestWorkflow()
+
+    # WHEN we run the node
+    events = list(workflow.run_node(node=TestNode))
+
+    # THEN we should get a rejected event with INVALID_INPUTS error code
+    assert len(events) == 2
+    assert isinstance(events[0], NodeExecutionInitiatedEvent)
+    assert isinstance(events[1], NodeExecutionRejectedEvent)
+
+    # AND the error should have the correct code and message
+    rejected_event = events[1]
+    assert rejected_event.body.error.code == WorkflowErrorCode.INVALID_INPUTS
+    assert rejected_event.body.error.message == "Invalid request to integration"
+
+    # AND the raw_data should contain the integration details
+    assert rejected_event.body.error.raw_data == {
+        "integration": {
+            "name": "test_integration",
+            "provider": "test_provider",
+        }
+    }
+
+
+def test_workflow_runner__handles_403_api_error_with_integration_details():
+    """
+    Tests that WorkflowRunner handles 403 API errors with integration details correctly.
+    """
+
+    # GIVEN a node that raises an ApiError with status_code 403 and integration details
+    class TestInputs(BaseInputs):
+        pass
+
+    class TestState(BaseState):
+        pass
+
+    class TestNode(BaseNode[TestState]):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+        def run(self) -> "TestNode.Outputs":
+            raise ApiError(
+                status_code=403,
+                body={
+                    "message": "You must authenticate with this integration",
+                    "integration": {
+                        "name": "test_integration",
+                        "provider": "test_provider",
+                    },
+                },
+            )
+
+    class TestWorkflow(BaseWorkflow[TestInputs, TestState]):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            result: str
+
+    workflow = TestWorkflow()
+
+    # WHEN we run the node
+    events = list(workflow.run_node(node=TestNode))
+
+    # THEN we should get a rejected event with INTEGRATION_CREDENTIALS_UNAVAILABLE error code
+    assert len(events) == 2
+    assert isinstance(events[0], NodeExecutionInitiatedEvent)
+    assert isinstance(events[1], NodeExecutionRejectedEvent)
+
+    # AND the error should have the correct code and message
+    rejected_event = events[1]
+    assert rejected_event.body.error.code == WorkflowErrorCode.INTEGRATION_CREDENTIALS_UNAVAILABLE
+    assert rejected_event.body.error.message == "You must authenticate with this integration"
+
+    # AND the raw_data should contain the integration details
+    assert rejected_event.body.error.raw_data == {
+        "integration": {
+            "name": "test_integration",
+            "provider": "test_provider",
+        }
+    }
+
+
+def test_workflow_runner__handles_400_api_error_without_integration_details():
+    """
+    Tests that WorkflowRunner handles 400 API errors without integration details
+    as generic errors (not INTEGRATION_CREDENTIALS_UNAVAILABLE).
+    """
+
+    # GIVEN a node that raises an ApiError with status_code 400 but no integration details
+    class TestInputs(BaseInputs):
+        pass
+
+    class TestState(BaseState):
+        pass
+
+    class TestNode(BaseNode[TestState]):
+        class Outputs(BaseNode.Outputs):
+            result: str
+
+        def run(self) -> "TestNode.Outputs":
+            raise ApiError(
+                status_code=400,
+                body={
+                    "message": "Invalid request parameters",
+                },
+            )
+
+    class TestWorkflow(BaseWorkflow[TestInputs, TestState]):
+        graph = TestNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            result: str
+
+    workflow = TestWorkflow()
+
+    # WHEN we run the node
+    events = list(workflow.run_node(node=TestNode))
+
+    # THEN we should get a rejected event
+    assert len(events) == 2
+    assert isinstance(events[0], NodeExecutionInitiatedEvent)
+    assert isinstance(events[1], NodeExecutionRejectedEvent)
+
+    # AND the error should NOT have INTEGRATION_CREDENTIALS_UNAVAILABLE code
+    rejected_event = events[1]
+    assert rejected_event.body.error.code != WorkflowErrorCode.INTEGRATION_CREDENTIALS_UNAVAILABLE

vellum/workflows/sandbox.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Generic, Optional, Sequence, Type, Union
+from typing import Any, Dict, Generic, Optional, Sequence, Union
 
 import dotenv
 
@@ -53,12 +53,12 @@ class WorkflowSandboxRunner(Generic[WorkflowType]):
         selected_inputs = self._inputs[index]
 
         raw_inputs: Union[BaseInputs, Dict[str, Any]]
-        trigger_class: Optional[Type[BaseTrigger]] = None
+        trigger_value: Optional[BaseTrigger] = None
         node_output_mocks = None
         if isinstance(selected_inputs, DatasetRow):
             raw_inputs = selected_inputs.inputs
-            trigger_class = selected_inputs.workflow_trigger
-            node_output_mocks = selected_inputs.node_output_mocks
+            trigger_value = selected_inputs.workflow_trigger
+            node_output_mocks = selected_inputs.mocks
         else:
             raw_inputs = selected_inputs
 
@@ -69,10 +69,7 @@ class WorkflowSandboxRunner(Generic[WorkflowType]):
         else:
             inputs_for_stream = raw_inputs
 
-        trigger_instance: Optional[BaseTrigger] = None
-        if trigger_class is not None:
-            # Instantiate the trigger with the inputs
-            trigger_instance = trigger_class(**raw_inputs) if isinstance(raw_inputs, dict) else trigger_class()
+        trigger_instance: Optional[BaseTrigger] = trigger_value
 
         events = self._workflow.stream(
             inputs=inputs_for_stream,
@@ -97,3 +94,5 @@ class WorkflowSandboxRunner(Generic[WorkflowType]):
             elif event.name == "node.execution.rejected":
                 self._logger.debug(f"Error: {event.error}")
                 self._logger.error(f"Failed to run Node: {event.node_definition.__name__}")
+            elif event.name == "workflow.execution.rejected":
+                self._logger.error(f"Workflow rejected! Error: {event.error}")

vellum/workflows/state/base.py

@@ -5,7 +5,7 @@ from dataclasses import field
 from datetime import datetime
 import logging
 from queue import Queue
-from threading import Lock
+from threading import RLock
 from uuid import UUID, uuid4
 from typing import (
     TYPE_CHECKING,
@@ -25,6 +25,7 @@ from typing import (
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, field_serializer, field_validator
+from pydantic.fields import FieldInfo
 from pydantic_core import core_schema
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
@@ -53,10 +54,14 @@ logger = logging.getLogger(__name__)
 class _Snapshottable:
     _snapshot_callback: Callable[[Optional[StateDelta]], None]
     _path: str
+    _lock: Optional[RLock]
 
-    def __bind__(self, path: str, snapshot_callback: Callable[[Optional[StateDelta]], None]) -> None:
+    def __bind__(
+        self, path: str, snapshot_callback: Callable[[Optional[StateDelta]], None], lock: Optional[RLock] = None
+    ) -> None:
         self._snapshot_callback = snapshot_callback
         self._path = path
+        self._lock = lock
 
 
 @dataclass_transform(kw_only_default=True)
@@ -65,7 +70,9 @@ class _BaseStateMeta(type):
         if not name.startswith("_"):
             instance = vars(cls).get(name, undefined)
             types = infer_types(cls, name)
-            return StateValueReference(name=name, types=types, instance=instance)
+            return StateValueReference(
+                name=name, types=types, instance=instance, state_class=cast(Type["BaseState"], cls)
+            )
 
         return super().__getattribute__(name)
 
@@ -86,7 +93,11 @@ class _BaseStateMeta(type):
 
 class _SnapshottableDict(dict, _Snapshottable):
     def __setitem__(self, key: Any, value: Any) -> None:
-        super().__setitem__(key, value)
+        if self._lock:
+            with self._lock:
+                super().__setitem__(key, value)
+        else:
+            super().__setitem__(key, value)
         self._snapshot_callback(SetStateDelta(name=f"{self._path}.{key}", delta=value))
 
     def __deepcopy__(self, memo: Any) -> "_SnapshottableDict":
@@ -96,19 +107,27 @@ class _SnapshottableDict(dict, _Snapshottable):
             y[deepcopy(key, memo)] = deepcopy(value, memo)
 
         y = _SnapshottableDict(y)
-        y.__bind__(self._path, self._snapshot_callback)
+        y.__bind__(self._path, self._snapshot_callback, self._lock)
         memo[id(self)] = y
         return y
 
 
 class _SnapshottableList(list, _Snapshottable):
     def __setitem__(self, index: Union[SupportsIndex, slice], value: Any) -> None:
-        super().__setitem__(index, value)
+        if self._lock:
+            with self._lock:
+                super().__setitem__(index, value)
+        else:
+            super().__setitem__(index, value)
         if isinstance(index, int):
             self._snapshot_callback(SetStateDelta(name=f"{self._path}.{index}", delta=value))
 
     def append(self, value: Any) -> None:
-        super().append(value)
+        if self._lock:
+            with self._lock:
+                super().append(value)
+        else:
+            super().append(value)
         self._snapshot_callback(AppendStateDelta(name=self._path, delta=value))
 
     def __deepcopy__(self, memo: Any) -> "_SnapshottableList":
@@ -119,27 +138,33 @@ class _SnapshottableList(list, _Snapshottable):
             append(deepcopy(a, memo))
 
         y = _SnapshottableList(y)
-        y.__bind__(self._path, self._snapshot_callback)
+        y.__bind__(self._path, self._snapshot_callback, self._lock)
         memo[id(self)] = y
         return y
 
 
-def _make_snapshottable(path: str, value: Any, snapshot_callback: Callable[[Optional[StateDelta]], None]) -> Any:
+def _make_snapshottable(
+    path: str,
+    value: Any,
+    snapshot_callback: Callable[[Optional[StateDelta]], None],
+    lock: Optional[RLock] = None,
+) -> Any:
     """
     Edits any value to make it snapshottable on edit. Made as a separate function from `BaseState` to
     avoid namespace conflicts with subclasses.
     """
     if isinstance(value, _Snapshottable):
+        value.__bind__(path, snapshot_callback, lock)
         return value
 
     if isinstance(value, dict):
         snapshottable_dict = _SnapshottableDict(value)
-        snapshottable_dict.__bind__(path, snapshot_callback)
+        snapshottable_dict.__bind__(path, snapshot_callback, lock)
         return snapshottable_dict
 
     if isinstance(value, list):
         snapshottable_list = _SnapshottableList(value)
-        snapshottable_list.__bind__(path, snapshot_callback)
+        snapshottable_list.__bind__(path, snapshot_callback, lock)
         return snapshottable_list
 
     return value
@@ -307,10 +332,10 @@ class StateMeta(UniversalBaseModel):
     id: UUID = field(default_factory=uuid4_default_factory)
     span_id: UUID = field(default_factory=uuid4_default_factory)
     updated_ts: datetime = field(default_factory=default_datetime_factory)
-    workflow_inputs: BaseInputs = field(default_factory=BaseInputs)
+    trigger_attributes: Optional[Dict[TriggerAttributeReference, Any]] = field(default=None)
+    workflow_inputs: Optional[BaseInputs] = field(default=None)
     external_inputs: Dict[ExternalInputReference, Any] = field(default_factory=dict)
     node_outputs: Dict[OutputReference, Any] = field(default_factory=dict)
-    trigger_attributes: Dict[TriggerAttributeReference, Any] = field(default_factory=dict)
     node_execution_cache: NodeExecutionCache = field(default_factory=NodeExecutionCache)
     parent: Optional["BaseState"] = None
     __snapshot_callback__: Optional[Callable[[Optional[StateDelta]], None]] = field(init=False, default=None)
@@ -318,10 +343,27 @@ class StateMeta(UniversalBaseModel):
     def model_post_init(self, context: Any) -> None:
         self.__snapshot_callback__ = None
 
-    def add_snapshot_callback(self, callback: Callable[[Optional[StateDelta]], None]) -> None:
-        self.node_outputs = _make_snapshottable("meta.node_outputs", self.node_outputs, callback)
-        self.external_inputs = _make_snapshottable("meta.external_inputs", self.external_inputs, callback)
-        self.trigger_attributes = _make_snapshottable("meta.trigger_attributes", self.trigger_attributes, callback)
+        # Auto-populate workflow_inputs with defaults only if trigger_attributes is None
+        # (trigger_attributes being set indicates a trigger-based workflow where we don't want default inputs)
+        if self.workflow_inputs is None and self.trigger_attributes is None:
+            workflow_definition = (
+                self.workflow_definition if is_workflow_class(self.workflow_definition) else import_workflow_class()
+            )
+            try:
+                inputs_class = workflow_definition.get_inputs_class()
+                self.workflow_inputs = inputs_class()
+            except Exception:
+                self.workflow_inputs = BaseInputs()
+
+    def add_snapshot_callback(
+        self, callback: Callable[[Optional[StateDelta]], None], lock: Optional[RLock] = None
+    ) -> None:
+        self.node_outputs = _make_snapshottable("meta.node_outputs", self.node_outputs, callback, lock)
+        self.external_inputs = _make_snapshottable("meta.external_inputs", self.external_inputs, callback, lock)
+        if self.trigger_attributes is not None:
+            self.trigger_attributes = _make_snapshottable(
+                "meta.trigger_attributes", self.trigger_attributes, callback, lock
+            )
         self.__snapshot_callback__ = callback
 
     def __setattr__(self, name: str, value: Any) -> None:
@@ -358,8 +400,10 @@ class StateMeta(UniversalBaseModel):
 
     @field_serializer("trigger_attributes")
     def serialize_trigger_attributes(
-        self, trigger_attributes: Dict[TriggerAttributeReference, Any], _info: Any
+        self, trigger_attributes: Optional[Dict[TriggerAttributeReference, Any]], _info: Any
     ) -> Dict[str, Any]:
+        if trigger_attributes is None:
+            return {}
         return {str(descriptor.id): value for descriptor, value in trigger_attributes.items()}
 
     @field_validator("node_outputs", mode="before")
@@ -452,11 +496,10 @@ class StateMeta(UniversalBaseModel):
     @field_validator("workflow_inputs", mode="before")
     @classmethod
     def deserialize_workflow_inputs(cls, workflow_inputs: Any, info: ValidationInfo):
-        workflow_definition = cls._get_workflow(info)
+        context = info.context if isinstance(info.context, dict) else {}
+        workflow_definition = context.get("workflow_definition") or cls._get_workflow(info)
 
         if workflow_definition:
-            if workflow_inputs is None:
-                return workflow_definition.get_inputs_class()()
             if isinstance(workflow_inputs, dict):
                 return workflow_definition.get_inputs_class()(**workflow_inputs)
 
@@ -538,7 +581,7 @@ class StateMeta(UniversalBaseModel):
 class BaseState(metaclass=_BaseStateMeta):
     meta: StateMeta = field(init=False)
 
-    __lock__: Lock = field(init=False)
+    __lock__: RLock = field(init=False)
     __is_quiet__: bool = field(init=False)
     __is_atomic__: bool = field(init=False)
     __snapshot_callback__: Callable[["BaseState", List[StateDelta]], None] = field(init=False)
@@ -549,16 +592,23 @@ class BaseState(metaclass=_BaseStateMeta):
         self.__is_atomic__ = False
         self.__snapshot_callback__ = lambda state, deltas: None
         self.__deltas__ = []
-        self.__lock__ = Lock()
+        self.__lock__ = RLock()
 
         self.meta = meta or StateMeta()
-        self.meta.add_snapshot_callback(self.__snapshot__)
+        self.meta.add_snapshot_callback(self.__snapshot__, self.__lock__)
 
         # Make all class attribute values snapshottable
         for name, value in self.__class__.__dict__.items():
             if not name.startswith("_") and name != "meta":
+                if isinstance(value, FieldInfo):
+                    if value.default_factory is not None:
+                        value = cast(Callable[[], Any], value.default_factory)()
+                    elif hasattr(value, "default") and value.default is not ...:
+                        value = value.default
+                    else:
+                        continue
                 # Bypass __is_quiet__ instead of `setattr`
-                snapshottable_value = _make_snapshottable(name, value, self.__snapshot__)
+                snapshottable_value = _make_snapshottable(name, value, self.__snapshot__, self.__lock__)
                 super().__setattr__(name, snapshottable_value)
 
         for name, value in kwargs.items():
@@ -570,11 +620,17 @@ class BaseState(metaclass=_BaseStateMeta):
         new_state = deepcopy_with_exclusions(
             self,
             exclusions={
-                "__lock__": Lock(),
+                "__lock__": RLock(),
             },
             memo=memo,
         )
-        new_state.meta.add_snapshot_callback(new_state.__snapshot__)
+        new_state.meta.add_snapshot_callback(new_state.__snapshot__, new_state.__lock__)
+
+        for name, value in list(vars(new_state).items()):
+            if not name.startswith("_") and name != "meta":
+                rebound_value = _make_snapshottable(name, value, new_state.__snapshot__, new_state.__lock__)
+                object.__setattr__(new_state, name, rebound_value)
+
         return new_state
 
     def __repr__(self) -> str:
@@ -615,8 +671,9 @@ class BaseState(metaclass=_BaseStateMeta):
             super().__setattr__(name, delta)
             return
 
-        snapshottable_value = _make_snapshottable(name, delta, self.__snapshot__)
-        super().__setattr__(name, snapshottable_value)
+        snapshottable_value = _make_snapshottable(name, delta, self.__snapshot__, self.__lock__)
+        with self.__lock__:
+            super().__setattr__(name, snapshottable_value)
         self.meta.updated_ts = datetime_now()
         self.__snapshot__(SetStateDelta(name=name, delta=delta))
 
@@ -662,7 +719,9 @@ class BaseState(metaclass=_BaseStateMeta):
             return
 
         try:
-            self.__snapshot_callback__(deepcopy(self), self.__deltas__)
+            with self.__lock__:
+                state_copy = deepcopy(self)
+            self.__snapshot_callback__(state_copy, self.__deltas__)
         except Exception:
             logger.exception("Failed to snapshot Workflow state.")
 

vellum/workflows/state/context.py

@@ -1,13 +1,17 @@
 from dataclasses import dataclass
 from functools import cached_property
 import json
+import logging
 from queue import Queue
+import traceback
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
 
 from vellum import Vellum, __version__
+from vellum.client.types import SeverityEnum
 from vellum.workflows.context import ExecutionContext, get_execution_context, set_execution_context
-from vellum.workflows.events.types import ExternalParentContext
+from vellum.workflows.events.node import NodeExecutionLogBody, NodeExecutionLogEvent
+from vellum.workflows.events.types import ExternalParentContext, NodeParentContext
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
@@ -21,6 +25,8 @@ if TYPE_CHECKING:
     from vellum.workflows.state.base import BaseState
     from vellum.workflows.workflows.base import BaseWorkflow
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class WorkflowDeploymentMetadata:
@@ -43,6 +49,7 @@ class WorkflowContext:
         generated_files: Optional[dict[str, str]] = None,
         namespace: Optional[str] = None,
         store_class: Optional[Type[Store]] = None,
+        event_max_size: Optional[int] = None,
     ):
         self._vellum_client = vellum_client
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
@@ -50,6 +57,7 @@ class WorkflowContext:
         self._execution_context = get_execution_context()
         self._namespace = namespace
         self._store_class = store_class if store_class is not None else Store
+        self._event_max_size = event_max_size
 
         if execution_context is not None:
             self._execution_context.trace_id = execution_context.trace_id
@@ -93,6 +101,10 @@ class WorkflowContext:
     def store_class(self) -> Type[Store]:
         return self._store_class
 
+    @property
+    def event_max_size(self) -> Optional[int]:
+        return self._event_max_size
+
     @property
     def monitoring_url(self) -> Optional[str]:
         """
@@ -140,6 +152,61 @@ class WorkflowContext:
             # For custom domains, assume the same pattern: api.* -> app.*
             return api_url.replace("api.", "app.", 1)
 
+    def emit_log_event(
+        self,
+        severity: SeverityEnum,
+        message: str,
+        attributes: Optional[Dict[str, Any]] = None,
+        exc_info: Optional[bool] = None,
+    ) -> None:
+        """Emit a log event for a particular node.
+
+        This is in active development and may have breaking changes.
+        """
+        from vellum.workflows.nodes.bases import BaseNode
+
+        if self._event_queue is None:
+            return
+
+        execution_context = get_execution_context()
+        parent_context = execution_context.parent_context
+        while parent_context:
+            if isinstance(parent_context, NodeParentContext):
+                break
+            parent_context = parent_context.parent
+
+        if not isinstance(parent_context, NodeParentContext):
+            return
+
+        try:
+            node_class = parent_context.node_definition.decode()
+        except Exception:
+            logger.exception("Failed to decode node definition.")
+            return
+
+        if not isinstance(node_class, type) or not issubclass(node_class, BaseNode):
+            logger.warning("Node definition is not a subclass of BaseNode.")
+            return
+
+        if exc_info:
+            attributes = {
+                **(attributes or {}),
+                "exc_info": traceback.format_exc(),
+            }
+
+        self._event_queue.put(
+            NodeExecutionLogEvent(
+                trace_id=execution_context.trace_id,
+                span_id=parent_context.span_id,
+                body=NodeExecutionLogBody(
+                    node_definition=node_class,
+                    severity=severity,
+                    message=message,
+                    attributes=attributes,
+                ),
+            )
+        )
+
     def _emit_subworkflow_event(self, event: "WorkflowEvent") -> None:
         if self._event_queue:
             self._event_queue.put(event)
@@ -147,6 +214,9 @@ class WorkflowContext:
     def _register_event_queue(self, event_queue: Queue["WorkflowEvent"]) -> None:
         self._event_queue = event_queue
 
+    def _register_event_max_size(self, event_max_size: Optional[int]) -> None:
+        self._event_max_size = event_max_size
+
     def _register_node_output_mocks(self, node_output_mocks: MockNodeExecutionArg) -> None:
         for mock in node_output_mocks:
             if isinstance(mock, MockNodeExecution):
@@ -290,10 +360,11 @@ class WorkflowContext:
         return None
 
     @classmethod
-    def create_from(cls, context):
+    def create_from(cls, context: "WorkflowContext") -> "WorkflowContext":
         return cls(
             vellum_client=context.vellum_client,
             generated_files=context.generated_files,
             namespace=context.namespace,
             store_class=context.store_class,
+            event_max_size=context.event_max_size,
         )