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

vellum/workflows/nodes/utils.py

@@ -8,6 +8,7 @@ from typing import Any, Callable, Dict, ForwardRef, List, Optional, Type, TypeVa
 from pydantic import BaseModel, create_model
 
 from vellum.client.types.function_call import FunctionCall
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
@@ -15,7 +16,12 @@ from vellum.workflows.nodes import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.ports.port import Port
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.types.code_execution_node_wrappers import DictWrapper, ListWrapper, StringValueWrapper
+from vellum.workflows.types.code_execution_node_wrappers import (
+    DictWrapper,
+    ListWrapper,
+    StringValueWrapper,
+    clean_for_dict_wrapper,
+)
 from vellum.workflows.types.core import Json
 from vellum.workflows.types.generics import NodeType
 
@@ -123,11 +129,11 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
             # raise ValueError("Invalid JSON Array format for result_as_str")
             raise NodeException(
                 code=WorkflowErrorCode.INVALID_OUTPUTS,
-                message="Invalid JSON Array format for result_as_str",
+                message="Invalid JSON format: expected a valid JSON array",
             )
 
         if not isinstance(data, list):
-            raise ValueError(f"Expected a list of items for result_as_str, received {data.__class__.__name__}")
+            raise ValueError(f"Expected a list of items, but received {data.__class__.__name__}")
 
         inner_type = get_args(output_type)[0]
         if issubclass(inner_type, BaseModel):
@@ -135,7 +141,7 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
         else:
             return data
 
-    if output_type is Json:
+    if output_type is Json or output_type is Any or output_type is dict or get_origin(output_type) is dict:
         try:
             data = json.loads(result_as_str)
             # If we got a FunctionCallVellumValue, return just the value
@@ -145,7 +151,7 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
         except json.JSONDecodeError:
             raise NodeException(
                 code=WorkflowErrorCode.INVALID_OUTPUTS,
-                message="Invalid JSON format for result_as_str",
+                message="Invalid JSON format: unable to parse the provided data",
             )
 
     if get_origin(output_type) is Union:
@@ -171,7 +177,7 @@ def parse_type_from_str(result_as_str: str, output_type: Any) -> Any:
         except json.JSONDecodeError:
             raise NodeException(
                 code=WorkflowErrorCode.INVALID_OUTPUTS,
-                message="Invalid JSON format for result_as_str",
+                message="Invalid JSON format: unable to parse the provided data",
             )
 
     raise ValueError(f"Unsupported output type: {output_type}")
@@ -248,6 +254,19 @@ def cast_to_output_type(result: Any, output_type: Any) -> Any:
     if result is None:
         return _get_default_value(output_type)
 
+    if result is undefined:
+        return _get_default_value(output_type)
+
+    # Attempt JSON parse if type is Any
+    if output_type is Any:
+        if isinstance(result, str):
+            try:
+                return json.loads(result)
+            except (json.JSONDecodeError, TypeError):
+                # If JSON parsing fails, fall back to original result
+                pass
+        return result
+
     clean_output_type = _clean_output_type(output_type)
     DynamicModel = create_model("Output", output_type=(clean_output_type, ...))
 
@@ -261,3 +280,23 @@ def cast_to_output_type(result: Any, output_type: Any) -> Any:
             code=WorkflowErrorCode.INVALID_OUTPUTS,
             message=f"Expected an output of type '{output_type_name}', but received '{result_type_name}'",
         )
+
+
+def wrap_inputs_for_backward_compatibility(inputs: Dict[str, Any]) -> Dict[str, Any]:
+    """Wrap inputs with backward-compatible wrapper classes for legacy .value and .type support."""
+
+    def _wrap_single_value(value: Any) -> Any:
+        if isinstance(value, list):
+            return ListWrapper(
+                [
+                    (
+                        item.model_dump()
+                        if isinstance(item, BaseModel)
+                        else clean_for_dict_wrapper(item) if isinstance(item, (dict, list, str)) else item
+                    )
+                    for item in value
+                ]
+            )
+        return clean_for_dict_wrapper(value)
+
+    return {name: _wrap_single_value(value) for name, value in inputs.items()}

vellum/workflows/outputs/base.py

@@ -1,4 +1,6 @@
-from typing import TYPE_CHECKING, Any, Generic, Iterator, Set, Tuple, Type, TypeVar, Union, cast
+from dataclasses import field
+import inspect
+from typing import Any, Dict, Generic, Iterator, Set, Tuple, Type, TypeVar, Union, cast
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler
@@ -8,12 +10,11 @@ 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
+from vellum.workflows.executable import BaseExecutable
 from vellum.workflows.references.output import OutputReference
+from vellum.workflows.types.generics import is_node_instance
 from vellum.workflows.types.utils import get_class_attr_names, infer_types
 
-if TYPE_CHECKING:
-    from vellum.workflows.nodes.bases.base import BaseNode
-
 _Delta = TypeVar("_Delta")
 _Accumulated = TypeVar("_Accumulated")
 
@@ -79,6 +80,10 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
 
         return data
 
+    def __vellum_encode__(self) -> dict:
+        """Return a JSON-serializable representation of this output."""
+        return self.serialize()
+
     def __repr__(self) -> str:
         if self.value is not undefined:
             return f"{self.__class__.__name__}({self.name}={self.value})"
@@ -99,6 +104,10 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
 
 @dataclass_transform(kw_only_default=True)
 class _BaseOutputsMeta(type):
+    def __new__(cls, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        dct["__parent_class__"] = type(None)
+        return super().__new__(cls, name, bases, dct)
+
     def __eq__(cls, other: Any) -> bool:
         """
         We need to include custom eq logic to prevent infinite loops during ipython reloading.
@@ -110,17 +119,21 @@ class _BaseOutputsMeta(type):
         if not cls.__qualname__.endswith(".Outputs") or not other.__qualname__.endswith(".Outputs"):
             return super().__eq__(other)
 
-        self_outputs_class = cast(Type["BaseNode.Outputs"], cls)
-        other_outputs_class = cast(Type["BaseNode.Outputs"], other)
+        self_outputs_class = cast(Type["BaseOutputs"], cls)
+        other_outputs_class = cast(Type["BaseOutputs"], other)
 
-        if not hasattr(self_outputs_class, "_node_class") or not hasattr(other_outputs_class, "_node_class"):
+        if not hasattr(self_outputs_class, "__parent_class__") or not hasattr(other_outputs_class, "__parent_class__"):
             return super().__eq__(other)
 
-        if self_outputs_class._node_class is None or other_outputs_class._node_class is None:
+        if self_outputs_class.__parent_class__ is None or other_outputs_class.__parent_class__ is None:
             return super().__eq__(other)
 
-        return getattr(self_outputs_class._node_class, "__qualname__") == getattr(
-            other_outputs_class._node_class, "__qualname__"
+        if self_outputs_class.__parent_class__ is type(None) or other_outputs_class.__parent_class__ is type(None):
+            return super().__eq__(other)
+
+        return (
+            self_outputs_class.__parent_class__.__qualname__ == other_outputs_class.__parent_class__.__qualname__
+            and self_outputs_class.__parent_class__.__module__ == other_outputs_class.__parent_class__.__module__
         )
 
     def __setattr__(cls, name: str, value: Any) -> None:
@@ -185,6 +198,8 @@ class _BaseOutputsMeta(type):
 
 
 class BaseOutputs(metaclass=_BaseOutputsMeta):
+    __parent_class__: Type[BaseExecutable] = field(init=False)
+
     def __init__(self, **kwargs: Any) -> None:
         declared_fields = {descriptor.name for descriptor in self.__class__}
         provided_fields = set(kwargs.keys())
@@ -198,8 +213,30 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
         for name, value in kwargs.items():
             setattr(self, name, value)
 
-        if hasattr(self, "_outputs_post_init") and callable(self._outputs_post_init):
-            self._outputs_post_init(**kwargs)
+        # If climb up the to the caller's frame, and if it's a BaseNode instance, it should
+        # have a state attribute that we can use to resolve the output descriptors.
+        frame = inspect.currentframe()
+        if not frame:
+            return
+
+        caller_frame = frame.f_back
+        if not caller_frame or "self" not in caller_frame.f_locals:
+            return
+
+        caller_self = caller_frame.f_locals["self"]
+        if not is_node_instance(caller_self):
+            return
+
+        for node_output_descriptor in self.__class__:
+            if node_output_descriptor.name in kwargs:
+                continue
+
+            if isinstance(node_output_descriptor.instance, BaseDescriptor):
+                setattr(
+                    self,
+                    node_output_descriptor.name,
+                    node_output_descriptor.instance.resolve(caller_self.state),
+                )
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, (dict, BaseOutputs)):
@@ -223,6 +260,10 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
     def __getitem__(self, key: str) -> Any:
         return getattr(self, key)
 
+    def __vellum_encode__(self) -> Dict[str, Any]:
+        """Return a JSON-serializable representation of this outputs object."""
+        return {descriptor.name: value for descriptor, value in self if value is not undefined}
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler

vellum/workflows/ports/node_ports.py

@@ -40,6 +40,9 @@ class NodePorts(metaclass=_NodePortsMeta):
 
         invoked_ports: Set[Port] = set()
         all_ports = [port for port in self.__class__]
+        if not all_ports:
+            return set()
+
         enforce_single_invoked_conditional_port = validate_ports(all_ports)
 
         for port in all_ports:

vellum/workflows/ports/port.py

@@ -7,8 +7,9 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.descriptors.exceptions import InvalidExpressionException
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors.types import WorkflowErrorCode
-from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.graph import Graph, GraphTarget
+from vellum.workflows.graph.graph import NoPortsNode
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.core import ConditionType
 
@@ -60,12 +61,29 @@ class Port:
         return iter(self._edges)
 
     def __rshift__(self, other: GraphTarget) -> Graph:
+        # Check for trigger target (class-level only)
+        from vellum.workflows.triggers.base import BaseTrigger
+
+        # Check if other is a trigger class
+        if isinstance(other, type) and issubclass(other, BaseTrigger):
+            raise TypeError(
+                f"Cannot create edge targeting trigger {other.__name__}. "
+                f"Triggers must be at the start of a graph path, not as targets. "
+                f"Did you mean: {other.__name__} >> {self.node_class.__name__}?"
+            )
+
         if isinstance(other, set) or isinstance(other, Graph):
             return Graph.from_port(self) >> other
 
         if isinstance(other, Port):
             return Graph.from_port(self) >> Graph.from_port(other)
 
+        if isinstance(other, NoPortsNode):
+            raise ValueError(
+                f"Cannot create edge to {other.node_class.__name__} because it has no ports defined. "
+                f"Nodes with empty Ports classes cannot be connected to other nodes."
+            )
+
         edge = Edge(from_port=self, to_node=other)
         if edge not in self._edges:
             self._edges.append(edge)
@@ -89,7 +107,10 @@ class Port:
             if self._condition is None:
                 return False
 
-            value = self._condition.resolve(state)
+            if isinstance(self._condition, BaseDescriptor):
+                value = self._condition.resolve(state)
+            else:
+                value = self._condition
             return bool(value)
         except InvalidExpressionException as e:
             raise NodeException(
@@ -102,18 +123,12 @@ class Port:
             "name": self.name,
         }
 
+    def __vellum_encode__(self) -> dict:
+        """Return a JSON-serializable representation of this port."""
+        return self.serialize()
+
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
         return core_schema.is_instance_schema(cls)
-
-    def validate(self):
-        if (
-            not self.default
-            and self._condition_type in (ConditionType.IF, ConditionType.ELIF)
-            and self._condition is None
-        ):
-            raise WorkflowInitializationException(
-                f"Class {self.node_class.__name__}'s {self.name} should have a defined condition and cannot be empty."
-            )

vellum/workflows/ports/tests/test_port.py

@@ -0,0 +1,45 @@
+import pytest
+from unittest.mock import Mock
+
+from vellum.workflows.ports.port import Port
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.types.core import ConditionType
+
+
+@pytest.mark.parametrize(
+    "condition_value,expected",
+    [
+        # Truthy values
+        (True, True),
+        (1, True),
+        ("hello", True),
+        ([1, 2, 3], True),
+        ({"key": "value"}, True),
+        (1.5, True),
+        (-1, True),
+        # Falsy values
+        (False, False),
+        (0, False),
+        ("", False),
+        ([], False),
+        ({}, False),
+        (0.0, False),
+        (None, False),
+    ],
+)
+def test_port_on_if_with_non_descriptor_values(condition_value, expected):
+    """Test that Port.on_if() correctly handles non-descriptor values."""
+
+    # Given a port with a non-descriptor value
+    port = Port.on_if(condition_value)
+
+    # Then the port properties should be correct
+    assert port._condition_type == ConditionType.IF
+    assert port._condition == condition_value
+
+    # When we resolve the condition
+    mock_state = Mock(spec=BaseState)
+    result = port.resolve_condition(mock_state)
+
+    # Then the result should be correct
+    assert result == expected

vellum/workflows/ports/utils.py

@@ -1,6 +1,8 @@
 from collections import Counter
 from typing import List
 
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.ports.port import Port
 from vellum.workflows.types.core import ConditionType
 
@@ -31,14 +33,20 @@ def get_port_groups(ports: List[Port]) -> List[List[ConditionType]]:
         else:
             # If we see an ELIF or ELSE without a preceding IF, that's an error
             if not current_port_group:
-                raise ValueError(f"Class {ports_class} must have ports in the following order: on_if, on_elif, on_else")
+                raise NodeException(
+                    message=f"Class {ports_class} must have ports in the following order: on_if, on_elif, on_else",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                )
             current_port_group.append(port_type)
 
     if current_port_group and current_port_group[0] == ConditionType.IF:
         port_groups.append(current_port_group)
     elif current_port_group:
         # If the last group doesn't start with IF, that's an error
-        raise ValueError(f"Class {ports_class} must have ports in the following order: on_if, on_elif, on_else")
+        raise NodeException(
+            message=f"Class {ports_class} must have ports in the following order: on_if, on_elif, on_else",
+            code=WorkflowErrorCode.INVALID_INPUTS,
+        )
 
     return port_groups
 
@@ -51,7 +59,10 @@ def validate_ports(ports: List[Port]) -> bool:
         # Check that each port group is in the correct order
         sorted_group = sorted(group, key=lambda port_type: PORT_TYPE_PRIORITIES[port_type])
         if sorted_group != group:
-            raise ValueError(f"Class {ports_class} must have ports in the following order: on_if, on_elif, on_else")
+            raise NodeException(
+                message=f"Class {ports_class} must have ports in the following order: on_if, on_elif, on_else",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
 
         # Count the types in this port group
         counter = Counter(group)
@@ -61,13 +72,22 @@ def validate_ports(ports: List[Port]) -> bool:
 
         # Apply the rules to each port group
         if number_of_if_ports != 1:
-            raise ValueError(f"Class {ports_class} must have exactly one on_if condition")
+            raise NodeException(
+                message=f"Class {ports_class} must have exactly one on_if condition",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
 
         if number_of_elif_ports > 0 and number_of_if_ports != 1:
-            raise ValueError(f"Class {ports_class} containing on_elif ports must have exactly one on_if condition")
+            raise NodeException(
+                message=f"Class {ports_class} containing on_elif ports must have exactly one on_if condition",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
 
         if number_of_else_ports > 1:
-            raise ValueError(f"Class {ports_class} must have at most one on_else condition")
+            raise NodeException(
+                message=f"Class {ports_class} must have at most one on_else condition",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
 
     enforce_single_invoked_conditional_port = len(port_groups) <= 1
     return enforce_single_invoked_conditional_port

vellum/workflows/references/__init__.py

@@ -4,6 +4,7 @@ from .lazy import LazyReference
 from .node import NodeReference
 from .output import OutputReference
 from .state_value import StateValueReference
+from .trigger import TriggerAttributeReference
 from .vellum_secret import VellumSecretReference
 from .workflow_input import WorkflowInputReference
 
@@ -14,6 +15,7 @@ __all__ = [
     "NodeReference",
     "OutputReference",
     "StateValueReference",
+    "TriggerAttributeReference",
     "VellumSecretReference",
     "WorkflowInputReference",
 ]

vellum/workflows/references/constant.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Dict, Generic, TypeVar
 
 from vellum.workflows.descriptors.base import BaseDescriptor
 
@@ -19,3 +19,6 @@ class ConstantValueReference(BaseDescriptor[_T], Generic[_T]):
 
     def resolve(self, state: "BaseState") -> _T:
         return self._value
+
+    def __vellum_encode__(self) -> Dict[str, Any]:
+        return {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": self._value}}

vellum/workflows/references/environment_variable.py

@@ -1,6 +1,7 @@
 import os
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional
 
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 
 if TYPE_CHECKING:
@@ -8,17 +9,34 @@ if TYPE_CHECKING:
 
 
 class EnvironmentVariableReference(BaseDescriptor[str]):
-    def __init__(self, *, name: str, default: Optional[str]):
-        super().__init__(name=name, types=(str,))
-        self._default = default
+    def __init__(
+        self,
+        *,
+        name: str,
+        # DEPRECATED - to be removed in 0.15.0 release
+        default: Optional[str] = None,
+        serialize_as_constant: bool = False,
+    ):
+        super().__init__(name=name, types=(str,), is_sensitive=True)
+        self._serialize_as_constant = serialize_as_constant
 
-    def resolve(self, state: "BaseState") -> str:
+    def resolve(self, state: "BaseState") -> Any:
         env_value = os.environ.get(self.name)
         if env_value is not None:
             return env_value
 
-        if self._default is not None:
-            return self._default
+        return undefined
 
-        # Fetch Vellum Environment Variable named `self.name` once that project is done
-        raise ValueError(f"No environment variable named '{self.name}' found")
+    def __vellum_encode__(self) -> Dict[str, Any]:
+        return {
+            "type": "ENVIRONMENT_VARIABLE",
+            "environment_variable": self.name,
+        }
+
+    @property
+    def serialize_as_constant(self) -> bool:
+        return self._serialize_as_constant
+
+    @serialize_as_constant.setter
+    def serialize_as_constant(self, value: bool):
+        self._serialize_as_constant = value

vellum/workflows/references/output.py

@@ -35,13 +35,11 @@ class OutputReference(BaseDescriptor[_OutputType], Generic[_OutputType]):
 
     @cached_property
     def id(self) -> UUID:
-        self._outputs_class = self._outputs_class
-
-        node_class = getattr(self._outputs_class, "_node_class", None)
-        if not node_class:
+        parent_class = self._outputs_class.__parent_class__
+        if not parent_class:
             return uuid4()
 
-        output_ids = getattr(node_class, "__output_ids__", {})
+        output_ids = getattr(parent_class, "__output_ids__", {})
         if not isinstance(output_ids, dict):
             return uuid4()
 
@@ -78,7 +76,7 @@ class OutputReference(BaseDescriptor[_OutputType], Generic[_OutputType]):
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, type(self)):
             return False
-        return super().__eq__(other) and id(self._outputs_class) == id(other._outputs_class)
+        return super().__eq__(other) and self._outputs_class == other._outputs_class
 
     def __hash__(self) -> int:
         return hash((self._outputs_class, self._name))

vellum/workflows/references/trigger.py

@@ -0,0 +1,77 @@
+"""Descriptor for referring to trigger attributes in workflow graphs."""
+
+from __future__ import annotations
+
+from uuid import UUID
+from typing import TYPE_CHECKING, Any, Generic, Optional, Tuple, Type, TypeVar, cast
+
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.utils.uuids import get_trigger_attribute_id
+
+if TYPE_CHECKING:
+    from vellum.workflows.state.base import BaseState
+    from vellum.workflows.triggers.base import BaseTrigger
+
+_T = TypeVar("_T")
+
+
+class TriggerAttributeReference(BaseDescriptor[_T], Generic[_T]):
+    """Reference to a trigger attribute defined via type annotations."""
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        types: Tuple[Type[_T], ...],
+        instance: Optional[_T],
+        trigger_class: Type[BaseTrigger],
+    ) -> None:
+        super().__init__(name=name, types=types, instance=instance)
+        self._trigger_class = trigger_class
+
+    @property
+    def trigger_class(self) -> Type[BaseTrigger]:
+        return self._trigger_class
+
+    @property
+    def id(self) -> UUID:
+        """Generate deterministic UUID from trigger class qualname and attribute name."""
+        return get_trigger_attribute_id(self._trigger_class, self.name)
+
+    def resolve(self, state: BaseState) -> _T:
+        trigger_attributes = getattr(state.meta, "trigger_attributes", {})
+        if self in trigger_attributes:
+            return cast(_T, trigger_attributes[self])
+
+        if state.meta.parent:
+            return self.resolve(state.meta.parent)
+
+        if type(None) in self.types:
+            return cast(_T, None)
+
+        raise NodeException(
+            message=f"Missing trigger attribute '{self.name}' for {self._trigger_class.__name__}",
+            code=WorkflowErrorCode.INVALID_INPUTS,
+        )
+
+    def __repr__(self) -> str:
+        return f"{self._trigger_class.__qualname__}.{self.name}"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, TriggerAttributeReference):
+            return False
+        return super().__eq__(other) and self._trigger_class == other._trigger_class
+
+    def __hash__(self) -> int:
+        return hash((self._trigger_class, self._name))
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source_type: Type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.CoreSchema:
+        return core_schema.is_instance_schema(cls)

vellum/workflows/resolvers/base.py

@@ -1,11 +1,25 @@
 from abc import ABC, abstractmethod
-from typing import Iterator
+from uuid import UUID
+from typing import TYPE_CHECKING, Iterator, Optional, Type, Union
 
 from vellum.workflows.events.workflow import WorkflowEvent
+from vellum.workflows.resolvers.types import LoadStateResult
 from vellum.workflows.state.base import BaseState
 
+if TYPE_CHECKING:
+    from vellum.workflows.state.context import WorkflowContext
+    from vellum.workflows.workflows.base import BaseWorkflow
+
 
 class BaseWorkflowResolver(ABC):
+    def __init__(self):
+        self._context: Optional["WorkflowContext"] = None
+        self._workflow_class: Optional[Type["BaseWorkflow"]] = None
+
+    def register_workflow_instance(self, workflow_instance: "BaseWorkflow") -> None:
+        self._workflow_class = type(workflow_instance)
+        self._context = workflow_instance.context
+
     @abstractmethod
     def get_latest_execution_events(self) -> Iterator[WorkflowEvent]:
         pass
@@ -13,3 +27,7 @@ class BaseWorkflowResolver(ABC):
     @abstractmethod
     def get_state_snapshot_history(self) -> Iterator[BaseState]:
         pass
+
+    @abstractmethod
+    def load_state(self, previous_execution_id: Optional[Union[UUID, str]] = None) -> Optional[LoadStateResult]:
+        pass