vellum-ai 0.11.9__py3-none-any.whl → 0.12.0__py3-none-any.whl

vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py

@@ -3,7 +3,7 @@ from typing import ClassVar, Generator, Generic, Iterator, List, Optional, Union
 
 from vellum import AdHocExecutePromptEvent, ExecutePromptEvent, PromptOutput
 from vellum.core import RequestOptions
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode, vellum_error_to_workflow_error
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
@@ -29,7 +29,7 @@ class BasePromptNode(BaseNode, Generic[StateType]):
         if outputs is None:
             raise NodeException(
                 message="Expected to receive outputs from Prompt",
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
     def _process_prompt_event_stream(self) -> Generator[BaseOutput, None, Optional[List[PromptOutput]]]:
@@ -45,15 +45,7 @@ class BasePromptNode(BaseNode, Generic[StateType]):
                 outputs = event.outputs
                 yield BaseOutput(name="results", value=event.outputs)
             elif event.state == "REJECTED":
-                if event.error.code in VellumErrorCode._value2member_map_:
-                    raise NodeException(
-                        message=event.error.message,
-                        code=VellumErrorCode(event.error.code),
-                    )
-                else:
-                    raise NodeException(
-                        message=event.error.message,
-                        code=VellumErrorCode.INTERNAL_ERROR,
-                    )
+                workflow_error = vellum_error_to_workflow_error(event.error)
+                raise NodeException.of(workflow_error)
 
         return outputs

vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py

@@ -17,7 +17,7 @@ from vellum import (
 from vellum.client import RequestOptions
 from vellum.workflows.constants import OMIT
 from vellum.workflows.context import get_parent_context
-from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
 from vellum.workflows.nodes.displayable.bases.inline_prompt_node.constants import DEFAULT_PROMPT_PARAMETERS
@@ -120,7 +120,7 @@ class BaseInlinePromptNode(BasePromptNode, Generic[StateType]):
             else:
                 raise NodeException(
                     message=f"Unrecognized input type for input '{input_name}': {input_value.__class__}",
-                    code=VellumErrorCode.INVALID_INPUTS,
+                    code=WorkflowErrorCode.INVALID_INPUTS,
                 )
 
         return input_variables, input_values

vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py

@@ -14,7 +14,7 @@ from vellum import (
 from vellum.client import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
 from vellum.workflows.context import get_parent_context
-from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
 from vellum.workflows.types.generics import StateType
@@ -101,7 +101,7 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
             else:
                 raise NodeException(
                     message=f"Unrecognized input type for input '{input_name}': {input_value.__class__}",
-                    code=VellumErrorCode.INVALID_INPUTS,
+                    code=WorkflowErrorCode.INVALID_INPUTS,
                 )
 
         return compiled_inputs

vellum/workflows/nodes/displayable/bases/search_node.py

@@ -12,7 +12,7 @@ from vellum import (
     SearchWeightsRequest,
 )
 from vellum.core import ApiError, RequestOptions
-from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
@@ -83,12 +83,12 @@ class BaseSearchNode(BaseNode[StateType], Generic[StateType]):
         except NotFoundError:
             raise NodeException(
                 message=f"Document Index '{self.document_index}' not found",
-                code=VellumErrorCode.INVALID_INPUTS,
+                code=WorkflowErrorCode.INVALID_INPUTS,
             )
         except ApiError:
             raise NodeException(
                 message=f"An error occurred while searching against Document Index '{self.document_index}'",  # noqa: E501
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
     def run(self) -> Outputs:

vellum/workflows/nodes/displayable/code_execution_node/node.py

@@ -20,7 +20,7 @@ from vellum import (
     VellumValue,
 )
 from vellum.core import RequestOptions
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
@@ -73,7 +73,8 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
     request_options: Optional[RequestOptions] = None - The request options to use for the custom script.
     """
 
-    filepath: ClassVar[str]
+    filepath: ClassVar[Optional[str]] = None
+    code: ClassVar[Optional[str]] = None
 
     code_inputs: ClassVar[EntityInputsInterface]
     runtime: CodeExecutionRuntime = "PYTHON_3_11_6"
@@ -101,7 +102,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
 
         if code_execution.output.type != expected_output_type:
             raise NodeException(
-                code=VellumErrorCode.INVALID_OUTPUTS,
+                code=WorkflowErrorCode.INVALID_OUTPUTS,
                 message=f"Expected an output of type '{expected_output_type}', received '{code_execution.output.type}'",
             )
 
@@ -184,18 +185,33 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
             else:
                 raise NodeException(
                     message=f"Unrecognized input type for input '{input_name}'",
-                    code=VellumErrorCode.INVALID_INPUTS,
+                    code=WorkflowErrorCode.INVALID_INPUTS,
                 )
 
         return compiled_inputs
 
     def _resolve_code(self) -> str:
+        if self.code and self.filepath:
+            raise NodeException(
+                message="Cannot specify both `code` and `filepath` for a CodeExecutionNode",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+
+        if self.code:
+            return self.code
+
+        if not self.filepath:
+            raise NodeException(
+                message="Must specify either `code` or `filepath` for a CodeExecutionNode",
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            )
+
         root = inspect.getfile(self.__class__)
         code = read_file_from_path(node_filepath=root, script_filepath=self.filepath)
         if not code:
             raise NodeException(
                 message=f"Filepath '{self.filepath}' does not exist",
-                code=VellumErrorCode.INVALID_INPUTS,
+                code=WorkflowErrorCode.INVALID_INPUTS,
             )
 
         return code

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -1,6 +1,8 @@
+import pytest
 import os
 
 from vellum import CodeExecutorResponse, NumberVellumValue, StringInput
+from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
 from vellum.workflows.references.vellum_secret import VellumSecretReference
@@ -62,6 +64,145 @@ def main(word: str) -> int:
     )
 
 
+def test_run_workflow__code_attribute(vellum_client):
+    """Confirm that CodeExecutionNodes can use the `code` attribute to specify the code to execute."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        code = """\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": 5, "log": "hello"}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            StringInput(name="word", value="hello"),
+        ],
+        code="""\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="NUMBER",
+        packages=[],
+        request_options=None,
+    )
+
+
+def test_run_workflow__code_and_filepath_defined(vellum_client):
+    """Confirm that CodeExecutionNodes raise an error if both `code` and `filepath` are defined."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    fixture = os.path.abspath(os.path.join(__file__, "../fixtures/main.py"))
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        filepath = fixture
+        code = """\
+def main(word: str) -> int:
+    print(word)  # noqa: T201
+    return len(word)
+"""
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Cannot specify both `code` and `filepath` for a CodeExecutionNode"
+
+
+def test_run_workflow__code_and_filepath_not_defined(vellum_client):
+    """Confirm that CodeExecutionNodes raise an error if neither `code` nor `filepath` are defined."""
+
+    # GIVEN a node that subclasses CodeExecutionNode
+    class Inputs(BaseInputs):
+        word: str
+
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, int]):
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": Inputs.word,
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="hello",
+        output=NumberVellumValue(value=5),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(word="hello")),
+        )
+    )
+    with pytest.raises(NodeException) as exc_info:
+        node.run()
+
+    # THEN the node should have produced the exception we expected
+    assert exc_info.value.message == "Must specify either `code` or `filepath` for a CodeExecutionNode"
+
+
 def test_run_workflow__vellum_secret(vellum_client):
     """Confirm that CodeExecutionNodes can use Vellum Secrets"""
 

vellum/workflows/nodes/displayable/guardrail_node/node.py

@@ -4,7 +4,7 @@ from typing import Any, ClassVar, Dict, Generic, List, Optional, Union, cast
 from vellum import ChatHistoryInput, ChatMessage, JsonInput, MetricDefinitionInput, NumberInput, StringInput
 from vellum.core import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
@@ -46,7 +46,7 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
         if not isinstance(score, float):
             raise NodeException(
                 message="Metric execution must have one output named 'score' with type 'float'",
-                code=VellumErrorCode.INVALID_OUTPUTS,
+                code=WorkflowErrorCode.INVALID_OUTPUTS,
             )
 
         metric_outputs.pop("score")
@@ -90,7 +90,7 @@ class GuardrailNode(BaseNode[StateType], Generic[StateType]):
             else:
                 raise NodeException(
                     message=f"Unrecognized input type for input '{input_name}'",
-                    code=VellumErrorCode.INVALID_INPUTS,
+                    code=WorkflowErrorCode.INVALID_INPUTS,
                 )
 
         return compiled_inputs

vellum/workflows/nodes/displayable/inline_prompt_node/node.py

@@ -1,6 +1,6 @@
 from typing import Iterator
 
-from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BaseInlinePromptNode as BaseInlinePromptNode
 from vellum.workflows.outputs import BaseOutput
@@ -34,7 +34,7 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
         if not outputs:
             raise NodeException(
                 message="Expected to receive outputs from Prompt",
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
         string_output = next((output for output in outputs if output.type == "STRING"), None)
@@ -43,7 +43,7 @@ class InlinePromptNode(BaseInlinePromptNode[StateType]):
             is_plural = len(output_types) > 1
             raise NodeException(
                 message=f"Expected to receive a non-null string output from Prompt. Only found outputs of type{'s' if is_plural else ''}: {', '.join(output_types)}",  # noqa: E501
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
         yield BaseOutput(name="text", value=string_output.value)

vellum/workflows/nodes/displayable/prompt_deployment_node/node.py

@@ -1,6 +1,6 @@
 from typing import Iterator
 
-from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BasePromptDeploymentNode as BasePromptDeploymentNode
 from vellum.workflows.outputs import BaseOutput
@@ -37,7 +37,7 @@ class PromptDeploymentNode(BasePromptDeploymentNode[StateType]):
         if not outputs:
             raise NodeException(
                 message="Expected to receive outputs from Prompt",
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
         string_output = next((output for output in outputs if output.type == "STRING"), None)
@@ -46,7 +46,7 @@ class PromptDeploymentNode(BasePromptDeploymentNode[StateType]):
             is_plural = len(output_types) > 1
             raise NodeException(
                 message=f"Expected to receive a non-null string output from Prompt. Only found outputs of type{'s' if is_plural else ''}: {', '.join(output_types)}",  # noqa: E501
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
         yield BaseOutput(name="text", value=string_output.value)

vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -14,7 +14,8 @@ from vellum import (
 from vellum.core import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
 from vellum.workflows.context import get_parent_context
-from vellum.workflows.errors import VellumErrorCode
+from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.errors.types import workflow_event_error_to_workflow_error
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases.base_subworkflow_node.node import BaseSubworkflowNode
 from vellum.workflows.outputs.base import BaseOutput
@@ -84,7 +85,7 @@ class SubworkflowDeploymentNode(BaseSubworkflowNode[StateType], Generic[StateTyp
             else:
                 raise NodeException(
                     message=f"Unrecognized input type for input '{input_name}'",
-                    code=VellumErrorCode.INVALID_INPUTS,
+                    code=WorkflowErrorCode.INVALID_INPUTS,
                 )
 
         return compiled_inputs
@@ -136,23 +137,15 @@ class SubworkflowDeploymentNode(BaseSubworkflowNode[StateType], Generic[StateTyp
                 if not error:
                     raise NodeException(
                         message="Expected to receive an error from REJECTED event",
-                        code=VellumErrorCode.INTERNAL_ERROR,
-                    )
-                elif error.code in VellumErrorCode._value2member_map_:
-                    raise NodeException(
-                        message=error.message,
-                        code=VellumErrorCode(error.code),
-                    )
-                else:
-                    raise NodeException(
-                        message=error.message,
-                        code=VellumErrorCode.INTERNAL_ERROR,
+                        code=WorkflowErrorCode.INTERNAL_ERROR,
                     )
+                workflow_error = workflow_event_error_to_workflow_error(error)
+                raise NodeException.of(workflow_error)
 
         if outputs is None:
             raise NodeException(
                 message="Expected to receive outputs from Workflow Deployment",
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=WorkflowErrorCode.INTERNAL_ERROR,
             )
 
         # For any outputs somehow in our final fulfilled outputs array,

vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py

@@ -11,8 +11,8 @@ from vellum import (
     StringVellumValue,
     VellumError,
 )
-from vellum.workflows.errors import VellumError as SdkVellumError
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors import WorkflowError as SdkVellumError
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
@@ -104,7 +104,7 @@ def test_inline_text_prompt_node__catch_provider_error(vellum_adhoc_prompt_clien
     class State(BaseState):
         pass
 
-    @TryNode.wrap(on_error_code=VellumErrorCode.PROVIDER_ERROR)
+    @TryNode.wrap(on_error_code=WorkflowErrorCode.PROVIDER_ERROR)
     class MyInlinePromptNode(InlinePromptNode):
         ml_model = "gpt-4o"
         prompt_inputs = {}
@@ -143,7 +143,7 @@ def test_inline_text_prompt_node__catch_provider_error(vellum_adhoc_prompt_clien
             name="error",
             value=SdkVellumError(
                 message="OpenAI failed",
-                code=VellumErrorCode.PROVIDER_ERROR,
+                code=WorkflowErrorCode.PROVIDER_ERROR,
             ),
         )
         in outputs

vellum/workflows/nodes/utils.py

@@ -2,7 +2,6 @@ from functools import cache
 from typing import Type
 
 from vellum.workflows.nodes import BaseNode
-from vellum.workflows.references import NodeReference
 from vellum.workflows.types.generics import NodeType
 
 ADORNMENT_MODULE_NAME = "<adornment>"
@@ -10,20 +9,17 @@ ADORNMENT_MODULE_NAME = "<adornment>"
 
 @cache
 def get_wrapped_node(node: Type[NodeType]) -> Type[BaseNode]:
-    if hasattr(node, "subworkflow"):
-        subworkflow = node.subworkflow
-        if isinstance(subworkflow, NodeReference) and subworkflow.instance:
-            graph = subworkflow.instance.graph
-            if issubclass(graph, BaseNode):
-                return graph
+    wrapped_node = getattr(node, "__wrapped_node__", None)
+    if wrapped_node is None:
+        raise AttributeError("Wrapped node not found")
 
-    raise TypeError("Wrapped subworkflow contains more than one node")
+    return wrapped_node
 
 
 def has_wrapped_node(node: Type[NodeType]) -> bool:
     try:
         get_wrapped_node(node)
-    except TypeError:
+    except AttributeError:
         return False
 
     return True

vellum/workflows/references/external_input.py

@@ -5,7 +5,7 @@ from pydantic_core import core_schema
 
 from vellum.workflows.constants import UNDEF
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 
 if TYPE_CHECKING:
@@ -40,7 +40,7 @@ class ExternalInputReference(BaseDescriptor[_InputType], Generic[_InputType]):
         if state.meta.parent:
             return self.resolve(state.meta.parent)
 
-        raise NodeException(f"Missing required Node Input: {self._name}", code=VellumErrorCode.INVALID_INPUTS)
+        raise NodeException(f"Missing required Node Input: {self._name}", code=WorkflowErrorCode.INVALID_INPUTS)
 
     @classmethod
     def __get_pydantic_core_schema__(

vellum/workflows/references/node.py

@@ -4,7 +4,7 @@ from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 
 if TYPE_CHECKING:
@@ -30,7 +30,7 @@ class NodeReference(BaseDescriptor[_T]):
     def resolve(self, state: "BaseState") -> _T:
         raise NodeException(
             f"NodeDescriptors cannot be resolved during runtime. Got: {self._name}",
-            code=VellumErrorCode.INTERNAL_ERROR,
+            code=WorkflowErrorCode.INTERNAL_ERROR,
         )
 
     def __repr__(self) -> str:

vellum/workflows/references/state_value.py

@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING, TypeVar, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 
 if TYPE_CHECKING:
@@ -20,4 +20,4 @@ class StateValueReference(BaseDescriptor[_T]):
         if state.meta.parent:
             return self.resolve(state.meta.parent)
 
-        raise NodeException(f"Missing required Workflow state: {self._name}", code=VellumErrorCode.INVALID_STATE)
+        raise NodeException(f"Missing required Workflow state: {self._name}", code=WorkflowErrorCode.INVALID_STATE)

vellum/workflows/references/workflow_input.py

@@ -1,7 +1,7 @@
 from typing import TYPE_CHECKING, Generic, Optional, Tuple, Type, TypeVar, cast
 
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 
 if TYPE_CHECKING:
@@ -35,7 +35,7 @@ class WorkflowInputReference(BaseDescriptor[_InputType], Generic[_InputType]):
         if state.meta.parent:
             return self.resolve(state.meta.parent)
 
-        raise NodeException(f"Missing required Workflow input: {self._name}", code=VellumErrorCode.INVALID_INPUTS)
+        raise NodeException(f"Missing required Workflow input: {self._name}", code=WorkflowErrorCode.INVALID_INPUTS)
 
     def __repr__(self) -> str:
         return f"{self._inputs_class.__qualname__}.{self.name}"