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

vellum/workflows/nodes/mocks.py

@@ -1,108 +1,28 @@
-from functools import reduce
-from uuid import UUID
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Sequence, Type, Union
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Type, Union
 
 from pydantic import ConfigDict, SerializationInfo, ValidationError, field_serializer, model_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.array_vellum_value import ArrayVellumValue
-from vellum.client.types.vellum_value import VellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.events.types import default_serializer
-from vellum.workflows.exceptions import WorkflowInitializationException
+from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
+from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
 
 if TYPE_CHECKING:
     from vellum.workflows import BaseWorkflow
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 
-class _RawLogicalCondition(UniversalBaseModel):
-    type: Literal["LOGICAL_CONDITION"] = "LOGICAL_CONDITION"
-    lhs_variable_id: UUID
-    operator: Literal["==", ">", ">=", "<", "<=", "!="]
-    rhs_variable_id: UUID
-
-
-class _RawLogicalConditionGroup(UniversalBaseModel):
-    type: Literal["LOGICAL_CONDITION_GROUP"] = "LOGICAL_CONDITION_GROUP"
-    conditions: List["_RawLogicalExpression"]
-    combinator: Literal["AND", "OR"]
-    negated: bool
-
-
-_RawLogicalExpression = Union[_RawLogicalCondition, _RawLogicalConditionGroup]
-
-
-class _RawLogicalExpressionVariable(UniversalBaseModel):
-    id: UUID
-
-
-class _RawMockWorkflowNodeExecutionConstantValuePointer(_RawLogicalExpressionVariable):
-    type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
-    variable_value: VellumValue
-
-
-class _RawMockWorkflowNodeExecutionNodeExecutionCounterPointer(_RawLogicalExpressionVariable):
-    type: Literal["EXECUTION_COUNTER"] = "EXECUTION_COUNTER"
-    node_id: UUID
-
-
-class _RawMockWorkflowNodeExecutionInputVariablePointer(_RawLogicalExpressionVariable):
-    type: Literal["INPUT_VARIABLE"] = "INPUT_VARIABLE"
-    input_variable_id: UUID
-
-
-class _RawMockWorkflowNodeExecutionNodeOutputPointer(_RawLogicalExpressionVariable):
-    type: Literal["NODE_OUTPUT"] = "NODE_OUTPUT"
-    node_id: UUID
-    input_id: UUID
-
-
-class _RawMockWorkflowNodeExecutionNodeInputPointer(_RawLogicalExpressionVariable):
-    type: Literal["NODE_INPUT"] = "NODE_INPUT"
-    node_id: UUID
-    input_id: UUID
-
-
-_RawMockWorkflowNodeExecutionValuePointer = Union[
-    _RawMockWorkflowNodeExecutionConstantValuePointer,
-    _RawMockWorkflowNodeExecutionNodeExecutionCounterPointer,
-    _RawMockWorkflowNodeExecutionInputVariablePointer,
-    _RawMockWorkflowNodeExecutionNodeOutputPointer,
-    _RawMockWorkflowNodeExecutionNodeInputPointer,
-]
-
-
-class _RawMockWorkflowNodeWhenCondition(UniversalBaseModel):
-    expression: _RawLogicalExpression
-    variables: List[_RawMockWorkflowNodeExecutionValuePointer]
-
-
-class _RawMockWorkflowNodeThenOutput(UniversalBaseModel):
-    output_id: UUID
-    value: _RawMockWorkflowNodeExecutionValuePointer
-
-
-class _RawMockWorkflowNodeExecution(UniversalBaseModel):
-    when_condition: _RawMockWorkflowNodeWhenCondition
-    then_outputs: List[_RawMockWorkflowNodeThenOutput]
-
-
-class _RawMockWorkflowNodeConfig(UniversalBaseModel):
-    type: Literal["WORKFLOW_NODE_OUTPUT"] = "WORKFLOW_NODE_OUTPUT"
-    node_id: UUID
-    mock_executions: List[_RawMockWorkflowNodeExecution]
-
-
 class MockNodeExecution(UniversalBaseModel):
     when_condition: BaseDescriptor
     then_outputs: BaseOutputs
+    disabled: Optional[bool] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -110,7 +30,12 @@ class MockNodeExecution(UniversalBaseModel):
     def serialize_full_model(self, handler: Callable[[Any], Any], info: SerializationInfo) -> Dict[str, Any]:
         """Serialize the model and add node_id field computed from then_outputs."""
         serialized = handler(self)
-        serialized["node_id"] = str(self.then_outputs.__class__.__parent_class__.__id__)
+        node_class = self.then_outputs.__class__.__parent_class__
+        unadorned_node = get_unadorned_node(node_class)  # type: ignore[arg-type]
+        serialized["node_id"] = str(unadorned_node.__id__)
+        serialized["type"] = "NODE_EXECUTION"
+        if self.disabled is None:
+            del serialized["disabled"]
         return serialized
 
     @field_serializer("then_outputs")
@@ -129,145 +54,54 @@ class MockNodeExecution(UniversalBaseModel):
     def validate_all(
         raw_mock_workflow_node_configs: Optional[List[Any]],
         workflow: Type["BaseWorkflow"],
+        descriptor_validator: Optional[Callable[[dict, Type["BaseWorkflow"]], BaseDescriptor]] = None,
     ) -> Optional[List["MockNodeExecution"]]:
         if not raw_mock_workflow_node_configs:
             return None
 
         ArrayVellumValue.model_rebuild()
-        try:
-            mock_workflow_node_configs = [
-                _RawMockWorkflowNodeConfig.model_validate(raw_mock_workflow_node_config)
-                for raw_mock_workflow_node_config in raw_mock_workflow_node_configs
-            ]
-        except ValidationError as e:
-            raise WorkflowInitializationException(
-                message="Failed to validate mock node executions",
-                code=WorkflowErrorCode.INVALID_INPUTS,
-                workflow_definition=workflow,
-            ) from e
-
-        nodes = {node.__id__: node for node in workflow.get_nodes()}
-        node_output_name_by_id = {
-            node.__output_ids__[output.name]: output.name for node in workflow.get_nodes() for output in node.Outputs
-        }
-
-        # We need to support the old way that the Vellum App's WorkflowRunner used to define Node Mocks in order to
-        # avoid needing to update the mock resolution strategy that it and the frontend uses. The path towards
-        # cleaning this up will go as follows:
-        # 1. ✅ Release Mock support in SDK-Enabled Workflows
-        # 2. ✅ Deprecate Mock support in non-SDK enabled Workflows, encouraging users to migrate to SDK Workflows
-        # 3. ✅ Remove the old mock resolution strategy
-        # 4. Update this SDK to handle the new mock resolution strategy with WorkflowValueDescriptors
-        # 5. Cutover the Vellum App to the new mock resolution strategy
-        # 6. Remove the old mock resolution strategy from this SDK
-        def _translate_raw_logical_expression(
-            raw_logical_expression: _RawLogicalExpression,
-            raw_variables: List[_RawMockWorkflowNodeExecutionValuePointer],
-        ) -> BaseDescriptor:
-            if raw_logical_expression.type == "LOGICAL_CONDITION":
-                return _translate_raw_logical_condition(raw_logical_expression, raw_variables)
-            else:
-                return _translate_raw_logical_condition_group(raw_logical_expression, raw_variables)
-
-        def _translate_raw_logical_condition_group(
-            raw_logical_condition_group: _RawLogicalConditionGroup,
-            raw_variables: List[_RawMockWorkflowNodeExecutionValuePointer],
-        ) -> BaseDescriptor:
-            if not raw_logical_condition_group.conditions:
-                return ConstantValueReference(True)
-
-            conditions = [
-                _translate_raw_logical_expression(condition, raw_variables)
-                for condition in raw_logical_condition_group.conditions
-            ]
-            return reduce(
-                lambda acc, condition: (
-                    acc and condition if raw_logical_condition_group.combinator == "AND" else acc or condition
-                ),
-                conditions,
-            )
-
-        def _translate_raw_logical_condition(
-            raw_logical_condition: _RawLogicalCondition,
-            raw_variables: List[_RawMockWorkflowNodeExecutionValuePointer],
-        ) -> BaseDescriptor:
-            variable_by_id = {v.id: v for v in raw_variables}
-            lhs = _translate_raw_logical_expression_variable(variable_by_id[raw_logical_condition.lhs_variable_id])
-            rhs = _translate_raw_logical_expression_variable(variable_by_id[raw_logical_condition.rhs_variable_id])
-            if raw_logical_condition.operator == ">":
-                return lhs.greater_than(rhs)
-            elif raw_logical_condition.operator == ">=":
-                return lhs.greater_than_or_equal_to(rhs)
-            elif raw_logical_condition.operator == "<":
-                return lhs.less_than(rhs)
-            elif raw_logical_condition.operator == "<=":
-                return lhs.less_than_or_equal_to(rhs)
-            elif raw_logical_condition.operator == "==":
-                return lhs.equals(rhs)
-            elif raw_logical_condition.operator == "!=":
-                return lhs.does_not_equal(rhs)
-            else:
-                raise WorkflowInitializationException(
-                    message=f"Unsupported logical operator: {raw_logical_condition.operator}",
-                    workflow_definition=workflow,
-                )
-
-        def _translate_raw_logical_expression_variable(
-            raw_variable: _RawMockWorkflowNodeExecutionValuePointer,
-        ) -> BaseDescriptor:
-            if raw_variable.type == "CONSTANT_VALUE":
-                return ConstantValueReference(raw_variable.variable_value.value)
-            elif raw_variable.type == "EXECUTION_COUNTER":
-                node = nodes[raw_variable.node_id]
-                return node.Execution.count
-            else:
+        mock_node_executions: list[MockNodeExecution] = []
+        for raw_mock_workflow_node_config in raw_mock_workflow_node_configs:
+            if not isinstance(raw_mock_workflow_node_config, dict):
                 raise WorkflowInitializationException(
-                    message=f"Unsupported logical expression type: {raw_variable.type}",
+                    message=f"Invalid mock node execution type: {type(raw_mock_workflow_node_config)}",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
                     workflow_definition=workflow,
                 )
 
-        mock_node_executions = []
-        for mock_workflow_node_config in mock_workflow_node_configs:
-            for mock_execution in mock_workflow_node_config.mock_executions:
-                try:
-                    when_condition = _translate_raw_logical_expression(
-                        mock_execution.when_condition.expression,
-                        mock_execution.when_condition.variables,
+            try:
+                mock_node_executions.append(
+                    MockNodeExecution.model_validate(
+                        raw_mock_workflow_node_config,
+                        context={
+                            "workflow": workflow,
+                            "node_id": raw_mock_workflow_node_config.get("node_id"),
+                            "descriptor_validator": lambda value: (
+                                descriptor_validator(value, workflow)
+                                if descriptor_validator
+                                else ConstantValueReference(False)
+                            ),
+                        },
                     )
-
-                    then_outputs = nodes[mock_workflow_node_config.node_id].Outputs()
-                    for then_output in mock_execution.then_outputs:
-                        node_output_name = node_output_name_by_id.get(then_output.output_id)
-                        if node_output_name is None:
-                            node_id = mock_workflow_node_config.node_id
-                            raise WorkflowInitializationException(
-                                message=f"Output {then_output.output_id} not found in node {node_id}",
-                                workflow_definition=workflow,
-                            )
-
-                        resolved_output_reference = _translate_raw_logical_expression_variable(then_output.value)
-                        if isinstance(resolved_output_reference, ConstantValueReference):
-                            setattr(
-                                then_outputs,
-                                node_output_name,
-                                resolved_output_reference._value,
-                            )
-                        else:
-                            ref_type = type(resolved_output_reference)
-                            raise WorkflowInitializationException(
-                                message=f"Unsupported resolved output reference type: {ref_type}",
-                                workflow_definition=workflow,
-                            )
-
-                    mock_node_executions.append(
-                        MockNodeExecution(
-                            when_condition=when_condition,
-                            then_outputs=then_outputs,
-                        )
+                )
+            except WorkflowInitializationException as e:
+                # If the node is not found in the workflow, skip it with a warning
+                node_id = raw_mock_workflow_node_config.get("node_id")
+                raw_data = e.raw_data or {}
+                if raw_data.get("node_ref") == node_id:
+                    logger.warning(
+                        "Skipping mock for node %s: node not found in workflow %s",
+                        node_id,
+                        workflow.__name__,
                     )
-                except Exception as e:
-                    logger.exception("Failed to validate mock node execution", exc_info=e)
                     continue
+                raise
+            except (ValidationError, NodeException) as e:
+                raise WorkflowInitializationException(
+                    message="Failed to validate mock node executions",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                    workflow_definition=workflow,
+                ) from e
 
         return mock_node_executions
 

vellum/workflows/nodes/tests/test_mocks.py

@@ -1,98 +1,6 @@
-import uuid
-
-from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows import BaseWorkflow
-from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.mocks import MockNodeExecution
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
-from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
-
-
-def test_mocks__parse_from_app():
-    # GIVEN a PromptNode
-    class PromptNode(InlinePromptNode):
-        pass
-
-    # AND a workflow class with that PromptNode
-    class MyWorkflow(BaseWorkflow):
-        graph = PromptNode
-
-    # AND a mock workflow node execution from the app
-    raw_mock_workflow_node_execution = [
-        {
-            "type": "WORKFLOW_NODE_OUTPUT",
-            "node_id": str(PromptNode.__id__),
-            "mock_executions": [
-                {
-                    "when_condition": {
-                        "expression": {
-                            "type": "LOGICAL_CONDITION_GROUP",
-                            "combinator": "AND",
-                            "negated": False,
-                            "conditions": [
-                                {
-                                    "type": "LOGICAL_CONDITION",
-                                    "lhs_variable_id": "e60902d5-6892-4916-80c1-f0130af52322",
-                                    "operator": ">=",
-                                    "rhs_variable_id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                                }
-                            ],
-                        },
-                        "variables": [
-                            {
-                                "type": "EXECUTION_COUNTER",
-                                "node_id": str(PromptNode.__id__),
-                                "id": "e60902d5-6892-4916-80c1-f0130af52322",
-                            },
-                            {
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "NUMBER", "value": 0},
-                                "id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                            },
-                        ],
-                    },
-                    "then_outputs": [
-                        {
-                            "output_id": "bbd955e8-b939-45f9-8617-f1405dddd2b3",
-                            "value": {
-                                "id": "27006b2a-fa81-430c-a0b2-c66a9351fc68",
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "STRING", "value": "Hello"},
-                            },
-                        },
-                        {
-                            "output_id": "e50bf6e9-417b-4c04-babb-94e7dccb9f2c",
-                            "value": {
-                                "id": "4559c778-6e27-4cfe-a460-734ba62a5082",
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "ARRAY", "value": [{"type": "STRING", "value": "Hello"}]},
-                            },
-                        },
-                    ],
-                }
-            ],
-        }
-    ]
-
-    # WHEN we parse the mock workflow node execution
-    node_output_mocks = MockNodeExecution.validate_all(
-        raw_mock_workflow_node_execution,
-        MyWorkflow,
-    )
-
-    # THEN we get a list of MockNodeExecution objects
-    assert node_output_mocks
-    assert len(node_output_mocks) == 1
-    assert node_output_mocks[0] == MockNodeExecution(
-        when_condition=PromptNode.Execution.count.greater_than_or_equal_to(0.0),
-        then_outputs=PromptNode.Outputs(
-            text="Hello",
-            results=[
-                StringVellumValue(value="Hello"),
-            ],
-        ),
-    )
 
 
 def test_mocks__parse_none_still_runs():
@@ -120,88 +28,3 @@ def test_mocks__parse_none_still_runs():
 
     # THEN it was successful
     assert final_event.name == "workflow.execution.fulfilled"
-
-
-def test_mocks__use_id_from_display():
-    # GIVEN a Base Node
-    class StartNode(BaseNode):
-        class Outputs(BaseNode.Outputs):
-            foo: str
-
-    # AND a workflow class with that Node
-    class MyWorkflow(BaseWorkflow):
-        graph = StartNode
-
-        class Outputs(BaseWorkflow.Outputs):
-            final_value = StartNode.Outputs.foo
-
-    # AND a display class on that Base Node
-    node_output_id = uuid.uuid4()
-
-    class StartNodeDisplay(BaseNodeDisplay[StartNode]):
-        output_display = {StartNode.Outputs.foo: NodeOutputDisplay(id=node_output_id, name="foo")}
-
-    # AND a mock workflow node execution from the app
-    raw_mock_workflow_node_execution = [
-        {
-            "type": "WORKFLOW_NODE_OUTPUT",
-            "node_id": str(StartNode.__id__),
-            "mock_executions": [
-                {
-                    "when_condition": {
-                        "expression": {
-                            "type": "LOGICAL_CONDITION_GROUP",
-                            "combinator": "AND",
-                            "negated": False,
-                            "conditions": [
-                                {
-                                    "type": "LOGICAL_CONDITION",
-                                    "lhs_variable_id": "e60902d5-6892-4916-80c1-f0130af52322",
-                                    "operator": ">=",
-                                    "rhs_variable_id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                                }
-                            ],
-                        },
-                        "variables": [
-                            {
-                                "type": "EXECUTION_COUNTER",
-                                "node_id": str(StartNode.__id__),
-                                "id": "e60902d5-6892-4916-80c1-f0130af52322",
-                            },
-                            {
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "NUMBER", "value": 0},
-                                "id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                            },
-                        ],
-                    },
-                    "then_outputs": [
-                        {
-                            "output_id": str(node_output_id),
-                            "value": {
-                                "id": "27006b2a-fa81-430c-a0b2-c66a9351fc68",
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "STRING", "value": "Hello"},
-                            },
-                        },
-                    ],
-                }
-            ],
-        }
-    ]
-
-    # WHEN we parsed the raw data on `MockNodeExecution`
-    node_output_mocks = MockNodeExecution.validate_all(
-        raw_mock_workflow_node_execution,
-        MyWorkflow,
-    )
-
-    # THEN we get the expected list of MockNodeExecution objects
-    assert node_output_mocks
-    assert len(node_output_mocks) == 1
-    assert node_output_mocks[0] == MockNodeExecution(
-        when_condition=StartNode.Execution.count.greater_than_or_equal_to(0.0),
-        then_outputs=StartNode.Outputs(
-            foo="Hello",
-        ),
-    )

vellum/workflows/nodes/utils.py

@@ -3,7 +3,20 @@ import inspect
 import json
 import sys
 from types import ModuleType
-from typing import Any, Callable, Dict, ForwardRef, List, Optional, Type, TypeVar, Union, get_args, get_origin
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    ForwardRef,
+    List,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    get_args,
+    get_origin,
+)
 
 from pydantic import BaseModel, create_model
 
@@ -12,8 +25,6 @@ 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
-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 (
@@ -23,13 +34,16 @@ from vellum.workflows.types.code_execution_node_wrappers import (
     clean_for_dict_wrapper,
 )
 from vellum.workflows.types.core import Json
-from vellum.workflows.types.generics import NodeType
+from vellum.workflows.types.generics import NodeType, import_base_adornment_node
+
+if TYPE_CHECKING:
+    from vellum.workflows.nodes import BaseNode
 
 ADORNMENT_MODULE_NAME = "<adornment>"
 
 
 @cache
-def get_unadorned_node(node: Type[BaseNode]) -> Type[BaseNode]:
+def get_unadorned_node(node: Type["BaseNode"]) -> Type["BaseNode"]:
     wrapped_node = get_wrapped_node(node)
     if wrapped_node is not None:
         return get_unadorned_node(wrapped_node)
@@ -46,20 +60,21 @@ def get_unadorned_port(port: Port) -> Port:
     return getattr(unadorned_node.Ports, port.name)
 
 
-def get_wrapped_node(node: Type[NodeType]) -> Optional[Type[BaseNode]]:
+def get_wrapped_node(node: Type[NodeType]) -> Optional[Type["BaseNode"]]:
+    BaseAdornmentNode = import_base_adornment_node()
     if not issubclass(node, BaseAdornmentNode):
         return None
 
     return node.__wrapped_node__
 
 
-AdornableNode = TypeVar("AdornableNode", bound=BaseNode)
+AdornableNode = TypeVar("AdornableNode", bound="BaseNode")
 
 
 def create_adornment(
     adornable_cls: Type[AdornableNode], attributes: Optional[dict[str, Any]] = None
 ) -> Callable[..., Type["AdornableNode"]]:
-    def decorator(inner_cls: Type[BaseNode]) -> Type["AdornableNode"]:
+    def decorator(inner_cls: Type["BaseNode"]) -> Type["AdornableNode"]:
         # Investigate how to use dependency injection to avoid circular imports
         # https://app.shortcut.com/vellum/story/4116
         from vellum.workflows import BaseWorkflow

vellum/workflows/outputs/base.py

@@ -3,7 +3,7 @@ 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
+from pydantic import GetCoreSchemaHandler, ValidationInfo
 from pydantic_core import core_schema
 
 from vellum.workflows.constants import undefined
@@ -12,7 +12,7 @@ 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.generics import import_workflow_class, is_node_instance
 from vellum.workflows.types.utils import get_class_attr_names, infer_types
 
 _Delta = TypeVar("_Delta")
@@ -268,4 +268,37 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
     def __get_pydantic_core_schema__(
         cls, source_type: Type[Any], handler: GetCoreSchemaHandler
     ) -> core_schema.CoreSchema:
-        return core_schema.is_instance_schema(cls)
+        def validate(value: Any, info: ValidationInfo) -> Any:
+            if isinstance(value, cls):
+                return value
+
+            if not isinstance(value, dict):
+                raise TypeError(f"Value must be an instance of {cls.__name__} or a dict")
+
+            context = info.context
+            if not isinstance(context, dict):
+                raise TypeError(f"Unexpected type for context: {type(context)}")
+
+            workflow = context.get("workflow")
+            node_id = context.get("node_id")
+            base_workflow_class = import_workflow_class()
+            if (
+                not workflow
+                or not node_id
+                or not isinstance(node_id, str)
+                or not isinstance(workflow, type)
+                or not issubclass(workflow, base_workflow_class)
+            ):
+                return cls(**value)
+
+            node_class = workflow.resolve_node_ref(node_id)
+            declared_fields = {descriptor.name for descriptor in node_class.Outputs}
+            filtered_value = {k: v for k, v in value.items() if k in declared_fields}
+            return node_class.Outputs(**filtered_value)
+
+        return core_schema.union_schema(
+            [
+                core_schema.is_instance_schema(cls),
+                core_schema.with_info_after_validator_function(validate, core_schema.dict_schema()),
+            ]
+        )

vellum/workflows/references/environment_variable.py

@@ -13,12 +13,10 @@ class EnvironmentVariableReference(BaseDescriptor[str]):
         self,
         *,
         name: str,
-        # DEPRECATED - to be removed in 0.15.0 release
+        # DEPRECATED - to be removed in 2.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") -> Any:
         env_value = os.environ.get(self.name)
@@ -32,11 +30,3 @@ class EnvironmentVariableReference(BaseDescriptor[str]):
             "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/lazy.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Callable, Generic, TypeVar, Union, get_args
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.types.generics import is_workflow_class
 
 if TYPE_CHECKING:
     from vellum.workflows.state.base import BaseState
@@ -36,6 +37,13 @@ class LazyReference(BaseDescriptor[_T], Generic[_T]):
                 if str(output_reference) == self._get:
                     return value
 
+            # Check workflow outputs
+            workflow_definition = state.meta.workflow_definition
+            if is_workflow_class(workflow_definition):
+                for output_reference in workflow_definition.Outputs:
+                    if str(output_reference) == self._get:
+                        return resolve_value(output_reference.instance, state)  # type: ignore[return-value]
+
             child_reference = self.resolve(state.meta.parent) if state.meta.parent else None
 
             # Fix typing surrounding the return value of node outputs/output descriptors