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

vellum/workflows/sandbox.py

@@ -1,25 +1,47 @@
-from typing import Generic, Sequence
+from typing import Any, Dict, Generic, Optional, Sequence, Type, Union
 
 import dotenv
 
 from vellum.workflows.events.workflow import WorkflowEventStream
 from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.inputs.dataset_row import DatasetRow
 from vellum.workflows.logging import load_logger
+from vellum.workflows.triggers.base import BaseTrigger
 from vellum.workflows.types.generics import WorkflowType
 from vellum.workflows.workflows.event_filters import root_workflow_event_filter
 
 
 class WorkflowSandboxRunner(Generic[WorkflowType]):
-    def __init__(self, workflow: WorkflowType, inputs: Sequence[BaseInputs]):
-        if not inputs:
-            raise ValueError("Inputs are required to have at least one defined inputs")
-
-        self._workflow = workflow
-        self._inputs = inputs
-
+    def __init__(
+        self,
+        workflow: WorkflowType,
+        inputs: Optional[Sequence[BaseInputs]] = None,  # DEPRECATED - remove in v2.0.0
+        dataset: Optional[Sequence[Union[BaseInputs, DatasetRow]]] = None,
+    ):
         dotenv.load_dotenv()
         self._logger = load_logger()
 
+        if dataset is not None and inputs is not None:
+            raise ValueError(
+                "Cannot specify both 'dataset' and 'inputs' parameters. " "Use 'dataset' as 'inputs' is deprecated."
+            )
+
+        if dataset is not None:
+            actual_inputs = dataset
+        elif inputs is not None:
+            self._logger.warning(
+                "The 'inputs' parameter is deprecated and will be removed in v2.0.0. " "Please use 'dataset' instead."
+            )
+            actual_inputs = inputs
+        else:
+            raise ValueError("Either 'dataset' or 'inputs' parameter is required")
+
+        if not actual_inputs:
+            raise ValueError("Dataset/inputs are required to have at least one defined input")
+
+        self._workflow = workflow
+        self._inputs = actual_inputs
+
     def run(self, index: int = 0):
         if index < 0:
             self._logger.warning("Index is less than 0, running first input")
@@ -30,9 +52,33 @@ class WorkflowSandboxRunner(Generic[WorkflowType]):
 
         selected_inputs = self._inputs[index]
 
+        raw_inputs: Union[BaseInputs, Dict[str, Any]]
+        trigger_class: Optional[Type[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
+        else:
+            raw_inputs = selected_inputs
+
+        inputs_for_stream: BaseInputs
+        if isinstance(raw_inputs, dict):
+            inputs_class = type(self._workflow).get_inputs_class()
+            inputs_for_stream = inputs_class(**raw_inputs)
+        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()
+
         events = self._workflow.stream(
-            inputs=selected_inputs,
+            inputs=inputs_for_stream,
             event_filter=root_workflow_event_filter,
+            trigger=trigger_instance,
+            node_output_mocks=node_output_mocks,
         )
 
         self._process_events(events)

vellum/workflows/state/base.py

@@ -7,17 +7,37 @@ import logging
 from queue import Queue
 from threading import Lock
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    SupportsIndex,
+    Tuple,
+    Type,
+    Union,
+    cast,
+)
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, field_serializer, field_validator
 from pydantic_core import core_schema
 
-from vellum.core.pydantic_utilities import UniversalBaseModel
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.constants import undefined
 from vellum.workflows.inputs.base import BaseInputs
-from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
+from vellum.workflows.references import (
+    ExternalInputReference,
+    OutputReference,
+    StateValueReference,
+    TriggerAttributeReference,
+)
+from vellum.workflows.state.delta import AppendStateDelta, SetStateDelta, StateDelta
 from vellum.workflows.types.definition import CodeResourceDefinition, serialize_type_encoder_with_id
 from vellum.workflows.types.generics import StateType, import_workflow_class, is_workflow_class
 from vellum.workflows.types.stack import Stack
@@ -31,16 +51,12 @@ logger = logging.getLogger(__name__)
 
 
 class _Snapshottable:
-    _snapshot_callback: Callable[[], None]
+    _snapshot_callback: Callable[[Optional[StateDelta]], None]
+    _path: str
 
-    def __deepcopy__(self, memo: Any) -> "_Snapshottable":
-        return deepcopy_with_exclusions(
-            self,
-            memo=memo,
-            exclusions={
-                "_snapshot_callback": self._snapshot_callback,
-            },
-        )
+    def __bind__(self, path: str, snapshot_callback: Callable[[Optional[StateDelta]], None]) -> None:
+        self._snapshot_callback = snapshot_callback
+        self._path = path
 
 
 @dataclass_transform(kw_only_default=True)
@@ -71,10 +87,44 @@ class _BaseStateMeta(type):
 class _SnapshottableDict(dict, _Snapshottable):
     def __setitem__(self, key: Any, value: Any) -> None:
         super().__setitem__(key, value)
-        self._snapshot_callback()
+        self._snapshot_callback(SetStateDelta(name=f"{self._path}.{key}", delta=value))
+
+    def __deepcopy__(self, memo: Any) -> "_SnapshottableDict":
+        y: dict = {}
+        memo[id(self)] = y
+        for key, value in self.items():
+            y[deepcopy(key, memo)] = deepcopy(value, memo)
+
+        y = _SnapshottableDict(y)
+        y.__bind__(self._path, self._snapshot_callback)
+        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 isinstance(index, int):
+            self._snapshot_callback(SetStateDelta(name=f"{self._path}.{index}", delta=value))
+
+    def append(self, value: Any) -> None:
+        super().append(value)
+        self._snapshot_callback(AppendStateDelta(name=self._path, delta=value))
 
+    def __deepcopy__(self, memo: Any) -> "_SnapshottableList":
+        y: list = []
+        memo[id(self)] = y
+        append = y.append
+        for a in self:
+            append(deepcopy(a, memo))
 
-def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> Any:
+        y = _SnapshottableList(y)
+        y.__bind__(self._path, self._snapshot_callback)
+        memo[id(self)] = y
+        return y
+
+
+def _make_snapshottable(path: str, value: Any, snapshot_callback: Callable[[Optional[StateDelta]], None]) -> Any:
     """
     Edits any value to make it snapshottable on edit. Made as a separate function from `BaseState` to
     avoid namespace conflicts with subclasses.
@@ -84,9 +134,14 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
 
     if isinstance(value, dict):
         snapshottable_dict = _SnapshottableDict(value)
-        snapshottable_dict._snapshot_callback = snapshot_callback
+        snapshottable_dict.__bind__(path, snapshot_callback)
         return snapshottable_dict
 
+    if isinstance(value, list):
+        snapshottable_list = _SnapshottableList(value)
+        snapshottable_list.__bind__(path, snapshot_callback)
+        return snapshottable_list
+
     return value
 
 
@@ -105,8 +160,6 @@ class NodeExecutionCache:
 
     # Derived fields - no need to serialize
     __node_execution_lookup__: NodeExecutionLookup  # execution_id -> node_class
-    __node_to_fork_id__: Dict[Type["BaseNode"], UUID]  # node_class -> fork_id
-    __all_fork_ids__: Set[UUID]
 
     def __init__(self) -> None:
         self._dependencies_invoked = defaultdict(set)
@@ -114,8 +167,6 @@ class NodeExecutionCache:
         self._node_executions_initiated = defaultdict(set)
         self._node_executions_queued = defaultdict(list)
         self.__node_execution_lookup__ = {}
-        self.__node_to_fork_id__ = {}
-        self.__all_fork_ids__ = set()
 
     @classmethod
     def deserialize(cls, raw_data: dict, nodes: Dict[Union[str, UUID], Type["BaseNode"]]):
@@ -225,6 +276,10 @@ class NodeExecutionCache:
             },
         }
 
+    def __vellum_encode__(self) -> Dict[str, Any]:
+        """Return a JSON-serializable representation of this cache."""
+        return self.dump()
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler
@@ -255,16 +310,18 @@ class StateMeta(UniversalBaseModel):
     workflow_inputs: BaseInputs = field(default_factory=BaseInputs)
     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[[], None]] = field(init=False, default=None)
+    __snapshot_callback__: Optional[Callable[[Optional[StateDelta]], None]] = field(init=False, default=None)
 
     def model_post_init(self, context: Any) -> None:
         self.__snapshot_callback__ = None
 
-    def add_snapshot_callback(self, callback: Callable[[], None]) -> None:
-        self.node_outputs = _make_snapshottable(self.node_outputs, callback)
-        self.external_inputs = _make_snapshottable(self.external_inputs, callback)
+    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)
         self.__snapshot_callback__ = callback
 
     def __setattr__(self, name: str, value: Any) -> None:
@@ -274,7 +331,7 @@ class StateMeta(UniversalBaseModel):
 
         super().__setattr__(name, value)
         if callable(self.__snapshot_callback__):
-            self.__snapshot_callback__()
+            self.__snapshot_callback__(SetStateDelta(name=f"meta.{name}", delta=value))
 
     @field_serializer("workflow_definition")
     def serialize_workflow_definition(self, workflow_definition: Type["BaseWorkflow"], _info: Any) -> Dict[str, Any]:
@@ -299,6 +356,12 @@ class StateMeta(UniversalBaseModel):
     def serialize_node_outputs(self, node_outputs: Dict[OutputReference, Any], _info: Any) -> Dict[str, Any]:
         return {str(descriptor.id): value for descriptor, value in node_outputs.items()}
 
+    @field_serializer("trigger_attributes")
+    def serialize_trigger_attributes(
+        self, trigger_attributes: Dict[TriggerAttributeReference, Any], _info: Any
+    ) -> Dict[str, Any]:
+        return {str(descriptor.id): value for descriptor, value in trigger_attributes.items()}
+
     @field_validator("node_outputs", mode="before")
     @classmethod
     def deserialize_node_outputs(cls, node_outputs: Any, info: ValidationInfo):
@@ -333,6 +396,41 @@ class StateMeta(UniversalBaseModel):
 
         return node_outputs
 
+    @field_validator("trigger_attributes", mode="before")
+    @classmethod
+    def deserialize_trigger_attributes(cls, trigger_attributes: Any, info: ValidationInfo):
+        if isinstance(trigger_attributes, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return trigger_attributes
+
+            trigger_attribute_map: Dict[Union[str, UUID], TriggerAttributeReference] = {}
+            for subgraph in workflow_definition.get_subgraphs():
+                for trigger_class in subgraph.triggers:
+                    for attribute in trigger_class:
+                        trigger_attribute_map[str(attribute)] = attribute
+                        attr_id = getattr(attribute, "id", None)
+                        if isinstance(attr_id, UUID):
+                            trigger_attribute_map[attr_id] = attribute
+                            trigger_attribute_map[str(attr_id)] = attribute
+
+            deserialized_attributes: Dict[TriggerAttributeReference, Any] = {}
+            for key, value in trigger_attributes.items():
+                reference: Optional[TriggerAttributeReference]
+                if is_valid_uuid(key):
+                    reference = trigger_attribute_map.get(UUID(key))
+                else:
+                    reference = trigger_attribute_map.get(key)
+
+                if not reference:
+                    continue
+
+                deserialized_attributes[reference] = value
+
+            return deserialized_attributes
+
+        return trigger_attributes
+
     @field_validator("node_execution_cache", mode="before")
     @classmethod
     def deserialize_node_execution_cache(cls, node_execution_cache: Any, info: ValidationInfo):
@@ -396,14 +494,24 @@ class StateMeta(UniversalBaseModel):
         if not memo:
             memo = {}
 
+        node_output_keys = list(self.node_outputs.keys())
         new_node_outputs = {
-            descriptor: value if isinstance(value, Queue) else deepcopy(value, memo)
-            for descriptor, value in self.node_outputs.items()
+            descriptor: (
+                self.node_outputs[descriptor]
+                if isinstance(self.node_outputs[descriptor], Queue)
+                else deepcopy(self.node_outputs[descriptor], memo)
+            )
+            for descriptor in node_output_keys
         }
 
+        external_input_keys = list(self.external_inputs.keys())
         new_external_inputs = {
-            descriptor: value if isinstance(value, Queue) else deepcopy(value, memo)
-            for descriptor, value in self.external_inputs.items()
+            descriptor: (
+                self.external_inputs[descriptor]
+                if isinstance(self.external_inputs[descriptor], Queue)
+                else deepcopy(self.external_inputs[descriptor], memo)
+            )
+            for descriptor in external_input_keys
         }
 
         memo[id(self.node_outputs)] = new_node_outputs
@@ -432,11 +540,15 @@ class BaseState(metaclass=_BaseStateMeta):
 
     __lock__: Lock = field(init=False)
     __is_quiet__: bool = field(init=False)
-    __snapshot_callback__: Callable[["BaseState"], None] = field(init=False)
+    __is_atomic__: bool = field(init=False)
+    __snapshot_callback__: Callable[["BaseState", List[StateDelta]], None] = field(init=False)
+    __deltas__: List[StateDelta] = field(init=False)
 
     def __init__(self, meta: Optional[StateMeta] = None, **kwargs: Any) -> None:
         self.__is_quiet__ = True
-        self.__snapshot_callback__ = lambda state: None
+        self.__is_atomic__ = False
+        self.__snapshot_callback__ = lambda state, deltas: None
+        self.__deltas__ = []
         self.__lock__ = Lock()
 
         self.meta = meta or StateMeta()
@@ -446,7 +558,7 @@ class BaseState(metaclass=_BaseStateMeta):
         for name, value in self.__class__.__dict__.items():
             if not name.startswith("_") and name != "meta":
                 # Bypass __is_quiet__ instead of `setattr`
-                snapshottable_value = _make_snapshottable(value, self.__snapshot__)
+                snapshottable_value = _make_snapshottable(name, value, self.__snapshot__)
                 super().__setattr__(name, snapshottable_value)
 
         for name, value in kwargs.items():
@@ -498,15 +610,15 @@ class BaseState(metaclass=_BaseStateMeta):
     def __getitem__(self, key: str) -> Any:
         return self.__dict__[key]
 
-    def __setattr__(self, name: str, value: Any) -> None:
+    def __setattr__(self, name: str, delta: Any) -> None:
         if name.startswith("_"):
-            super().__setattr__(name, value)
+            super().__setattr__(name, delta)
             return
 
-        snapshottable_value = _make_snapshottable(value, self.__snapshot__)
+        snapshottable_value = _make_snapshottable(name, delta, self.__snapshot__)
         super().__setattr__(name, snapshottable_value)
         self.meta.updated_ts = datetime_now()
-        self.__snapshot__()
+        self.__snapshot__(SetStateDelta(name=name, delta=delta))
 
     def __add__(self, other: StateType) -> StateType:
         """
@@ -535,7 +647,7 @@ class BaseState(metaclass=_BaseStateMeta):
 
         return cast(StateType, latest_state)
 
-    def __snapshot__(self) -> None:
+    def __snapshot__(self, delta: Optional[StateDelta] = None) -> None:
         """
         Snapshots the current state to the workflow emitter. The invoked callback is overridden by the
         workflow runner.
@@ -543,11 +655,19 @@ class BaseState(metaclass=_BaseStateMeta):
         if self.__is_quiet__:
             return
 
+        if delta:
+            self.__deltas__.append(delta)
+
+        if self.__is_atomic__:
+            return
+
         try:
-            self.__snapshot_callback__(deepcopy(self))
+            self.__snapshot_callback__(deepcopy(self), self.__deltas__)
         except Exception:
             logger.exception("Failed to snapshot Workflow state.")
 
+        self.__deltas__.clear()
+
     @contextmanager
     def __quiet__(self):
         prev = self.__is_quiet__
@@ -559,14 +679,18 @@ class BaseState(metaclass=_BaseStateMeta):
 
     @contextmanager
     def __atomic__(self):
-        prev = self.__is_quiet__
-        self.__is_quiet__ = True
+        prev = self.__is_atomic__
+        self.__is_atomic__ = True
         try:
             yield
         finally:
-            self.__is_quiet__ = prev
+            self.__is_atomic__ = prev
             self.__snapshot__()
 
+    def __vellum_encode__(self) -> Dict[str, Any]:
+        """Return a JSON-serializable representation of this state."""
+        return dict(self)
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler