vellum-ai 0.14.21__py3-none-any.whl → 0.14.22__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.14.21",
+            "X-Fern-SDK-Version": "0.14.22",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

vellum/utils/templating/constants.py

@@ -10,7 +10,7 @@ import pydash
 import pytz
 import yaml
 
-from vellum.utils.templating.custom_filters import is_valid_json_string, replace
+from vellum.utils.templating.custom_filters import is_valid_json_string, replace, safe_tojson
 
 DEFAULT_JINJA_GLOBALS: Dict[str, Any] = {
     "datetime": datetime,
@@ -24,9 +24,14 @@ DEFAULT_JINJA_GLOBALS: Dict[str, Any] = {
     "yaml": yaml,
 }
 
-FilterFunc = Union[Callable[[Union[str, bytes]], bool], Callable[[Any, Any, Any], str]]
+FilterFunc = Union[
+    Callable[[Union[str, bytes]], bool],  # is_valid_json_string
+    Callable[[Any, Any, Any], str],  # replace
+    Callable[[Any], str],  # safe_tojson
+]
 
 DEFAULT_JINJA_CUSTOM_FILTERS: Dict[str, FilterFunc] = {
     "is_valid_json_string": is_valid_json_string,
     "replace": replace,
+    "tojson": safe_tojson,
 }

vellum/utils/templating/custom_filters.py

@@ -1,6 +1,9 @@
 import json
 from typing import Any, Union
 
+from jinja2 import Undefined
+from jinja2.utils import htmlsafe_json_dumps
+
 from vellum.workflows.state.encoder import DefaultStateEncoder
 
 
@@ -31,3 +34,9 @@ def replace(s: Any, old: Any, new: Any) -> str:
     old_str = encode_to_str(old)
     new_str = encode_to_str(new)
     return s_str.replace(old_str, new_str)
+
+
+def safe_tojson(value: Any) -> str:
+    if isinstance(value, Undefined):
+        return ""
+    return htmlsafe_json_dumps(value, cls=DefaultStateEncoder)

vellum/workflows/README.md

@@ -77,7 +77,8 @@ code to UI and vice versa.
     ```bash
     python my_workflow.py
 
-Note: To use most out-of-box Nodes, and to push/pull to/from the Velłum UI, you'll need a Vellum account and API key. [Talk to us](https://www.vellum.ai/landing-pages/request-demo) or visit our [pricing page](https://www.vellum.ai/pricing) for more information.
+Note: To use most out-of-box Nodes, and to push/pull to/from the Vellum UI, you'll need a Vellum account and API key.
+You can [sign up for free here](https://app.vellum.ai/signup?f=wsdk&utm_source=github&utm_medium=repo_quickstart&utm_campaign=sdk).
 
 
 ## Documentation
@@ -86,5 +87,5 @@ Complete documentation for the Vellum Workflows SDK can be found at https://docs
 
 ## Stability
 
-This SDK is currently in <Availability type="beta" /> and is subject to change. If you'd like to pariticpate in
+This SDK is currently in <Availability type="beta" /> and is subject to change. If you'd like to participate in
 our beta program, please [contact us](https://docs.vellum.ai/home/getting-started/support).

vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py

@@ -299,3 +299,21 @@ def test_templating_node__empty_string_to_list():
 
     # THEN the output should be an empty list, not raise an exception
     assert outputs.result == []
+
+
+def test_api_error_templating_node():
+    class UndefinedTemplatingNode(TemplatingNode[BaseState, str]):
+        template = """{{ foo | tojson }}"""
+        inputs = {
+            "bar": "bar",
+            # foo is not define
+        }
+
+    # GIVEN a templating node with an undefined value
+    node = UndefinedTemplatingNode()
+
+    # WHEN the node is run
+    outputs = node.run()
+
+    # THEN the output should be empty string
+    assert outputs.result == ""

vellum/workflows/nodes/displayable/api_node/tests/test_api_node.py

@@ -20,7 +20,7 @@ def test_run_workflow__secrets(vellum_client):
         method = APIRequestMethod.POST
         authorization_type = AuthorizationType.BEARER_TOKEN
         url = "https://api.vellum.ai"
-        body = {
+        json = {
             "key": "value",
         }
         headers = {
@@ -32,6 +32,7 @@ def test_run_workflow__secrets(vellum_client):
     terminal = node.run()
 
     assert vellum_client.execute_api.call_count == 1
+    assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
     bearer_token = vellum_client.execute_api.call_args.kwargs["bearer_token"]
     assert bearer_token == ClientVellumSecret(name="secret")
     assert terminal.headers == {"X-Response-Header": "bar"}
@@ -45,7 +46,7 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
         method = APIRequestMethod.GET
         authorization_type = AuthorizationType.BEARER_TOKEN
         url = "https://api.vellum.ai"
-        body = {
+        json = {
             "key": "value",
         }
         headers = {
@@ -64,6 +65,7 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
 
     # AND the API call should have been made
     assert vellum_client.execute_api.call_count == 1
+    assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
 
 
 def test_api_node_defaults_to_get_method(vellum_client):

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

@@ -57,7 +57,7 @@ class BaseAPINode(BaseNode, Generic[StateType]):
             if isinstance(headers[header], VellumSecret):
                 vellum_instance = True
         if vellum_instance or bearer_token:
-            return self._vellum_execute_api(bearer_token, data, headers, method, url)
+            return self._vellum_execute_api(bearer_token, json, headers, method, url)
         else:
             return self._local_execute_api(data, headers, json, method, url)
 

{vellum_ai-0.14.21.dist-info → vellum_ai-0.14.22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.21
+Version: 0.14.22
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -90,8 +90,7 @@ Description-Content-Type: text/markdown
 
 ## Get Started
 
-Most functionality within the SDK requires a Vellum account and API key. To sign up, [talk to us](https://www.vellum.ai/landing-pages/request-demo)
-or visit our [pricing page](https://www.vellum.ai/pricing).
+Most functionality within the Vellum SDK requires a Vellum account and API key. You can sign up for free [here](https://app.vellum.ai/signup?f=wsdk&utm_source=github&utm_medium=repo_readme&utm_campaign=sdk) or visit our [pricing page](https://www.vellum.ai/pricing) for paid options.
 
 Even without a Vellum account, you can use the Workflows SDK to define the control flow of your AI systems. [Learn
 more below](#workflows-sdk).

{vellum_ai-0.14.21.dist-info → vellum_ai-0.14.22.dist-info}/RECORD

@@ -33,7 +33,7 @@ vellum_ee/workflows/display/nodes/types.py,sha256=St1BB6no528OyELGiyRabWao0GGw6m
 vellum_ee/workflows/display/nodes/utils.py,sha256=sloya5TpXsnot1HURc9L51INwflRqUzHxRVnCS9Cd-4,973
 vellum_ee/workflows/display/nodes/vellum/__init__.py,sha256=nUIgH2s0-7IbQRNrBhLPyRNe8YIrx3Yo9HeeW-aXXFk,1668
 vellum_ee/workflows/display/nodes/vellum/api_node.py,sha256=lumtypr0JEZfiA32hcs_olTmnT0wIUzPBy0pnZVfyU4,8532
-vellum_ee/workflows/display/nodes/vellum/base_adornment_node.py,sha256=oJAQrAm5iFQq0_fX94sMbS3RQEK1M1VsoUck4vsPs9A,5820
+vellum_ee/workflows/display/nodes/vellum/base_adornment_node.py,sha256=Bzqplrnx-bDIRD1JgU7036m8pSWSO45zEReNv8RJTu4,6379
 vellum_ee/workflows/display/nodes/vellum/code_execution_node.py,sha256=IYx0nll0t-tsPcjfgpfHMZE4FJgMFIsOiaQanGLYF0Q,4410
 vellum_ee/workflows/display/nodes/vellum/conditional_node.py,sha256=ybLIa4uclqVIy3VAQvI1ivg2tnK5Ug_1R5a69DFqL7E,11104
 vellum_ee/workflows/display/nodes/vellum/error_node.py,sha256=I1Jkp2htRINJATtv1e-zs9BrReFX842djpiVgBPHDYg,2186
@@ -93,10 +93,10 @@ vellum_ee/workflows/display/utils/expressions.py,sha256=9FpOslDI-RCR5m4TgAu9KCHh
 vellum_ee/workflows/display/utils/vellum.py,sha256=EVPQUSsZ3OIeLTEbV6LHPor37t9fnj9kJxDqP4PmTLQ,8234
 vellum_ee/workflows/display/vellum.py,sha256=bevbLCd2KqJBKqJ3lQayeRfjY7x1Djf57F9iJ-6KBJw,5162
 vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
-vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=2PyHM8bpgRqNtImCtN9fi5VqYvkOi5saVM9OXlwGmGQ,20400
+vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=cvpDj5gDzoO1Sdt8BCo1acvZsNb30bpX-xGftl3u0xE,19940
 vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
-vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=Ny9VWjCxc1_w8Z5xiQEJmNmEwtjlDlxgFCrmCWSmFmA,8511
-vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=wLXrYY0ccR37IG2MWwV49ELmwquVwDlDe3W4KXwOitM,16691
+vellum_ee/workflows/display/workflows/tests/test_workflow_display.py,sha256=TwxfmIpCL1xrJOjA54d52q5Ko0CFUT2bW60DVD0wrlY,10095
+vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=AS-vMrM93KEjb02-ye0Il29l3bX74o0Q8P2Nvf9NIZ0,16683
 vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum_ee/workflows/server/virtual_file_loader.py,sha256=ET-Q83W5Cgqzqz3qtFNwtS2nJEIcm3VtvR5kffsT3VY,2262
 vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -126,7 +126,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
 vellum/client/__init__.py,sha256=tKtdM1_GqmGq1gpi9ydWD_T-MM7fPn8QdHh8ww19cNI,117564
 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=r9FX6fq7T8lmJtpc4VW9IORBb0_3s14dOpTkKYA_l68,1869
+vellum/client/core/client_wrapper.py,sha256=DpiUwVloITIZTAtfFwbXdgMJxWZ-qVWlr-Qx6Kc1pZc,1869
 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
@@ -1292,8 +1292,8 @@ vellum/types/workspace_read.py,sha256=9CvgvK8Li8vL6qC5KX7f3-nEHslJ4lw2w07bvXcrjA
 vellum/types/workspace_secret_read.py,sha256=Z6QNXHxVHRdrLXSI31KxngePRwJTVoJYMXVbtPQwrxs,159
 vellum/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/utils/templating/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/utils/templating/constants.py,sha256=8OHMO6WFAEimbIaiHc5gy6s91D7_KvW-vTlEMWwvl_M,711
-vellum/utils/templating/custom_filters.py,sha256=XVHriwazejRZmxB_eg4xHgCxl7AiQQ2sx-hRLMmylfU,885
+vellum/utils/templating/constants.py,sha256=mvAcJloHe1D9-LzM_jpzVZEJVYy326OCUMtSqD_vmo0,838
+vellum/utils/templating/custom_filters.py,sha256=qdbDSBPVIz-jKbetD4OnZ4In6O-SdeV0oZsixYoiWOY,1116
 vellum/utils/templating/exceptions.py,sha256=cDp140PP4OnInW4qAvg3KqiSiF70C71UyEAKRBR1Abo,46
 vellum/utils/templating/render.py,sha256=P2t9qU4w_WdXAVLM5Nj3bc1-XlIKOkwK-czQ80pHBag,2172
 vellum/utils/templating/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1301,7 +1301,7 @@ vellum/utils/templating/tests/test_custom_filters.py,sha256=mkJwc7t1gE13SKgPxhF-
 vellum/utils/typing.py,sha256=wx_daFqD69cYkuJTVnvNrpjhqC3uuhbnyJ9_bIwC9OU,327
 vellum/utils/uuid.py,sha256=Ch6wWRgwICxLxJCTl5iE3EdRlZj2zADR-zUMUtjcMWM,214
 vellum/version.py,sha256=jq-1PlAYxN9AXuaZqbYk9ak27SgE2lw9Ia5gx1b1gVI,76
-vellum/workflows/README.md,sha256=MLNm-ihc0ao6I8gwwOhXQQBf0jOf-EsA9C519ALYI1o,3610
+vellum/workflows/README.md,sha256=hZdTKBIcsTKPofK68oPkBhyt0nnRh0csqC12k4FMHHA,3597
 vellum/workflows/__init__.py,sha256=CssPsbNvN6rDhoLuqpEv7MMKGa51vE6dvAh6U31Pcio,71
 vellum/workflows/constants.py,sha256=2yg4_uo5gpqViy3ZLSwfC8qTybleYCtOnhA4Rj6bacM,1310
 vellum/workflows/context.py,sha256=DwSf8lO9NHABiqOoD3exgrjUoRuNsKtutaL5TgRbD-A,1441
@@ -1390,7 +1390,7 @@ vellum/workflows/nodes/core/retry_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TI
 vellum/workflows/nodes/core/retry_node/tests/test_node.py,sha256=fNgDufkIsrTC-6ftvogqSpWhqqBj9iNESdfK19B1Yx0,5159
 vellum/workflows/nodes/core/templating_node/__init__.py,sha256=GmyuYo81_A1_Bz6id69ozVFS6FKiuDsZTiA3I6MaL2U,70
 vellum/workflows/nodes/core/templating_node/node.py,sha256=iqBmr2i-f-BqhisNQJiDfewjol0ur7-XpupLStyMJsg,3731
-vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=MHofz-BwAgt7EXkab8VIyacYznDEIJ7Er7MJUaxNQQo,9614
+vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py,sha256=nXkgGDBg4wP36NwykdMEVWwx_xjv8oGT2rYkwuCB_VU,10075
 vellum/workflows/nodes/core/try_node/__init__.py,sha256=JVD4DrldTIqFQQFrubs9KtWCCc0YCAc7Fzol5ZWIWeM,56
 vellum/workflows/nodes/core/try_node/node.py,sha256=RbxL0NRXS0IxRP0MJAnLABolF6dkwVniiqsagzy-lwk,4445
 vellum/workflows/nodes/core/try_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -1399,10 +1399,10 @@ vellum/workflows/nodes/displayable/__init__.py,sha256=6F_4DlSwvHuilWnIalp8iDjjDX
 vellum/workflows/nodes/displayable/api_node/__init__.py,sha256=MoxdQSnidIj1Nf_d-hTxlOxcZXaZnsWFDbE-PkTK24o,56
 vellum/workflows/nodes/displayable/api_node/node.py,sha256=cp0nAukcOpM6TcNhbz12h08TMJxp_LM-MLDl1dAzYsk,2534
 vellum/workflows/nodes/displayable/api_node/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vellum/workflows/nodes/displayable/api_node/tests/test_api_node.py,sha256=fiskhfcI4c6CxFlbSWb1JKsuq-98zAeUWOExc848ALw,3130
+vellum/workflows/nodes/displayable/api_node/tests/test_api_node.py,sha256=J21H1dQT0BJ0oAalaA-9mgKv-NRcCJaTImhnKXp-cX4,3294
 vellum/workflows/nodes/displayable/bases/__init__.py,sha256=0mWIx3qUrzllV7jqt7wN03vWGMuI1WrrLZeMLT2Cl2c,304
 vellum/workflows/nodes/displayable/bases/api_node/__init__.py,sha256=1jwx4WC358CLA1jgzl_UD-rZmdMm2v9Mps39ndwCD7U,64
-vellum/workflows/nodes/displayable/bases/api_node/node.py,sha256=I4zdxSXJ7-pjVcNNrSAEN_hAnhQGUBM6u5RczCActkw,4343
+vellum/workflows/nodes/displayable/bases/api_node/node.py,sha256=70pLGU0UzWvSbKwNkx3YlUYrDSkl7MmhVHoI8bzN79c,4343
 vellum/workflows/nodes/displayable/bases/base_prompt_node/__init__.py,sha256=Org3xTvgp1pA0uUXFfnJr29D3HzCey2lEdYF4zbIUgo,70
 vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py,sha256=HGNoGLJ9lbqflGdYFDIiuHFyi0iJ-agJu4kkJ7D3dGs,3212
 vellum/workflows/nodes/displayable/bases/inline_prompt_node/__init__.py,sha256=Hl35IAoepRpE-j4cALaXVJIYTYOF3qszyVbxTj4kS1s,82
@@ -1522,8 +1522,8 @@ vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnad
 vellum/workflows/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 vellum/workflows/workflows/tests/test_base_workflow.py,sha256=NRteiICyJvDM5zrtUfq2fZoXcGQVaWC9xmNlLLVW0cU,7979
 vellum/workflows/workflows/tests/test_context.py,sha256=VJBUcyWVtMa_lE5KxdhgMu0WYNYnUQUDvTF7qm89hJ0,2333
-vellum_ai-0.14.21.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
-vellum_ai-0.14.21.dist-info/METADATA,sha256=UqD3qeDGaeZTzlZPvYTpGbLEGaN7RYFGaTsxUDoboCo,5408
-vellum_ai-0.14.21.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-vellum_ai-0.14.21.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
-vellum_ai-0.14.21.dist-info/RECORD,,
+vellum_ai-0.14.22.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
+vellum_ai-0.14.22.dist-info/METADATA,sha256=n1LtcR20giCKtaq1b6qilNyA4MKvG4EON_hG8EC8A9U,5484
+vellum_ai-0.14.22.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+vellum_ai-0.14.22.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
+vellum_ai-0.14.22.dist-info/RECORD,,

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

@@ -3,6 +3,7 @@ import types
 from uuid import UUID
 from typing import Any, Callable, Dict, Generic, Optional, Type, TypeVar, cast
 
+from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.nodes.utils import get_wrapped_node
 from vellum.workflows.types.core import JsonArray, JsonObject
@@ -17,7 +18,62 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
 _BaseAdornmentNodeType = TypeVar("_BaseAdornmentNodeType", bound=BaseAdornmentNode)
 
 
+def _recursively_replace_wrapped_node(node_class: Type[BaseNode], wrapped_node_display_class: Type["BaseNodeDisplay"]):
+    if not issubclass(node_class, BaseAdornmentNode):
+        return
+
+    # We must edit the node display class to use __wrapped_node__ everywhere it
+    # references the adorned node class, which is three places:
+    wrapped_node_class = node_class.__wrapped_node__
+    if not wrapped_node_class:
+        raise ValueError(f"Node {node_class.__name__} must be used as an adornment with the `wrap` method.")
+
+    # 1. The node display class' parameterized type
+    original_base_node_display = get_original_base(wrapped_node_display_class)
+    original_base_node_display.__args__ = (wrapped_node_class,)
+    wrapped_node_display_class._node_display_registry[wrapped_node_class] = wrapped_node_display_class
+    wrapped_node_display_class.__annotate_node__()
+
+    # 2. The node display class' output displays
+    for old_output in node_class.Outputs:
+        new_output = getattr(wrapped_node_class.Outputs, old_output.name, None)
+        if new_output is None:
+            # If the adornment is adding a new output, such as TryNode adding an "error" output,
+            # we skip it, since it should not be included in the adorned node's output displays
+            wrapped_node_display_class.output_display.pop(old_output, None)
+            continue
+
+        if old_output not in wrapped_node_display_class.output_display:
+            # If the adorned node doesn't have an output display defined for this output, we define one
+            wrapped_node_display_class.output_display[new_output] = NodeOutputDisplay(
+                id=wrapped_node_class.__output_ids__[old_output.name],
+                name=old_output.name,
+            )
+        else:
+            wrapped_node_display_class.output_display[new_output] = wrapped_node_display_class.output_display.pop(
+                old_output
+            )
+
+    # 3. The node display class' port displays
+    old_ports = list(wrapped_node_display_class.port_displays.keys())
+    for old_port in old_ports:
+        new_port = getattr(wrapped_node_class.Ports, old_port.name)
+        wrapped_node_display_class.port_displays[new_port] = wrapped_node_display_class.port_displays.pop(old_port)
+
+    # Now we keep drilling down recursively
+    if (
+        issubclass(wrapped_node_display_class, BaseAdornmentNodeDisplay)
+        and wrapped_node_display_class.__wrapped_node_display__
+    ):
+        _recursively_replace_wrapped_node(
+            wrapped_node_class,
+            wrapped_node_display_class.__wrapped_node_display__,
+        )
+
+
 class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Generic[_BaseAdornmentNodeType]):
+    __wrapped_node_display__: Optional[Type[BaseNodeDisplay]] = None
+
     def serialize(
         self,
         display_context: "WorkflowDisplayContext",
@@ -52,10 +108,6 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
             if not issubclass(node_class, BaseAdornmentNode):
                 raise ValueError(f"Node {node_class.__name__} must be wrapped with a {BaseAdornmentNode.__name__}")
 
-            wrapped_node_class = node_class.__wrapped_node__
-            if not wrapped_node_class:
-                raise ValueError(f"Node {node_class.__name__} must be used as an adornment with the `wrap` method.")
-
             # `mypy` is wrong here, `cls` is indexable bc it's Generic
             BaseAdornmentDisplay = cls[node_class]  # type: ignore[index]
 
@@ -64,6 +116,7 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
                     ns[key] = kwarg
 
                 ns["node_id"] = node_id or uuid4_from_hash(node_class.__qualname__)
+                ns["__wrapped_node_display__"] = wrapped_node_display_class
 
             AdornmentDisplay = types.new_class(
                 re.sub(r"^Base", "", cls.__name__), bases=(BaseAdornmentDisplay,), exec_body=exec_body
@@ -71,43 +124,11 @@ class BaseAdornmentNodeDisplay(BaseNodeVellumDisplay[_BaseAdornmentNodeType], Ge
 
             setattr(wrapped_node_display_class, "__adorned_by__", AdornmentDisplay)
 
-            # We must edit the node display class to use __wrapped_node__ everywhere it
-            # references the adorned node class, which is three places:
-
-            # 1. The node display class' parameterized type
-            original_base_node_display = get_original_base(wrapped_node_display_class)
-            original_base_node_display.__args__ = (wrapped_node_class,)
-            wrapped_node_display_class._node_display_registry[wrapped_node_class] = wrapped_node_display_class
-            wrapped_node_display_class.__annotate_node__()
-
-            # 2. The node display class' output displays
-            for old_output in node_class.Outputs:
-                new_output = getattr(wrapped_node_class.Outputs, old_output.name, None)
-                if new_output is None:
-                    # If the adornment is adding a new output, such as TryNode adding an "error" output,
-                    # we skip it, since it should not be included in the adorned node's output displays
-                    wrapped_node_display_class.output_display.pop(old_output, None)
-                    continue
-
-                if old_output not in wrapped_node_display_class.output_display:
-                    # If the adorned node doesn't have an output display defined for this output, we define one
-                    wrapped_node_display_class.output_display[new_output] = NodeOutputDisplay(
-                        id=wrapped_node_class.__output_ids__[old_output.name],
-                        name=old_output.name,
-                    )
-                else:
-                    wrapped_node_display_class.output_display[new_output] = (
-                        wrapped_node_display_class.output_display.pop(old_output)
-                    )
-
-            # 3. The node display class' port displays
-            old_ports = list(wrapped_node_display_class.port_displays.keys())
-            for old_port in old_ports:
-                new_port = getattr(wrapped_node_class.Ports, old_port.name)
-                wrapped_node_display_class.port_displays[new_port] = wrapped_node_display_class.port_displays.pop(
-                    old_port
-                )
-
-            return wrapped_node_display_class
+            _recursively_replace_wrapped_node(
+                node_class,
+                wrapped_node_display_class,
+            )
+
+            return AdornmentDisplay
 
         return decorator

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -150,8 +150,6 @@ class BaseWorkflowDisplay(
             if node_output in node_output_displays:
                 continue
 
-            # 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_display.get_node_output_display(node_output)
 
     def _enrich_node_port_displays(
@@ -272,8 +270,8 @@ class BaseWorkflowDisplay(
                 continue
 
             edge_display_overrides = self.edge_displays.get((edge.from_port, edge.to_node))
-            edge_displays[(edge.from_port, edge.to_node)] = self._generate_edge_display(
-                edge, node_displays, overrides=edge_display_overrides
+            edge_displays[(edge.from_port, edge.to_node)] = edge_display_overrides or self._generate_edge_display(
+                edge, node_displays
             )
 
         for edge in self._workflow.get_unused_edges():
@@ -281,8 +279,8 @@ class BaseWorkflowDisplay(
                 continue
 
             edge_display_overrides = self.edge_displays.get((edge.from_port, edge.to_node))
-            edge_displays[(edge.from_port, edge.to_node)] = self._generate_edge_display(
-                edge, node_displays, overrides=edge_display_overrides
+            edge_displays[(edge.from_port, edge.to_node)] = edge_display_overrides or self._generate_edge_display(
+                edge, node_displays
             )
 
         workflow_output_displays: Dict[BaseDescriptor, WorkflowOutputDisplay] = {}
@@ -443,30 +441,20 @@ class BaseWorkflowDisplay(
 
         return additional_node_displays
 
-    def _generate_edge_display(
-        self,
-        edge: Edge,
-        node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
-        overrides: Optional[EdgeDisplay] = None,
-    ) -> EdgeDisplay:
+    def _generate_edge_display(self, edge: Edge, node_displays: Dict[Type[BaseNode], BaseNodeDisplay]) -> EdgeDisplay:
         source_node = get_unadorned_node(edge.from_port.node_class)
         target_node = get_unadorned_node(edge.to_node)
 
         source_node_id = node_displays[source_node].node_id
         target_node_id = node_displays[target_node].node_id
 
-        return self._generate_edge_display_from_source(source_node_id, target_node_id, overrides)
+        return self._generate_edge_display_from_source(source_node_id, target_node_id)
 
     def _generate_edge_display_from_source(
         self,
         source_node_id: UUID,
         target_node_id: UUID,
-        overrides: Optional[EdgeDisplay] = None,
     ) -> EdgeDisplay:
-        edge_id: UUID
-        if overrides:
-            edge_id = overrides.id
-        else:
-            edge_id = uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}")
-
-        return EdgeDisplay(id=edge_id)
+        return EdgeDisplay(
+            id=uuid4_from_hash(f"{self.workflow_id}|id|{source_node_id}|{target_node_id}"),
+        )

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -235,3 +235,43 @@ def test_get_event_display_context__templating_node_input_display():
     node_event_display = display_context.node_displays[MyNode.__id__]
 
     assert node_event_display.input_display.keys() == {"inputs.foo"}
+
+
+def test_get_event_display_context__node_display_for_mutiple_adornments():
+    # GIVEN a simple workflow with multiple adornments
+    @TryNode.wrap()
+    @RetryNode.wrap()
+    class MyNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            foo: str
+
+    class MyWorkflow(BaseWorkflow):
+        graph = MyNode
+
+    # AND a display class for the node
+    node_id = uuid4()
+    inner_node_id = uuid4()
+    innermost_node_id = uuid4()
+
+    @BaseTryNodeDisplay.wrap(node_id=node_id)
+    @BaseRetryNodeDisplay.wrap(node_id=inner_node_id)
+    class MyNodeDisplay(BaseNodeDisplay[MyNode]):
+        node_id = innermost_node_id
+
+    # WHEN we gather the event display context
+    display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
+
+    # THEN the subworkflow display should be included
+    assert node_id in display_context.node_displays
+    node_event_display = display_context.node_displays[node_id]
+    assert node_event_display.subworkflow_display
+
+    # AND the inner node should be included
+    assert inner_node_id in node_event_display.subworkflow_display.node_displays
+    inner_node_event_display = node_event_display.subworkflow_display.node_displays[inner_node_id]
+    assert inner_node_event_display.subworkflow_display
+
+    # AND the innermost node should be included
+    assert innermost_node_id in inner_node_event_display.subworkflow_display.node_displays
+    innermost_node_event_display = inner_node_event_display.subworkflow_display.node_displays[innermost_node_id]
+    assert not innermost_node_event_display.subworkflow_display

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -363,8 +363,8 @@ class VellumWorkflowDisplay(
         target_node_display = node_displays[entrypoint_target]
         target_node_id = target_node_display.node_id
 
-        edge_display = self._generate_edge_display_from_source(
-            entrypoint_node_id, target_node_id, overrides=edge_display_overrides
+        edge_display = edge_display_overrides or self._generate_edge_display_from_source(
+            entrypoint_node_id, target_node_id
         )
 
         return EntrypointVellumDisplay(id=entrypoint_id, edge_display=edge_display)