vellum-ai 0.14.37__py3-none-any.whl → 0.14.39__py3-none-any.whl

vellum/workflows/state/base.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import field
 from datetime import datetime
@@ -6,7 +7,7 @@ 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, Sequence, Set, Tuple, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, cast
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, field_serializer, field_validator
@@ -16,19 +17,14 @@ from vellum.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.constants import undefined
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
-from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references import ExternalInputReference, OutputReference, StateValueReference
-from vellum.workflows.types.generics import StateType
+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
-from vellum.workflows.types.utils import (
-    datetime_now,
-    deepcopy_with_exclusions,
-    get_class_attr_names,
-    get_class_by_qualname,
-    infer_types,
-)
+from vellum.workflows.types.utils import datetime_now, deepcopy_with_exclusions, get_class_attr_names, infer_types
 
 if TYPE_CHECKING:
+    from vellum.workflows import BaseWorkflow
     from vellum.workflows.nodes.bases import BaseNode
 
 logger = logging.getLogger(__name__)
@@ -94,38 +90,65 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
     return value
 
 
+NodeExecutionsFulfilled = Dict[Type["BaseNode"], Stack[UUID]]
+NodeExecutionsInitiated = Dict[Type["BaseNode"], Set[UUID]]
+NodeExecutionsQueued = Dict[Type["BaseNode"], List[UUID]]
+DependenciesInvoked = Dict[UUID, Set[Type["BaseNode"]]]
+
+
 class NodeExecutionCache:
-    _node_executions_fulfilled: Dict[Type["BaseNode"], Stack[UUID]]
-    _node_executions_initiated: Dict[Type["BaseNode"], Set[UUID]]
-    _node_executions_queued: Dict[Type["BaseNode"], List[UUID]]
-    _dependencies_invoked: Dict[UUID, Set[Type["BaseNode"]]]
+    _node_executions_fulfilled: NodeExecutionsFulfilled
+    _node_executions_initiated: NodeExecutionsInitiated
+    _node_executions_queued: NodeExecutionsQueued
+    _dependencies_invoked: DependenciesInvoked
 
-    def __init__(
-        self,
-        dependencies_invoked: 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:
+    def __init__(self) -> None:
         self._dependencies_invoked = defaultdict(set)
         self._node_executions_fulfilled = defaultdict(Stack[UUID])
         self._node_executions_initiated = defaultdict(set)
         self._node_executions_queued = defaultdict(list)
 
-        for execution_id, dependencies in (dependencies_invoked or {}).items():
-            self._dependencies_invoked[UUID(execution_id)] = {get_class_by_qualname(dep) for dep in dependencies}
+    @classmethod
+    def deserialize(cls, raw_data: dict, nodes: Dict[str, Type["BaseNode"]]):
+        cache = cls()
+
+        dependencies_invoked = raw_data.get("dependencies_invoked")
+        if isinstance(dependencies_invoked, dict):
+            for execution_id, dependencies in dependencies_invoked.items():
+                cache._dependencies_invoked[UUID(execution_id)] = {nodes[dep] for dep in dependencies if dep in nodes}
+
+        node_executions_fulfilled = raw_data.get("node_executions_fulfilled")
+        if isinstance(node_executions_fulfilled, dict):
+            for node, execution_ids in node_executions_fulfilled.items():
+                node_class = nodes.get(node)
+                if not node_class:
+                    continue
 
-        for node, execution_ids in (node_executions_fulfilled or {}).items():
-            node_class = get_class_by_qualname(node)
-            self._node_executions_fulfilled[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
+                cache._node_executions_fulfilled[node_class].extend(
+                    UUID(execution_id) for execution_id in execution_ids
+                )
+
+        node_executions_initiated = raw_data.get("node_executions_initiated")
+        if isinstance(node_executions_initiated, dict):
+            for node, execution_ids in node_executions_initiated.items():
+                node_class = nodes.get(node)
+                if not node_class:
+                    continue
+
+                cache._node_executions_initiated[node_class].update(
+                    {UUID(execution_id) for execution_id in execution_ids}
+                )
+
+        node_executions_queued = raw_data.get("node_executions_queued")
+        if isinstance(node_executions_queued, dict):
+            for node, execution_ids in node_executions_queued.items():
+                node_class = nodes.get(node)
+                if not node_class:
+                    continue
 
-        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})
+                cache._node_executions_queued[node_class].extend(UUID(execution_id) for execution_id in execution_ids)
 
-        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)
+        return cache
 
     def _invoke_dependency(
         self,
@@ -207,8 +230,8 @@ def default_datetime_factory() -> datetime:
 
 
 class StateMeta(UniversalBaseModel):
+    workflow_definition: Type["BaseWorkflow"] = field(default_factory=import_workflow_class)
     id: UUID = field(default_factory=uuid4_default_factory)
-    trace_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)
@@ -219,8 +242,6 @@ class StateMeta(UniversalBaseModel):
     __snapshot_callback__: Optional[Callable[[], None]] = field(init=False, default=None)
 
     def model_post_init(self, context: Any) -> None:
-        if self.parent:
-            self.trace_id = self.parent.meta.trace_id
         self.__snapshot_callback__ = None
 
     def add_snapshot_callback(self, callback: Callable[[], None]) -> None:
@@ -237,6 +258,25 @@ class StateMeta(UniversalBaseModel):
         if callable(self.__snapshot_callback__):
             self.__snapshot_callback__()
 
+    @field_serializer("workflow_definition")
+    def serialize_workflow_definition(self, workflow_definition: Type["BaseWorkflow"], _info: Any) -> Dict[str, Any]:
+        return serialize_type_encoder_with_id(workflow_definition)
+
+    @field_validator("workflow_definition", mode="before")
+    @classmethod
+    def deserialize_workflow_definition(cls, workflow_definition: Any, info: ValidationInfo):
+        if isinstance(workflow_definition, dict):
+            deserialized_workflow_definition = CodeResourceDefinition.model_validate(workflow_definition).decode()
+            if not is_workflow_class(deserialized_workflow_definition):
+                return import_workflow_class()
+
+            return deserialized_workflow_definition
+
+        if is_workflow_class(workflow_definition):
+            return workflow_definition
+
+        return import_workflow_class()
+
     @field_serializer("node_outputs")
     def serialize_node_outputs(self, node_outputs: Dict[OutputReference, Any], _info: Any) -> Dict[str, Any]:
         return {str(descriptor): value for descriptor, value in node_outputs.items()}
@@ -244,17 +284,16 @@ class StateMeta(UniversalBaseModel):
     @field_validator("node_outputs", mode="before")
     @classmethod
     def deserialize_node_outputs(cls, node_outputs: Any, info: ValidationInfo):
-        if isinstance(node_outputs, dict) and isinstance(info.context, dict):
-            raw_workflow_nodes = info.context.get("nodes")
-            workflow_node_outputs = {}
-            if isinstance(raw_workflow_nodes, list):
-                for node in raw_workflow_nodes:
-                    Outputs = getattr(node, "Outputs", None)
-                    if not isinstance(Outputs, type) or not issubclass(Outputs, BaseOutputs):
-                        continue
+        if isinstance(node_outputs, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return node_outputs
 
-                    for output in Outputs:
-                        workflow_node_outputs[str(output)] = output
+            raw_workflow_nodes = workflow_definition.get_nodes()
+            workflow_node_outputs = {}
+            for node in raw_workflow_nodes:
+                for output in node.Outputs:
+                    workflow_node_outputs[str(output)] = output
 
             node_output_keys = list(node_outputs.keys())
             deserialized_node_outputs = {}
@@ -269,12 +308,64 @@ class StateMeta(UniversalBaseModel):
 
         return node_outputs
 
+    @field_validator("node_execution_cache", mode="before")
+    @classmethod
+    def deserialize_node_execution_cache(cls, node_execution_cache: Any, info: ValidationInfo):
+        if isinstance(node_execution_cache, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return node_execution_cache
+
+            nodes_cache: Dict[str, Type["BaseNode"]] = {}
+            raw_workflow_nodes = workflow_definition.get_nodes()
+            for node in raw_workflow_nodes:
+                nodes_cache[str(node)] = node
+
+            return NodeExecutionCache.deserialize(node_execution_cache, nodes_cache)
+
+        return node_execution_cache
+
+    @field_validator("workflow_inputs", mode="before")
+    @classmethod
+    def deserialize_workflow_inputs(cls, workflow_inputs: Any, info: ValidationInfo):
+        workflow_definition = 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)
+
+        return workflow_inputs
+
     @field_serializer("external_inputs")
     def serialize_external_inputs(
         self, external_inputs: Dict[ExternalInputReference, Any], _info: Any
     ) -> Dict[str, Any]:
         return {str(descriptor): value for descriptor, value in external_inputs.items()}
 
+    @field_validator("parent", mode="before")
+    @classmethod
+    def deserialize_parent(cls, parent: Any, info: ValidationInfo):
+        if isinstance(parent, dict):
+            workflow_definition = cls._get_workflow(info)
+            if not workflow_definition:
+                return parent
+
+            parent_meta = parent.get("meta")
+            if not isinstance(parent_meta, dict):
+                return parent
+
+            parent_workflow_definition = cls.deserialize_workflow_definition(
+                parent_meta.get("workflow_definition"), info
+            )
+            if not is_workflow_class(parent_workflow_definition):
+                return parent
+
+            return parent_workflow_definition.deserialize_state(parent)
+
+        return parent
+
     def __deepcopy__(self, memo: Optional[Dict[int, Any]] = None) -> "StateMeta":
         if not memo:
             memo = {}
@@ -295,16 +386,30 @@ class StateMeta(UniversalBaseModel):
 
         return super().__deepcopy__(memo)
 
+    @classmethod
+    def _get_workflow(cls, info: ValidationInfo) -> Optional[Type["BaseWorkflow"]]:
+        if not info.context:
+            return None
+
+        if not isinstance(info.context, dict):
+            return None
+
+        workflow_definition = info.context.get("workflow_definition")
+        if not is_workflow_class(workflow_definition):
+            return None
+
+        return workflow_definition
+
 
 class BaseState(metaclass=_BaseStateMeta):
     meta: StateMeta = field(init=False)
 
     __lock__: Lock = field(init=False)
-    __is_initializing__: bool = field(init=False)
+    __is_quiet__: bool = field(init=False)
     __snapshot_callback__: Callable[["BaseState"], None] = field(init=False)
 
     def __init__(self, meta: Optional[StateMeta] = None, **kwargs: Any) -> None:
-        self.__is_initializing__ = True
+        self.__is_quiet__ = True
         self.__snapshot_callback__ = lambda state: None
         self.__lock__ = Lock()
 
@@ -314,14 +419,14 @@ class BaseState(metaclass=_BaseStateMeta):
         # Make all class attribute values snapshottable
         for name, value in self.__class__.__dict__.items():
             if not name.startswith("_") and name != "meta":
-                # Bypass __is_initializing__ instead of `setattr`
+                # Bypass __is_quiet__ instead of `setattr`
                 snapshottable_value = _make_snapshottable(value, self.__snapshot__)
                 super().__setattr__(name, snapshottable_value)
 
         for name, value in kwargs.items():
             setattr(self, name, value)
 
-        self.__is_initializing__ = False
+        self.__is_quiet__ = False
 
     def __deepcopy__(self, memo: Any) -> "BaseState":
         new_state = deepcopy_with_exclusions(
@@ -368,7 +473,7 @@ class BaseState(metaclass=_BaseStateMeta):
         return self.__dict__[key]
 
     def __setattr__(self, name: str, value: Any) -> None:
-        if name.startswith("_") or self.__is_initializing__:
+        if name.startswith("_"):
             super().__setattr__(name, value)
             return
 
@@ -409,11 +514,33 @@ class BaseState(metaclass=_BaseStateMeta):
         Snapshots the current state to the workflow emitter. The invoked callback is overridden by the
         workflow runner.
         """
+        if self.__is_quiet__:
+            return
+
         try:
             self.__snapshot_callback__(deepcopy(self))
         except Exception:
             logger.exception("Failed to snapshot Workflow state.")
 
+    @contextmanager
+    def __quiet__(self):
+        prev = self.__is_quiet__
+        self.__is_quiet__ = True
+        try:
+            yield
+        finally:
+            self.__is_quiet__ = prev
+
+    @contextmanager
+    def __atomic__(self):
+        prev = self.__is_quiet__
+        self.__is_quiet__ = True
+        try:
+            yield
+        finally:
+            self.__is_quiet__ = prev
+            self.__snapshot__()
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler

vellum/workflows/state/tests/test_state.py

@@ -1,17 +1,14 @@
 import pytest
-from collections import defaultdict
 from copy import deepcopy
 import json
 from queue import Queue
-from typing import Dict
+from typing import Dict, cast
 
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.state.encoder import DefaultStateEncoder
 
-snapshot_count: Dict[int, int] = defaultdict(int)
-
 
 @pytest.fixture()
 def mock_deepcopy(mocker):
@@ -27,9 +24,19 @@ class MockState(BaseState):
     foo: str
     nested_dict: Dict[str, int] = {}
 
-    def __snapshot__(self) -> None:
-        global snapshot_count
-        snapshot_count[id(self)] += 1
+    __snapshot_count__: int = 0
+
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.__snapshot_callback__ = lambda _: self.__mock_snapshot__()
+
+    def __mock_snapshot__(self) -> None:
+        self.__snapshot_count__ += 1
+
+    def __deepcopy__(self, memo: dict) -> "MockState":
+        new_state = cast(MockState, super().__deepcopy__(memo))
+        new_state.__snapshot_count__ = 0
+        return new_state
 
 
 class MockNode(BaseNode):
@@ -43,50 +50,50 @@ class MockNode(BaseNode):
 def test_state_snapshot__node_attribute_edit():
     # GIVEN an initial state instance
     state = MockState(foo="bar")
-    assert snapshot_count[id(state)] == 0
+    assert state.__snapshot_count__ == 0
 
     # WHEN we edit an attribute
     state.foo = "baz"
 
     # THEN the snapshot is emitted
-    assert snapshot_count[id(state)] == 1
+    assert state.__snapshot_count__ == 1
 
 
 def test_state_snapshot__node_output_edit():
     # GIVEN an initial state instance
     state = MockState(foo="bar")
-    assert snapshot_count[id(state)] == 0
+    assert state.__snapshot_count__ == 0
 
     # WHEN we add a Node Output to state
     for output in MockNode.Outputs:
         state.meta.node_outputs[output] = "hello"
 
     # THEN the snapshot is emitted
-    assert snapshot_count[id(state)] == 1
+    assert state.__snapshot_count__ == 1
 
 
 def test_state_snapshot__nested_dictionary_edit():
     # GIVEN an initial state instance
     state = MockState(foo="bar")
-    assert snapshot_count[id(state)] == 0
+    assert state.__snapshot_count__ == 0
 
     # WHEN we edit a nested dictionary
     state.nested_dict["hello"] = 1
 
     # THEN the snapshot is emitted
-    assert snapshot_count[id(state)] == 1
+    assert state.__snapshot_count__ == 1
 
 
 def test_state_snapshot__external_input_edit():
     # GIVEN an initial state instance
     state = MockState(foo="bar")
-    assert snapshot_count[id(state)] == 0
+    assert state.__snapshot_count__ == 0
 
     # WHEN we add an external input to state
     state.meta.external_inputs[MockNode.ExternalInputs.message] = "hello"
 
     # THEN the snapshot is emitted
-    assert snapshot_count[id(state)] == 1
+    assert state.__snapshot_count__ == 1
 
 
 def test_state_deepcopy():
@@ -103,7 +110,6 @@ def test_state_deepcopy():
     assert deepcopied_state.meta.node_outputs == state.meta.node_outputs
 
 
-@pytest.mark.skip(reason="https://app.shortcut.com/vellum/story/5654")
 def test_state_deepcopy__with_node_output_updates():
     # GIVEN an initial state instance
     state = MockState(foo="bar")
@@ -121,10 +127,10 @@ def test_state_deepcopy__with_node_output_updates():
     assert deepcopied_state.meta.node_outputs[MockNode.Outputs.baz] == "hello"
 
     # AND the original state has had the correct number of snapshots
-    assert snapshot_count[id(state)] == 2
+    assert state.__snapshot_count__ == 2
 
     # AND the copied state has had the correct number of snapshots
-    assert snapshot_count[id(deepcopied_state)] == 0
+    assert deepcopied_state.__snapshot_count__ == 0
 
 
 def test_state_json_serialization__with_node_output_updates():
@@ -158,10 +164,10 @@ def test_state_deepcopy__with_external_input_updates():
     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
+    assert state.__snapshot_count__ == 2
 
     # AND the copied state has had the correct number of snapshots
-    assert snapshot_count[id(deepcopied_state)] == 0
+    assert deepcopied_state.__snapshot_count__ == 0
 
 
 def test_state_json_serialization__with_queue():

vellum/workflows/types/definition.py

@@ -0,0 +1,71 @@
+import importlib
+import inspect
+from types import FrameType
+from uuid import UUID
+from typing import Annotated, Any, Dict, Optional, Union
+
+from pydantic import BeforeValidator
+
+from vellum.client.types.code_resource_definition import CodeResourceDefinition as ClientCodeResourceDefinition
+
+
+def serialize_type_encoder(obj: type) -> Dict[str, Any]:
+    return {
+        "name": obj.__name__,
+        "module": obj.__module__.split("."),
+    }
+
+
+def serialize_type_encoder_with_id(obj: Union[type, "CodeResourceDefinition"]) -> Dict[str, Any]:
+    if hasattr(obj, "__id__") and isinstance(obj, type):
+        return {
+            "id": getattr(obj, "__id__"),
+            **serialize_type_encoder(obj),
+        }
+    elif isinstance(obj, CodeResourceDefinition):
+        return obj.model_dump(mode="json")
+
+    raise AttributeError(f"The object of type '{type(obj).__name__}' must have an '__id__' attribute.")
+
+
+class CodeResourceDefinition(ClientCodeResourceDefinition):
+    id: UUID
+
+    @staticmethod
+    def encode(obj: type) -> "CodeResourceDefinition":
+        return CodeResourceDefinition(**serialize_type_encoder_with_id(obj))
+
+    def decode(self) -> Any:
+        if ".<locals>." in self.name:
+            # We are decoding a local class that should already be loaded in our stack frame. So
+            # we climb up to look for it.
+            frame = inspect.currentframe()
+            return self._resolve_local(frame)
+
+        try:
+            imported_module = importlib.import_module(".".join(self.module))
+        except ImportError:
+            return None
+
+        return getattr(imported_module, self.name, None)
+
+    def _resolve_local(self, frame: Optional[FrameType]) -> Any:
+        if not frame:
+            return None
+
+        frame_module = frame.f_globals.get("__name__")
+        if not isinstance(frame_module, str) or frame_module.split(".") != self.module:
+            return self._resolve_local(frame.f_back)
+
+        outer, inner = self.name.split(".<locals>.")
+        frame_outer = frame.f_code.co_name
+        if frame_outer != outer:
+            return self._resolve_local(frame.f_back)
+
+        return frame.f_locals.get(inner)
+
+
+VellumCodeResourceDefinition = Annotated[
+    CodeResourceDefinition,
+    BeforeValidator(lambda d: (d if type(d) is dict else serialize_type_encoder_with_id(d))),
+]

vellum/workflows/types/generics.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, TypeVar
+from functools import cache
+from typing import TYPE_CHECKING, Any, Type, TypeVar
+from typing_extensions import TypeGuard
 
 if TYPE_CHECKING:
     from vellum.workflows import BaseWorkflow
@@ -12,3 +14,34 @@ StateType = TypeVar("StateType", bound="BaseState")
 WorkflowType = TypeVar("WorkflowType", bound="BaseWorkflow")
 InputsType = TypeVar("InputsType", bound="BaseInputs")
 OutputsType = TypeVar("OutputsType", bound="BaseOutputs")
+
+
+@cache
+def _import_node_class() -> Type["BaseNode"]:
+    """
+    Helper function to help avoid circular imports.
+    """
+
+    from vellum.workflows.nodes import BaseNode
+
+    return BaseNode
+
+
+def import_workflow_class() -> Type["BaseWorkflow"]:
+    """
+    Helper function to help avoid circular imports.
+    """
+
+    from vellum.workflows.workflows import BaseWorkflow
+
+    return BaseWorkflow
+
+
+def is_node_class(obj: Any) -> TypeGuard[Type["BaseNode"]]:
+    base_node_class = _import_node_class()
+    return isinstance(obj, type) and issubclass(obj, base_node_class)
+
+
+def is_workflow_class(obj: Any) -> TypeGuard[Type["BaseWorkflow"]]:
+    base_workflow_class = import_workflow_class()
+    return isinstance(obj, type) and issubclass(obj, base_workflow_class)