vellum-ai 0.13.8__py3-none-any.whl → 0.13.9__py3-none-any.whl

vellum/client/core/client_wrapper.py

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

vellum/plugins/vellum_mypy.py

@@ -19,6 +19,7 @@ from mypy.plugin import AttributeContext, ClassDefContext, FunctionSigContext, M
 from mypy.types import AnyType, CallableType, FunctionLike, Instance, Type as MypyType, TypeAliasType, UnionType
 
 TypeResolver = Callable[[str, List[MypyType]], MypyType]
+NODE_ATTRIBUTE_REGEX = r"^[a-z].*$"
 
 DESCRIPTOR_PATHS: list[tuple[str, str, str]] = [
     (
@@ -49,7 +50,7 @@ DESCRIPTOR_PATHS: list[tuple[str, str, str]] = [
     (
         "vellum.workflows.nodes.bases.base.BaseNode",
         "vellum.workflows.references.node.NodeReference",
-        r"^[a-z].*$",
+        NODE_ATTRIBUTE_REGEX,
     ),
 ]
 
@@ -239,10 +240,13 @@ class VellumMypyPlugin(Plugin):
         the original defined type, and the descriptor version of the type.
         """
 
-        for sym in ctx.cls.info.names.values():
+        for key, sym in ctx.cls.info.names.items():
             if not isinstance(sym.node, Var):
                 continue
 
+            if not re.match(NODE_ATTRIBUTE_REGEX, key):
+                continue
+
             type_ = sym.node.type
             if not type_:
                 continue

vellum/workflows/nodes/bases/base.py

@@ -222,7 +222,6 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
     state: StateType
     _context: WorkflowContext
     _inputs: MappingProxyType[NodeReference, Any]
-    _is_wrapped_node: bool = False
 
     class ExternalInputs(BaseInputs):
         __descriptor_class__ = ExternalInputReference

vellum/workflows/nodes/utils.py

@@ -4,6 +4,7 @@ from types import ModuleType
 from typing import Any, Callable, Optional, Type, TypeVar
 
 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.types.generics import NodeType
 
@@ -12,7 +13,7 @@ ADORNMENT_MODULE_NAME = "<adornment>"
 
 @cache
 def get_unadorned_node(node: Type[BaseNode]) -> Type[BaseNode]:
-    wrapped_node = getattr(node, "__wrapped_node__", None)
+    wrapped_node = get_wrapped_node(node)
     if wrapped_node is not None:
         return get_unadorned_node(wrapped_node)
 
@@ -28,22 +29,11 @@ def get_unadorned_port(port: Port) -> Port:
     return getattr(unadorned_node.Ports, port.name)
 
 
-@cache
-def get_wrapped_node(node: Type[NodeType]) -> Type[BaseNode]:
-    wrapped_node = getattr(node, "__wrapped_node__", None)
-    if wrapped_node is None:
-        raise AttributeError("Wrapped node not found")
-
-    return wrapped_node
-
+def get_wrapped_node(node: Type[NodeType]) -> Optional[Type[BaseNode]]:
+    if not issubclass(node, BaseAdornmentNode):
+        return None
 
-def has_wrapped_node(node: Type[NodeType]) -> bool:
-    try:
-        get_wrapped_node(node)
-    except AttributeError:
-        return False
-
-    return True
+    return node.__wrapped_node__
 
 
 AdornableNode = TypeVar("AdornableNode", bound=BaseNode)
@@ -57,8 +47,6 @@ def create_adornment(
         # https://app.shortcut.com/vellum/story/4116
         from vellum.workflows import BaseWorkflow
 
-        inner_cls._is_wrapped_node = True
-
         class Subworkflow(BaseWorkflow):
             graph = inner_cls
 

{vellum_ai-0.13.8.dist-info → vellum_ai-0.13.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.13.8
+Version: 0.13.9
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.13.8.dist-info → vellum_ai-0.13.9.dist-info}/RECORD

@@ -21,7 +21,7 @@ vellum_ee/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 vellum_ee/workflows/display/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/base.py,sha256=3ZFUYRNKL24fBqXhKpa_Dq2W1a-a86J20dmJYA3H2eY,1755
 vellum_ee/workflows/display/nodes/__init__.py,sha256=5XOcZJXYUgaLS55QgRJzyQ_W1tpeprjnYAeYVezqoGw,160
-vellum_ee/workflows/display/nodes/base_node_display.py,sha256=pQ-QlaDdxME3roRrQj5eYRZwZXuE9im8gjo2QFwn954,6942
+vellum_ee/workflows/display/nodes/base_node_display.py,sha256=Ge3f_MNCxawRNR2EhBvygXTIWxKOpS40Dkgrp1S1WPQ,6991
 vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=BxA-YVUJvB36NM-Q5DNCwckeqymwLKMp9DVCaTyrPbs,2253
 vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=vyKeJAevAXvEAEtWeTEdBZXD3eJQYW_DEXLKVZ5KmYc,1325
 vellum_ee/workflows/display/nodes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -46,9 +46,10 @@ vellum_ee/workflows/display/nodes/vellum/search_node.py,sha256=TxcAGZDl_hvJ7Y1hU
 vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=lfevlHpGEX14dEDym6qmnkw3nvzQPTB1_D2ch12B_Rk,2701
 vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=JVIMPR3WpveOCWZubHKZkE04mavnTdb_9QY_r3XliRg,3424
 vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=Nqd6mxn0sgL4sp62wDYnQVc7COxrt5wXXIveRFoYQ8c,3913
-vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=hB8dcGMGkuC95kk9hmZUgHsCLwEA37fHTFXj0JzbRjM,4692
-vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=nIZ1DYTYPRaYsVKcel9a-Fm8EniJL0N7f5PowxVGTVU,8318
+vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py,sha256=mtzB8LJlFCHVFM4H5AanLp29gQfaVmnN4A4iaRGJHoI,2427
+vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=3uT7Gbc0f_mQ3u8uZuCWd0mJ4GtWbz2gbUMySYaVlNE,3774
+vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=Ywa8NMNvE4TlEZ_gQA6jhrqykH8-wWgKFa-uEokWhF8,5550
+vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=TFChHlSN8B0MOdgFTDmQPKkqwGGYxeF5lNmFp-JXzqo,7882
 vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=4jXe7Wyspw6CxghVqKAOu-_QunKFvY4U6--zOiuXvV0,3069
 vellum_ee/workflows/display/tests/workflow_serialization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -60,7 +61,7 @@ vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outp
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=wZekdhIZBwbBoFNAdC9bBLwUxVKk3EnFdNmMwYIdKGA,37308
 vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py,sha256=os6wCZjcOyig8EJu-pTAFWXa0odMxTaR0mrbL0SS_4Y,4480
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py,sha256=bXZWxOKAVjZlbP3iLHPHGA4aPs0EguKlQqmYPNGv3cU,16391
-vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=UPLEQPwsLzOqZdkXrB5Jo1FdCx0iQNf7ekfcq1byoFw,29392
+vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=uvABhpstyxNNOz50-XJMAr3SKp8gluPp1mUtJjvRL0c,29410
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py,sha256=Xn8t-SpvQocKdxNofDeDETlFnkCbBxdLGiUG8m6fM6U,48399
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=QVlkczxzaKuOhwbRvVP70otPDNnJcSGDfN79j92lFyk,5534
 vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py,sha256=_gv7vPxBWSOSRKMlXYv8GKj9h1JAXjXIVWkCE_XhkYE,5746
@@ -80,9 +81,9 @@ vellum_ee/workflows/display/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeR
 vellum_ee/workflows/display/utils/vellum.py,sha256=N5VSDYdtHUxvTeSj4zA8aMiAAKn4bAVgpKUDV_obNQ8,3632
 vellum_ee/workflows/display/vellum.py,sha256=8xXRI8b8Tt661H-iZreTQTvLNEKUr4lf-XaKhE7_yUY,8147
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=cLmW-oMMaP2BSwXCTVJ_KxhKmyfTRRJ_TGhkW3AQYXA,14324
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=VNYc9n_TvsYxtJ__BaTI9ZGEzTg8JClrZ-nVnFlgiBU,14051
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
-vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=V0edhtohqAWbaHvHkj9Sth4ieaIVejsrsRIr7aCCoVc,16871
+vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=ZGPhfb-4_7bLEYaKnESkZuwWrzh9lmT2wMq2jPQ-ciQ,16518
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
 vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -111,7 +112,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=8nZt88C9SVwWanjLbIQMU3rzb32h5UZfFMBx3VPHB50,111887
 vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
 vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
-vellum/client/core/client_wrapper.py,sha256=NnIgNorh26bh1LnC042Vcmf_kylY4hdObVYLpA5mq6k,1868
+vellum/client/core/client_wrapper.py,sha256=dxLror26l68ej4_dKgIfKbRZE9rDH7pYhicRFBkhKss,1868
 vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
 vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
 vellum/client/core/http_client.py,sha256=R0pQpCppnEtxccGvXl4uJ76s7ro_65Fo_erlNNLp_AI,19228
@@ -691,7 +692,7 @@ vellum/evaluations/utils/paginator.py,sha256=rEED_BJAXAM6tM1yMwHePNzszjq_tTq4NbQ
 vellum/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/plugins/pydantic.py,sha256=dNtZWHo-IdseG52C2RoTanxyTJg0AhPZrH-9lbNqwYg,2604
 vellum/plugins/utils.py,sha256=U9ZY9KdE3RRvbcG01hXxu9CvfJD6Fo7nJDgcHjQn0FI,606
-vellum/plugins/vellum_mypy.py,sha256=7JmyuTX723-XWnWoE6afiUNOkHyAufCUZwxck9BP2aI,27784
+vellum/plugins/vellum_mypy.py,sha256=QTuMSq6PiZW1dyTUZ5Bf1d4XkgFj0TKAgZLP8f4UgL4,27914
 vellum/prompts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/prompts/blocks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/prompts/blocks/compilation.py,sha256=ISuvDHaeVCPb1L7l4umORCECkDn0-rvE49hopz6c2gM,9222
@@ -1305,7 +1306,7 @@ vellum/workflows/inputs/base.py,sha256=1kMgr0WqCYdWUqgFvgSoAMw2067FAlgwhGXLgbIOr
 vellum/workflows/logging.py,sha256=_a217XogktV4Ncz6xKFz7WfYmZAzkfVRVuC0rWob8ls,437
 vellum/workflows/nodes/__init__.py,sha256=aVdQVv7Y3Ro3JlqXGpxwaU2zrI06plDHD2aumH5WUIs,1157
 vellum/workflows/nodes/bases/__init__.py,sha256=cniHuz_RXdJ4TQgD8CBzoiKDiPxg62ErdVpCbWICX64,58
-vellum/workflows/nodes/bases/base.py,sha256=CcWQBMR3cx5vftqQp_oJ_GncULJIbOLMyNP4KZNgU10,14487
+vellum/workflows/nodes/bases/base.py,sha256=P8j70SnO5-xABMbAPCHNKLook5Ip32FpzqfYEUZG8E8,14452
 vellum/workflows/nodes/bases/base_adornment_node.py,sha256=eFTgsPCYb3eyGS0-kw7C6crFnwFx437R5wh9-8bWYts,2905
 vellum/workflows/nodes/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/bases/tests/test_base_node.py,sha256=51CueFVty9XYASC0rKr1cXWejho5WElmhfhp6cCONy0,3811
@@ -1385,7 +1386,7 @@ vellum/workflows/nodes/experimental/README.md,sha256=eF6DfIL8t-HbF9-mcofOMymKrra
 vellum/workflows/nodes/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/nodes/experimental/openai_chat_completion_node/__init__.py,sha256=lsyD9laR9p7kx5-BXGH2gUTM242UhKy8SMV0SR6S2iE,90
 vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py,sha256=1EGeiaT-Zoo6pttQFKKBcdf3dmhAbjKGaErYD5FFwlc,10185
-vellum/workflows/nodes/utils.py,sha256=chSsmKe_BsvMIJpzSxO5TWYlr3sAuxiwkfB5azkuN5Q,2715
+vellum/workflows/nodes/utils.py,sha256=tjBsootwm7vUq9qU4ttDL3CWH0C9Sd6QlX7IhfCCy34,2512
 vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
 vellum/workflows/outputs/base.py,sha256=a7W6rNSDSawwGAXYjNTF2iHb9lnZu7WFSOagZIyy__k,7976
 vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
@@ -1440,8 +1441,8 @@ vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528
 vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
 vellum/workflows/workflows/base.py,sha256=k0kUWWko4fHyCqLSU_1cBK_pXZpl9MXekWiG-bdOAo0,18353
 vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
-vellum_ai-0.13.8.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.13.8.dist-info/METADATA,sha256=DslQiu2NfpU6727w1BqOFezu5H0uKbP0EiBP1ALlB1Y,5334
-vellum_ai-0.13.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.13.8.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.13.8.dist-info/RECORD,,
+vellum_ai-0.13.9.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.13.9.dist-info/METADATA,sha256=pIFtYGqbPEQ93I_doRy1gCDYIObzB2saS0Wd_2UNfwY,5334
+vellum_ai-0.13.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.13.9.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.13.9.dist-info/RECORD,,

vellum_ee/workflows/display/nodes/base_node_display.py

@@ -76,12 +76,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
         )
         return node_definition
 
-    def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
+    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
         explicit_display = self.output_display.get(output)
         if explicit_display:
-            return explicit_display
+            return self._node, explicit_display
 
-        return NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name)
+        return (self._node, NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name))
 
     def get_node_port_display(self, port: Port) -> PortDisplay:
         overrides = self.port_displays.get(port)

vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py

@@ -0,0 +1,63 @@
+from typing import Any, Dict, cast
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
+from vellum.workflows.nodes.core.try_node.node import TryNode
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
+
+
+def test_try_node_display__serialize_with_error_output() -> None:
+    # GIVEN a Base Node wrapped with a TryNode
+    @TryNode.wrap()
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            hello = "world"
+
+    # AND a displayable node referencing
+    class OtherNode(TemplatingNode):
+        template = "{{ hello }} {{ error }}"
+        inputs = {
+            "hello": MyNode.Outputs.hello,
+            "error": MyNode.Outputs.error,
+        }
+
+    # AND a workflow referencing the two nodes
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode >> OtherNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
+    serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
+
+    # THEN the correct inputs should be serialized on the node
+    serialized_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(OtherNode.__id__)
+    )
+    hello_input_value = next(input["value"] for input in serialized_node["inputs"] if input["key"] == "hello")
+    error_input_value = next(input["value"] for input in serialized_node["inputs"] if input["key"] == "error")
+    assert hello_input_value == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "data": {
+                    "node_id": str(MyNode.__wrapped_node__.__id__),
+                    "output_id": "c8fbe459-c9ee-4639-a82d-961180cf9411",
+                },
+                "type": "NODE_OUTPUT",
+            }
+        ],
+    }
+    assert error_input_value == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "data": {
+                    "node_id": str(MyNode.__wrapped_node__.__id__),
+                    "output_id": "efe6e307-3ea4-4862-a26f-4c4416bb4537",
+                },
+                "type": "NODE_OUTPUT",
+            }
+        ],
+    }

vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -1,13 +1,12 @@
 import pytest
 from uuid import UUID, uuid4
-from typing import List, cast
+from typing import List
 
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
-from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
@@ -92,12 +91,12 @@ def test_create_node_input_value_pointer_rules(
                 entrypoint_node_display=NodeDisplayData(),
             ),
             global_workflow_input_displays={
-                cast(WorkflowInputReference, Inputs.example_workflow_input): WorkflowInputsVellumDisplayOverrides(
+                Inputs.example_workflow_input: WorkflowInputsVellumDisplayOverrides(
                     id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
                 ),
             },
             global_node_output_displays={
-                cast(OutputReference, MyNodeA.Outputs.output): (
+                MyNodeA.Outputs.output: (
                     MyNodeA,
                     NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
                 ),

vellum_ee/workflows/display/nodes/vellum/try_node.py

@@ -1,15 +1,17 @@
 from uuid import UUID
-from typing import Any, Callable, ClassVar, Generic, Optional, Type, TypeVar, cast
+from typing import Any, Callable, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
 
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
 from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME, get_wrapped_node
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.types.core import JsonObject
 from vellum.workflows.types.utils import get_original_base
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
+from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
@@ -22,9 +24,8 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
     def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
         node = self._node
 
-        try:
-            inner_node = get_wrapped_node(node)
-        except AttributeError:
+        inner_node = get_wrapped_node(node)
+        if not inner_node:
             subworkflow = raise_if_descriptor(node.subworkflow)
             if not isinstance(subworkflow.graph, type) or not issubclass(subworkflow.graph, BaseNode):
                 raise NotImplementedError(
@@ -57,6 +58,22 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
 
         return serialized_node
 
+    def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
+        inner_node = self._node.__wrapped_node__
+        if not inner_node:
+            return super().get_node_output_display(output)
+
+        node_display_class = get_node_display_class(BaseNodeVellumDisplay, inner_node)
+        node_display = node_display_class()
+        if output.name == "error":
+            return inner_node, NodeOutputDisplay(
+                id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
+                name="error",
+            )
+
+        inner_output = getattr(inner_node.Outputs, output.name)
+        return node_display.get_node_output_display(inner_output)
+
     @classmethod
     def wrap(cls, error_output_id: Optional[UUID] = None) -> Callable[..., Type["BaseTryNodeDisplay"]]:
         _error_output_id = error_output_id

vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -28,8 +28,7 @@ from vellum.workflows.expressions.not_between import NotBetweenExpression
 from vellum.workflows.expressions.not_in import NotInExpression
 from vellum.workflows.expressions.or_ import OrExpression
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
-from vellum.workflows.references import NodeReference, OutputReference
+from vellum.workflows.references import NodeReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule
@@ -56,15 +55,6 @@ def create_node_input(
     pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
 ) -> NodeInput:
     input_id = str(input_id) if input_id else str(uuid4_from_hash(f"{node_id}|{input_name}"))
-    if (
-        isinstance(value, OutputReference)
-        and value.outputs_class._node_class
-        and has_wrapped_node(value.outputs_class._node_class)
-    ):
-        wrapped_node = get_wrapped_node(value.outputs_class._node_class)
-        if wrapped_node._is_wrapped_node:
-            value = getattr(wrapped_node.Outputs, value.name)
-
     rules = create_node_input_value_pointer_rules(value, display_context, pointer_type=pointer_type)
     return NodeInput(
         id=input_id,

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py

@@ -685,8 +685,7 @@ def test_serialize_workflow__try_wrapped():
                 "display_data": {"position": {"x": 0.0, "y": 0.0}},
             },
         ],
-        final_output_nodes,
-        ignore_order=True,
+        sorted(final_output_nodes, key=lambda x: x["id"], reverse=True),
     )
 
     # AND each edge should be serialized correctly

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -11,7 +11,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
+from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.ports import Port
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.types.core import JsonObject
@@ -138,19 +138,18 @@ class BaseWorkflowDisplay(
     ):
         """This method recursively adds nodes wrapped in decorators to the node_output_displays dictionary."""
 
+        inner_node = get_wrapped_node(node)
+        if inner_node:
+            inner_node_display = self._get_node_display(inner_node)
+            self._enrich_global_node_output_displays(inner_node, inner_node_display, node_output_displays)
+
         for node_output in node.Outputs:
             if node_output in node_output_displays:
                 continue
 
-            if has_wrapped_node(node):
-                inner_node = get_wrapped_node(node)
-                if inner_node._is_wrapped_node:
-                    inner_node_display = self._get_node_display(inner_node)
-                    self._enrich_global_node_output_displays(inner_node, inner_node_display, node_output_displays)
-
             # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
             # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
-            node_output_displays[node_output] = node, node_display.get_node_output_display(node_output)
+            node_output_displays[node_output] = node_display.get_node_output_display(node_output)
 
     def _enrich_node_port_displays(
         self,
@@ -160,16 +159,15 @@ class BaseWorkflowDisplay(
     ):
         """This method recursively adds nodes wrapped in decorators to the port_displays dictionary."""
 
+        inner_node = get_wrapped_node(node)
+        if inner_node:
+            inner_node_display = self._get_node_display(inner_node)
+            self._enrich_node_port_displays(inner_node, inner_node_display, port_displays)
+
         for port in node.Ports:
             if port in port_displays:
                 continue
 
-            if has_wrapped_node(node):
-                inner_node = get_wrapped_node(node)
-                if inner_node._is_wrapped_node:
-                    inner_node_display = self._get_node_display(inner_node)
-                    self._enrich_node_port_displays(inner_node, inner_node_display, port_displays)
-
             port_displays[port] = node_display.get_node_port_display(port)
 
     def _get_node_display(self, node: Type[BaseNode]) -> NodeDisplayType:
@@ -216,13 +214,11 @@ class BaseWorkflowDisplay(
             global_node_displays[node] = node_display
 
             # Nodes wrapped in a decorator need to be in our node display dictionary for later retrieval
-            if has_wrapped_node(node):
-                inner_node = get_wrapped_node(node)
+            inner_node = get_wrapped_node(node)
+            if inner_node:
                 inner_node_display = self._get_node_display(inner_node)
-
-                if inner_node._is_wrapped_node:
-                    node_displays[inner_node] = inner_node_display
-                    global_node_displays[inner_node] = inner_node_display
+                node_displays[inner_node] = inner_node_display
+                global_node_displays[inner_node] = inner_node_display
 
             self._enrich_global_node_output_displays(node, node_display, global_node_output_displays)
             self._enrich_node_port_displays(node, node_display, port_displays)

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -7,7 +7,7 @@ from vellum.workflows.edges import Edge
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
-from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node, has_wrapped_node
+from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port
 from vellum.workflows.ports import Port
 from vellum.workflows.references import WorkflowInputReference
 from vellum.workflows.references.output import OutputReference
@@ -95,11 +95,8 @@ class VellumWorkflowDisplay(
         )
 
         # Add all the nodes in the workflow
-        for node, node_display in self.display_context.node_displays.items():
-            if getattr(node, "_is_wrapped_node") is True:
-                # Nodes that are wrapped or decorated by other nodes are not serialized here
-                # They are instead serialized by the wrapper node
-                continue
+        for node in self._workflow.get_nodes():
+            node_display = self.display_context.node_displays[node]
 
             try:
                 serialized_node = node_display.serialize(self.display_context)
@@ -295,13 +292,9 @@ class VellumWorkflowDisplay(
             else uuid4_from_hash(f"{self.workflow_id}|id|{entrypoint_node_id}")
         )
 
-        if has_wrapped_node(entrypoint):
-            wrapped_node = get_wrapped_node(entrypoint)
-            if wrapped_node._is_wrapped_node:
-                entrypoint = wrapped_node
-
-        target_node_id = node_displays[entrypoint].node_id
-        target_handle_id = node_displays[entrypoint].get_target_handle_id()
+        entrypoint_target = get_unadorned_node(entrypoint)
+        target_node_id = node_displays[entrypoint_target].node_id
+        target_handle_id = node_displays[entrypoint_target].get_target_handle_id()
 
         edge_display = self._generate_edge_display_from_source(
             entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides