vellum-ai 0.10.4__py3-none-any.whl → 0.10.6__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -17,7 +17,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.10.4",
+            "X-Fern-SDK-Version": "0.10.6",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/workflows/nodes/__init__.py

@@ -1,5 +1,5 @@
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.core import (ErrorNode, InlineSubworkflowNode, MapNode, RetryNode, TemplatingNode, TryNode,)
+from vellum.workflows.nodes.core import ErrorNode, InlineSubworkflowNode, MapNode, RetryNode, TemplatingNode, TryNode
 from vellum.workflows.nodes.displayable import (
     APINode,
     CodeExecutionNode,
@@ -7,6 +7,7 @@ from vellum.workflows.nodes.displayable import (
     FinalOutputNode,
     GuardrailNode,
     InlinePromptNode,
+    NoteNode,
     PromptDeploymentNode,
     SearchNode,
     SubworkflowDeploymentNode,
@@ -28,20 +29,18 @@ __all__ = [
     "TemplatingNode",
     "TryNode",
     # Displayable Base Nodes
-    "BaseSearchNode",
     "BaseInlinePromptNode",
     "BasePromptDeploymentNode",
+    "BaseSearchNode",
     # Displayable Nodes
     "APINode",
     "CodeExecutionNode",
+    "ConditionalNode",
+    "FinalOutputNode",
     "GuardrailNode",
     "InlinePromptNode",
+    "NoteNode",
     "PromptDeploymentNode",
     "SearchNode",
-    "ConditionalNode",
-    "GuardrailNode",
     "SubworkflowDeploymentNode",
-    "FinalOutputNode",
-    "PromptDeploymentNode",
-    "SearchNode",
 ]

vellum/workflows/nodes/bases/base.py

@@ -215,7 +215,6 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
     #   https://app.shortcut.com/vellum/story/4008/auto-inherit-basenodeoutputs-in-outputs-classes
     class Outputs(BaseOutputs):
         _node_class: Optional[Type["BaseNode"]] = None
-        pass
 
     class Ports(NodePorts):
         default = Port(default=True)

vellum/workflows/nodes/core/inline_subworkflow_node/node.py

@@ -57,7 +57,7 @@ class InlineSubworkflowNode(BaseSubworkflowNode[StateType], Generic[StateType, W
         if outputs is None:
             raise NodeException(
                 message="Expected to receive outputs from Workflow Deployment",
-                code=VellumErrorCode.INTERNAL_ERROR,
+                code=VellumErrorCode.INVALID_OUTPUTS,
             )
 
         # For any outputs somehow in our final fulfilled outputs array,

vellum/workflows/nodes/core/templating_node/node.py

@@ -49,7 +49,11 @@ class _TemplatingNodeMeta(BaseNodeMeta):
         if not isinstance(parent, _TemplatingNodeMeta):
             raise ValueError("TemplatingNode must be created with the TemplatingNodeMeta metaclass")
 
-        parent.__dict__["Outputs"].__annotations__["result"] = parent.get_output_type()
+        annotations = parent.__dict__["Outputs"].__annotations__
+        parent.__dict__["Outputs"].__annotations__ = {
+            **annotations,
+            "result": parent.get_output_type(),
+        }
         return parent
 
     def get_output_type(cls) -> Type:

vellum/workflows/nodes/core/try_node/node.py

@@ -1,10 +1,13 @@
-from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Optional, Tuple, Type, TypeVar
+import sys
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, Iterator, Optional, Set, Tuple, Type, TypeVar, cast
 
 from vellum.workflows.errors.types import VellumError, VellumErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
-from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
+from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
 from vellum.workflows.types.generics import StateType
 
 if TYPE_CHECKING:
@@ -56,34 +59,60 @@ class TryNode(BaseNode[StateType], Generic[StateType], metaclass=_TryNodeMeta):
     class Outputs(BaseNode.Outputs):
         error: Optional[VellumError] = None
 
-    def run(self) -> Outputs:
+    def run(self) -> Iterator[BaseOutput]:
         subworkflow = self.subworkflow(
             parent_state=self.state,
             context=self._context,
         )
-        terminal_event = subworkflow.run()
-
-        if terminal_event.name == "workflow.execution.fulfilled":
-            outputs = self.Outputs()
-            for descriptor, value in terminal_event.outputs:
-                setattr(outputs, descriptor.name, value)
-            return outputs
-        elif terminal_event.name == "workflow.execution.paused":
+        subworkflow_stream = subworkflow.stream()
+
+        outputs: Optional[BaseOutputs] = None
+        exception: Optional[NodeException] = None
+        fulfilled_output_names: Set[str] = set()
+
+        for event in subworkflow_stream:
+            if exception:
+                continue
+
+            if event.name == "workflow.execution.streaming":
+                if event.output.is_fulfilled:
+                    fulfilled_output_names.add(event.output.name)
+                yield event.output
+            elif event.name == "workflow.execution.fulfilled":
+                outputs = event.outputs
+            elif event.name == "workflow.execution.paused":
+                exception = NodeException(
+                    code=VellumErrorCode.INVALID_OUTPUTS,
+                    message="Subworkflow unexpectedly paused within Try Node",
+                )
+            elif event.name == "workflow.execution.rejected":
+                if self.on_error_code and self.on_error_code != event.error.code:
+                    exception = NodeException(
+                        code=VellumErrorCode.INVALID_OUTPUTS,
+                        message=f"""Unexpected rejection: {event.error.code.value}.
+Message: {event.error.message}""",
+                    )
+                else:
+                    outputs = self.Outputs(error=event.error)
+
+        if exception:
+            raise exception
+
+        if outputs is None:
             raise NodeException(
                 code=VellumErrorCode.INVALID_OUTPUTS,
-                message="Subworkflow unexpectedly paused within Try Node",
-            )
-        elif self.on_error_code and self.on_error_code != terminal_event.error.code:
-            raise NodeException(
-                code=VellumErrorCode.INVALID_OUTPUTS,
-                message=f"""Unexpected rejection: {terminal_event.error.code.value}.
-Message: {terminal_event.error.message}""",
-            )
-        else:
-            return self.Outputs(
-                error=terminal_event.error,
+                message="Expected to receive outputs from Try Node's subworkflow",
             )
 
+        # For any outputs somehow in our final fulfilled outputs array,
+        # but not fulfilled by the stream.
+        for descriptor, value in outputs:
+            if descriptor.name not in fulfilled_output_names:
+                yield BaseOutput(
+                    name=descriptor.name,
+                    value=value,
+                )
+
     @classmethod
     def wrap(cls, on_error_code: Optional[VellumErrorCode] = None) -> Callable[..., Type["TryNode"]]:
         _on_error_code = on_error_code
@@ -101,11 +130,20 @@ Message: {terminal_event.error.message}""",
                 class Outputs(inner_cls.Outputs):  # type: ignore[name-defined]
                     pass
 
-            class WrappedNode(TryNode[StateType]):
-                on_error_code = _on_error_code
-
-                subworkflow = Subworkflow
-
+            dynamic_module = f"{inner_cls.__module__}.{inner_cls.__name__}.{ADORNMENT_MODULE_NAME}"
+            # This dynamic module allows calls to `type_hints` to work
+            sys.modules[dynamic_module] = ModuleType(dynamic_module)
+
+            # We use a dynamic wrapped node class to be uniquely tied to this `inner_cls` node during serialization
+            WrappedNode = type(
+                cls.__name__,
+                (TryNode,),
+                {
+                    "__module__": dynamic_module,
+                    "on_error_code": _on_error_code,
+                    "subworkflow": Subworkflow,
+                },
+            )
             return WrappedNode
 
         return decorator

vellum/workflows/nodes/core/try_node/tests/test_node.py

@@ -7,6 +7,7 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.outputs import BaseOutputs
+from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.state.base import BaseState, StateMeta
 from vellum.workflows.state.context import WorkflowContext
 
@@ -23,11 +24,15 @@ def test_try_node__on_error_code__successfully_caught():
 
     # WHEN the node is run and throws a PROVIDER_ERROR
     node = TestNode(state=BaseState())
-    outputs = node.run()
-
-    # THEN the exception is retried
-    assert outputs == {
-        "error": VellumError(message="This will be caught", code=VellumErrorCode.PROVIDER_ERROR),
+    outputs = [o for o in node.run()]
+
+    # THEN the exception is caught and returned
+    assert len(outputs) == 2
+    assert set(outputs) == {
+        BaseOutput(name="value"),
+        BaseOutput(
+            name="error", value=VellumError(message="This will be caught", code=VellumErrorCode.PROVIDER_ERROR)
+        ),
     }
 
 
@@ -44,7 +49,7 @@ def test_try_node__retry_on_error_code__missed():
     # WHEN the node is run and throws a different exception
     node = TestNode(state=BaseState())
     with pytest.raises(NodeException) as exc_info:
-        node.run()
+        list(node.run())
 
     # THEN the exception is not caught
     assert exc_info.value.message == "Unexpected rejection: INTERNAL_ERROR.\nMessage: This will be missed"
@@ -78,10 +83,11 @@ def test_try_node__use_parent_inputs_and_state():
             meta=StateMeta(workflow_inputs=Inputs(foo="foo")),
         ),
     )
-    outputs = node.run()
+    outputs = list(node.run())
 
     # THEN the data is used successfully
-    assert outputs == {"value": "foo bar"}
+    assert len(outputs) == 1
+    assert outputs[-1] == BaseOutput(name="value", value="foo bar")
 
 
 def test_try_node__use_parent_execution_context():
@@ -100,7 +106,8 @@ def test_try_node__use_parent_execution_context():
             _vellum_client=Vellum(api_key="test-key"),
         )
     )
-    outputs = node.run()
+    outputs = list(node.run())
 
     # THEN the inner node had access to the key
-    assert outputs == {"key": "test-key"}
+    assert len(outputs) == 1
+    assert outputs[-1] == BaseOutput(name="key", value="test-key")

vellum/workflows/nodes/displayable/__init__.py

@@ -9,6 +9,7 @@ from .final_output_node import FinalOutputNode
 from .guardrail_node import GuardrailNode
 from .inline_prompt_node import InlinePromptNode
 from .merge_node import MergeNode
+from .note_node import NoteNode
 from .prompt_deployment_node import PromptDeploymentNode
 from .search_node import SearchNode
 from .subworkflow_deployment_node import SubworkflowDeploymentNode
@@ -23,6 +24,7 @@ __all__ = [
     "GuardrailNode",
     "MapNode",
     "MergeNode",
+    "NoteNode",
     "SubworkflowDeploymentNode",
     "PromptDeploymentNode",
     "SearchNode",

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

@@ -8,7 +8,7 @@ from vellum.workflows.errors.types import VellumErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
-from vellum.workflows.types.core import JsonObject, VellumSecret
+from vellum.workflows.types.core import Json, JsonObject, VellumSecret
 from vellum.workflows.types.generics import StateType
 
 
@@ -26,11 +26,11 @@ class BaseAPINode(BaseNode, Generic[StateType]):
     url: str
     method: APIRequestMethod
     data: Optional[str] = None
-    json: Optional["JsonObject"] = None
+    json: Optional["Json"] = None
     headers: Optional[Dict[str, Union[str, VellumSecret]]] = None
 
     class Outputs(BaseOutputs):
-        json: Optional["JsonObject"]
+        json: Optional["Json"]
         headers: Dict[str, str]
         status_code: int
         text: str

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

@@ -19,7 +19,6 @@ from vellum import (
     VellumValue,
 )
 from vellum.core import RequestOptions
-
 from vellum.workflows.errors.types import VellumErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
@@ -44,7 +43,11 @@ class _CodeExecutionNodeMeta(BaseNodeMeta):
         if not isinstance(parent, _CodeExecutionNodeMeta):
             raise ValueError("CodeExecutionNode must be created with the CodeExecutionNodeMeta metaclass")
 
-        parent.__dict__["Outputs"].__annotations__["result"] = parent.get_output_type()
+        annotations = parent.__dict__["Outputs"].__annotations__
+        parent.__dict__["Outputs"].__annotations__ = {
+            **annotations,
+            "result": parent.get_output_type(),
+        }
         return parent
 
     def get_output_type(cls) -> Type:

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

@@ -16,9 +16,13 @@ class _FinalOutputNodeMeta(BaseNodeMeta):
 
         # We use the compiled class to infer the output type for the Outputs.value descriptor.
         if not isinstance(parent, _FinalOutputNodeMeta):
-            raise ValueError("CodeExecutionNode must be created with the CodeExecutionNodeMeta metaclass")
+            raise ValueError("FinalOutputNode must be created with the FinalOutputNodeMeta metaclass")
 
-        parent.__dict__["Outputs"].__annotations__["value"] = parent.get_output_type()
+        annotations = parent.__dict__["Outputs"].__annotations__
+        parent.__dict__["Outputs"].__annotations__ = {
+            **annotations,
+            "value": parent.get_output_type(),
+        }
         return parent
 
     def get_output_type(cls) -> Type:

vellum/workflows/nodes/displayable/note_node/__init__.py

@@ -0,0 +1,5 @@
+from .node import NoteNode
+
+__all__ = [
+    "NoteNode",
+]

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

@@ -0,0 +1,10 @@
+from vellum.workflows.nodes.bases import BaseNode
+
+
+class NoteNode(BaseNode):
+    """
+    A no-op Node purely used to display a note in the Vellum UI.
+    """
+
+    def run(self) -> BaseNode.Outputs:
+        raise RuntimeError("NoteNode should never be run")

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

@@ -11,13 +11,12 @@ from vellum import (
     StringVellumValue,
     VellumError,
 )
-
-from vellum.workflows.constants import UNDEF
-from vellum.workflows.errors import VellumError as WacVellumError
+from vellum.workflows.errors import VellumError as SdkVellumError
 from vellum.workflows.errors.types import VellumErrorCode
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
+from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.state import BaseState
 from vellum.workflows.state.base import StateMeta
 
@@ -136,13 +135,13 @@ def test_inline_text_prompt_node__catch_provider_error(vellum_adhoc_prompt_clien
             meta=StateMeta(workflow_inputs=Inputs(input="Say something.")),
         )
     )
-    outputs = node.run()
+    outputs = list(node.run())
 
     # THEN the node should have produced the outputs we expect
-    # We need mypy support for annotations to remove these type ignores
-    # https://app.shortcut.com/vellum/story/4890
-    assert outputs.error == WacVellumError(  # type: ignore[attr-defined]
-        message="OpenAI failed",
-        code=VellumErrorCode.PROVIDER_ERROR,
-    )
-    assert outputs.text is UNDEF  # type: ignore[attr-defined]
+    assert BaseOutput(
+        name="error",
+        value=SdkVellumError(
+            message="OpenAI failed",
+            code=VellumErrorCode.PROVIDER_ERROR,
+        ),
+    ) in outputs

vellum/workflows/nodes/utils.py

@@ -5,6 +5,8 @@ from vellum.workflows.nodes import BaseNode
 from vellum.workflows.references import NodeReference
 from vellum.workflows.types.generics import NodeType
 
+ADORNMENT_MODULE_NAME = "<adornment>"
+
 
 @cache
 def get_wrapped_node(node: Type[NodeType]) -> Type[BaseNode]:

vellum/workflows/outputs/base.py

@@ -5,6 +5,7 @@ from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 
 from vellum.workflows.constants import UNDEF
+from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.utils import get_class_attr_names, infer_types
 
@@ -76,6 +77,23 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
 
         return data
 
+    def __repr__(self) -> str:
+        if self.value is not UNDEF:
+            return f"{self.__class__.__name__}({self.name}={self.value})"
+        elif self.delta is not UNDEF:
+            return f"{self.__class__.__name__}({self.name}={self.delta})"
+        else:
+            return f"{self.__class__.__name__}(name='{self.name}')"
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, BaseOutput):
+            return False
+
+        return self.name == other.name and self.value == other.value and self.delta == other.delta
+    
+    def __hash__(self) -> int:
+        return hash((self._name, self._value, self._value))
+
 
 @dataclass_transform(kw_only_default=True)
 class _BaseOutputsMeta(type):
@@ -175,7 +193,9 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
         if not isinstance(other, dict):
             return super().__eq__(other)
 
-        outputs = {name: value for name, value in vars(self).items() if not name.startswith("_") and value is not UNDEF}
+        outputs = {
+            name: value for name, value in vars(self).items() if not name.startswith("_") and value is not UNDEF
+        }
         return outputs == other
 
     def __repr__(self) -> str:
@@ -184,7 +204,11 @@ class BaseOutputs(metaclass=_BaseOutputsMeta):
 
     def __iter__(self) -> Iterator[Tuple[OutputReference, Any]]:
         for output_descriptor in self.__class__:
-            yield (output_descriptor, getattr(self, output_descriptor.name, output_descriptor.instance))
+            output_value = getattr(self, output_descriptor.name, UNDEF)
+            if isinstance(output_value, BaseDescriptor):
+                output_value = UNDEF
+
+            yield (output_descriptor, output_value)
 
     def __getitem__(self, key: str) -> Any:
         return getattr(self, key)

vellum/workflows/runner/runner.py

@@ -170,32 +170,37 @@ class WorkflowRunner(Generic[StateType]):
                 streaming_output_queues: Dict[str, Queue] = {}
                 outputs = node.Outputs()
 
+                def initiate_node_streaming_output(output: BaseOutput) -> None:
+                    streaming_output_queues[output.name] = Queue()
+                    output_descriptor = OutputReference(
+                        name=output.name,
+                        types=(type(output.delta),),
+                        instance=None,
+                        outputs_class=node.Outputs,
+                    )
+                    node.state.meta.node_outputs[output_descriptor] = streaming_output_queues[output.name]
+                    self._work_item_event_queue.put(
+                        WorkItemEvent(
+                            node=node,
+                            event=NodeExecutionStreamingEvent(
+                                trace_id=node.state.meta.trace_id,
+                                span_id=span_id,
+                                body=NodeExecutionStreamingBody(
+                                    node_definition=node.__class__,
+                                    output=BaseOutput(name=output.name),
+                                ),
+                            ),
+                            invoked_ports=invoked_ports,
+                        )
+                    )
+
                 for output in node_run_response:
                     invoked_ports = output > ports
-                    if not output.is_fulfilled:
+                    if output.is_initiated:
+                        initiate_node_streaming_output(output)
+                    elif output.is_streaming:
                         if output.name not in streaming_output_queues:
-                            streaming_output_queues[output.name] = Queue()
-                            output_descriptor = OutputReference(
-                                name=output.name,
-                                types=(type(output.delta),),
-                                instance=None,
-                                outputs_class=node.Outputs,
-                            )
-                            node.state.meta.node_outputs[output_descriptor] = streaming_output_queues[output.name]
-                            self._work_item_event_queue.put(
-                                WorkItemEvent(
-                                    node=node,
-                                    event=NodeExecutionStreamingEvent(
-                                        trace_id=node.state.meta.trace_id,
-                                        span_id=span_id,
-                                        body=NodeExecutionStreamingBody(
-                                            node_definition=node.__class__,
-                                            output=BaseOutput(name=output.name),
-                                        ),
-                                    ),
-                                    invoked_ports=invoked_ports,
-                                )
-                            )
+                            initiate_node_streaming_output(output)
 
                         streaming_output_queues[output.name].put(output.delta)
                         self._work_item_event_queue.put(
@@ -212,7 +217,7 @@ class WorkflowRunner(Generic[StateType]):
                                 invoked_ports=invoked_ports,
                             )
                         )
-                    else:
+                    elif output.is_fulfilled:
                         if output.name in streaming_output_queues:
                             streaming_output_queues[output.name].put(UNDEF)
 
@@ -233,6 +238,11 @@ class WorkflowRunner(Generic[StateType]):
                         )
 
             for descriptor, output_value in outputs:
+                if output_value is UNDEF:
+                    if descriptor in node.state.meta.node_outputs:
+                        del node.state.meta.node_outputs[descriptor]
+                    continue
+
                 node.state.meta.node_outputs[descriptor] = output_value
 
             invoked_ports = ports(outputs, node.state)
@@ -540,11 +550,15 @@ class WorkflowRunner(Generic[StateType]):
         )
 
     def stream(self) -> WorkflowEventStream:
-        background_thread = Thread(target=self._run_background_thread)
+        background_thread = Thread(
+            target=self._run_background_thread, name=f"{self.workflow.__class__.__name__}.background_thread"
+        )
         background_thread.start()
 
         if self._cancel_signal:
-            cancel_thread = Thread(target=self._run_cancel_thread)
+            cancel_thread = Thread(
+                target=self._run_cancel_thread, name=f"{self.workflow.__class__.__name__}.cancel_thread"
+            )
             cancel_thread.start()
 
         event: WorkflowEvent
@@ -557,7 +571,7 @@ class WorkflowRunner(Generic[StateType]):
         self._initial_state.meta.is_terminated = False
 
         # The extra level of indirection prevents the runner from waiting on the caller to consume the event stream
-        stream_thread = Thread(target=self._stream)
+        stream_thread = Thread(target=self._stream, name=f"{self.workflow.__class__.__name__}.stream_thread")
         stream_thread.start()
 
         while stream_thread.is_alive():

vellum/workflows/types/tests/test_utils.py

@@ -1,6 +1,8 @@
 import pytest
 from typing import ClassVar, Generic, List, TypeVar, Union
 
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.utils import get_class_attr_names, infer_types
@@ -30,6 +32,11 @@ class ExampleGenericClass(Generic[T]):
 class ExampleInheritedClass(ExampleClass):
     theta: int
 
+@TryNode.wrap()
+class ExampleNode(BaseNode):
+    class Outputs(BaseNode.Outputs):
+        iota: str
+
 
 @pytest.mark.parametrize(
     "cls, attr_name, expected_type",
@@ -45,6 +52,7 @@ class ExampleInheritedClass(ExampleClass):
         (ExampleInheritedClass, "theta", (int,)),
         (ExampleInheritedClass, "alpha", (str,)),
         (ExampleInheritedClass, "beta", (int,)),
+        (ExampleNode.Outputs, "iota", (str,)),
     ],
     ids=[
         "str",
@@ -58,6 +66,7 @@ class ExampleInheritedClass(ExampleClass):
         "inherited_int",
         "inherited_parent_annotation",
         "inherited_parent_class_var",
+        "try_node_output",
     ],
 )
 def test_infer_types(cls, attr_name, expected_type):

vellum/workflows/types/utils.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 from datetime import datetime
 import importlib
+import sys
 from typing import (
     Any,
     ClassVar,
@@ -18,7 +19,6 @@ from typing import (
 )
 
 from vellum import ArrayVellumValue, ArrayVellumValueRequest, ChatMessagePromptBlock
-
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.types.core import Json, SpecialGenericAlias, UnderGenericAlias, UnionGenericAlias
 

vellum/workflows/utils/vellum_variables.py

@@ -1,3 +1,4 @@
+import typing
 from typing import List, Tuple, Type, Union, get_args, get_origin
 
 from vellum import (
@@ -17,8 +18,8 @@ from vellum import (
     VellumValueRequest,
     VellumVariableType,
 )
-
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.types.core import Json
 
 
 def primitive_type_to_vellum_variable_type(type_: Union[Type, BaseDescriptor]) -> VellumVariableType:
@@ -32,6 +33,17 @@ def primitive_type_to_vellum_variable_type(type_: Union[Type, BaseDescriptor]) -
             return "JSON"
 
         if len(types) != 1:
+            # Check explicitly for our internal JSON type.
+            # Matches the type found at vellum.workflows.utils.vellum_variables.Json
+            if types == [
+                bool,
+                int,
+                float,
+                str,
+                typing.List[typing.ForwardRef('Json')],  # type: ignore [misc]
+                typing.Dict[str, typing.ForwardRef('Json')],  # type: ignore [misc]
+            ]:
+                return "JSON"
             raise ValueError(f"Expected Descriptor to only have one type, got {types}")
 
         type_ = type_.types[0]