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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -18,7 +18,7 @@ class BaseClientWrapper:
18
18
  headers: typing.Dict[str, str] = {
19
19
  "X-Fern-Language": "Python",
20
20
  "X-Fern-SDK-Name": "vellum-ai",
21
- "X-Fern-SDK-Version": "0.13.8",
21
+ "X-Fern-SDK-Version": "0.13.9",
22
22
  }
23
23
  headers["X_API_KEY"] = self.api_key
24
24
  return headers
@@ -19,6 +19,7 @@ from mypy.plugin import AttributeContext, ClassDefContext, FunctionSigContext, M
19
19
  from mypy.types import AnyType, CallableType, FunctionLike, Instance, Type as MypyType, TypeAliasType, UnionType
20
20
 
21
21
  TypeResolver = Callable[[str, List[MypyType]], MypyType]
22
+ NODE_ATTRIBUTE_REGEX = r"^[a-z].*$"
22
23
 
23
24
  DESCRIPTOR_PATHS: list[tuple[str, str, str]] = [
24
25
  (
@@ -49,7 +50,7 @@ DESCRIPTOR_PATHS: list[tuple[str, str, str]] = [
49
50
  (
50
51
  "vellum.workflows.nodes.bases.base.BaseNode",
51
52
  "vellum.workflows.references.node.NodeReference",
52
- r"^[a-z].*$",
53
+ NODE_ATTRIBUTE_REGEX,
53
54
  ),
54
55
  ]
55
56
 
@@ -239,10 +240,13 @@ class VellumMypyPlugin(Plugin):
239
240
  the original defined type, and the descriptor version of the type.
240
241
  """
241
242
 
242
- for sym in ctx.cls.info.names.values():
243
+ for key, sym in ctx.cls.info.names.items():
243
244
  if not isinstance(sym.node, Var):
244
245
  continue
245
246
 
247
+ if not re.match(NODE_ATTRIBUTE_REGEX, key):
248
+ continue
249
+
246
250
  type_ = sym.node.type
247
251
  if not type_:
248
252
  continue
@@ -222,7 +222,6 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
222
222
  state: StateType
223
223
  _context: WorkflowContext
224
224
  _inputs: MappingProxyType[NodeReference, Any]
225
- _is_wrapped_node: bool = False
226
225
 
227
226
  class ExternalInputs(BaseInputs):
228
227
  __descriptor_class__ = ExternalInputReference
@@ -4,6 +4,7 @@ from types import ModuleType
4
4
  from typing import Any, Callable, Optional, Type, TypeVar
5
5
 
6
6
  from vellum.workflows.nodes import BaseNode
7
+ from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
7
8
  from vellum.workflows.ports.port import Port
8
9
  from vellum.workflows.types.generics import NodeType
9
10
 
@@ -12,7 +13,7 @@ ADORNMENT_MODULE_NAME = "<adornment>"
12
13
 
13
14
  @cache
14
15
  def get_unadorned_node(node: Type[BaseNode]) -> Type[BaseNode]:
15
- wrapped_node = getattr(node, "__wrapped_node__", None)
16
+ wrapped_node = get_wrapped_node(node)
16
17
  if wrapped_node is not None:
17
18
  return get_unadorned_node(wrapped_node)
18
19
 
@@ -28,22 +29,11 @@ def get_unadorned_port(port: Port) -> Port:
28
29
  return getattr(unadorned_node.Ports, port.name)
29
30
 
30
31
 
31
- @cache
32
- def get_wrapped_node(node: Type[NodeType]) -> Type[BaseNode]:
33
- wrapped_node = getattr(node, "__wrapped_node__", None)
34
- if wrapped_node is None:
35
- raise AttributeError("Wrapped node not found")
36
-
37
- return wrapped_node
38
-
32
+ def get_wrapped_node(node: Type[NodeType]) -> Optional[Type[BaseNode]]:
33
+ if not issubclass(node, BaseAdornmentNode):
34
+ return None
39
35
 
40
- def has_wrapped_node(node: Type[NodeType]) -> bool:
41
- try:
42
- get_wrapped_node(node)
43
- except AttributeError:
44
- return False
45
-
46
- return True
36
+ return node.__wrapped_node__
47
37
 
48
38
 
49
39
  AdornableNode = TypeVar("AdornableNode", bound=BaseNode)
@@ -57,8 +47,6 @@ def create_adornment(
57
47
  # https://app.shortcut.com/vellum/story/4116
58
48
  from vellum.workflows import BaseWorkflow
59
49
 
60
- inner_cls._is_wrapped_node = True
61
-
62
50
  class Subworkflow(BaseWorkflow):
63
51
  graph = inner_cls
64
52
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: vellum-ai
3
- Version: 0.13.8
3
+ Version: 0.13.9
4
4
  Summary:
5
5
  License: MIT
6
6
  Requires-Python: >=3.9,<4.0
@@ -21,7 +21,7 @@ vellum_ee/workflows/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
21
21
  vellum_ee/workflows/display/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
22
22
  vellum_ee/workflows/display/base.py,sha256=3ZFUYRNKL24fBqXhKpa_Dq2W1a-a86J20dmJYA3H2eY,1755
23
23
  vellum_ee/workflows/display/nodes/__init__.py,sha256=5XOcZJXYUgaLS55QgRJzyQ_W1tpeprjnYAeYVezqoGw,160
24
- vellum_ee/workflows/display/nodes/base_node_display.py,sha256=pQ-QlaDdxME3roRrQj5eYRZwZXuE9im8gjo2QFwn954,6942
24
+ vellum_ee/workflows/display/nodes/base_node_display.py,sha256=Ge3f_MNCxawRNR2EhBvygXTIWxKOpS40Dkgrp1S1WPQ,6991
25
25
  vellum_ee/workflows/display/nodes/base_node_vellum_display.py,sha256=BxA-YVUJvB36NM-Q5DNCwckeqymwLKMp9DVCaTyrPbs,2253
26
26
  vellum_ee/workflows/display/nodes/get_node_display_class.py,sha256=vyKeJAevAXvEAEtWeTEdBZXD3eJQYW_DEXLKVZ5KmYc,1325
27
27
  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
46
46
  vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py,sha256=lfevlHpGEX14dEDym6qmnkw3nvzQPTB1_D2ch12B_Rk,2701
47
47
  vellum_ee/workflows/display/nodes/vellum/templating_node.py,sha256=JVIMPR3WpveOCWZubHKZkE04mavnTdb_9QY_r3XliRg,3424
48
48
  vellum_ee/workflows/display/nodes/vellum/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
49
- vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=Nqd6mxn0sgL4sp62wDYnQVc7COxrt5wXXIveRFoYQ8c,3913
50
- vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=hB8dcGMGkuC95kk9hmZUgHsCLwEA37fHTFXj0JzbRjM,4692
51
- vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=nIZ1DYTYPRaYsVKcel9a-Fm8EniJL0N7f5PowxVGTVU,8318
49
+ vellum_ee/workflows/display/nodes/vellum/tests/test_try_node.py,sha256=mtzB8LJlFCHVFM4H5AanLp29gQfaVmnN4A4iaRGJHoI,2427
50
+ vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py,sha256=3uT7Gbc0f_mQ3u8uZuCWd0mJ4GtWbz2gbUMySYaVlNE,3774
51
+ vellum_ee/workflows/display/nodes/vellum/try_node.py,sha256=Ywa8NMNvE4TlEZ_gQA6jhrqykH8-wWgKFa-uEokWhF8,5550
52
+ vellum_ee/workflows/display/nodes/vellum/utils.py,sha256=TFChHlSN8B0MOdgFTDmQPKkqwGGYxeF5lNmFp-JXzqo,7882
52
53
  vellum_ee/workflows/display/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
53
54
  vellum_ee/workflows/display/tests/test_vellum_workflow_display.py,sha256=4jXe7Wyspw6CxghVqKAOu-_QunKFvY4U6--zOiuXvV0,3069
54
55
  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
60
61
  vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py,sha256=wZekdhIZBwbBoFNAdC9bBLwUxVKk3EnFdNmMwYIdKGA,37308
61
62
  vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py,sha256=os6wCZjcOyig8EJu-pTAFWXa0odMxTaR0mrbL0SS_4Y,4480
62
63
  vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py,sha256=bXZWxOKAVjZlbP3iLHPHGA4aPs0EguKlQqmYPNGv3cU,16391
63
- vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=UPLEQPwsLzOqZdkXrB5Jo1FdCx0iQNf7ekfcq1byoFw,29392
64
+ vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py,sha256=uvABhpstyxNNOz50-XJMAr3SKp8gluPp1mUtJjvRL0c,29410
64
65
  vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py,sha256=Xn8t-SpvQocKdxNofDeDETlFnkCbBxdLGiUG8m6fM6U,48399
65
66
  vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py,sha256=QVlkczxzaKuOhwbRvVP70otPDNnJcSGDfN79j92lFyk,5534
66
67
  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
80
81
  vellum_ee/workflows/display/utils/vellum.py,sha256=N5VSDYdtHUxvTeSj4zA8aMiAAKn4bAVgpKUDV_obNQ8,3632
81
82
  vellum_ee/workflows/display/vellum.py,sha256=8xXRI8b8Tt661H-iZreTQTvLNEKUr4lf-XaKhE7_yUY,8147
82
83
  vellum_ee/workflows/display/workflows/__init__.py,sha256=kapXsC67VJcgSuiBMa86FdePG5A9kMB5Pi4Uy1O2ob4,207
83
- vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=cLmW-oMMaP2BSwXCTVJ_KxhKmyfTRRJ_TGhkW3AQYXA,14324
84
+ vellum_ee/workflows/display/workflows/base_workflow_display.py,sha256=VNYc9n_TvsYxtJ__BaTI9ZGEzTg8JClrZ-nVnFlgiBU,14051
84
85
  vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py,sha256=kp0u8LN_2IwshLrhMImhpZx1hRyAcD5gXY-kDuuaGMQ,1269
85
- vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=V0edhtohqAWbaHvHkj9Sth4ieaIVejsrsRIr7aCCoVc,16871
86
+ vellum_ee/workflows/display/workflows/vellum_workflow_display.py,sha256=ZGPhfb-4_7bLEYaKnESkZuwWrzh9lmT2wMq2jPQ-ciQ,16518
86
87
  vellum_ee/workflows/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
87
88
  vellum_ee/workflows/server/virtual_file_loader.py,sha256=X_DdNK7MfyOjKWekk6YQpOSCT6klKcdjT6nVJcBH1sM,1481
88
89
  vellum_ee/workflows/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -111,7 +112,7 @@ vellum/client/README.md,sha256=JkCJjmMZl4jrPj46pkmL9dpK4gSzQQmP5I7z4aME4LY,4749
111
112
  vellum/client/__init__.py,sha256=8nZt88C9SVwWanjLbIQMU3rzb32h5UZfFMBx3VPHB50,111887
112
113
  vellum/client/core/__init__.py,sha256=SQ85PF84B9MuKnBwHNHWemSGuy-g_515gFYNFhvEE0I,1438
113
114
  vellum/client/core/api_error.py,sha256=RE8LELok2QCjABadECTvtDp7qejA1VmINCh6TbqPwSE,426
114
- vellum/client/core/client_wrapper.py,sha256=NnIgNorh26bh1LnC042Vcmf_kylY4hdObVYLpA5mq6k,1868
115
+ vellum/client/core/client_wrapper.py,sha256=dxLror26l68ej4_dKgIfKbRZE9rDH7pYhicRFBkhKss,1868
115
116
  vellum/client/core/datetime_utils.py,sha256=nBys2IsYrhPdszxGKCNRPSOCwa-5DWOHG95FB8G9PKo,1047
116
117
  vellum/client/core/file.py,sha256=X9IbmkZmB2bB_DpmZAO3crWdXagOakAyn6UCOCImCPg,2322
117
118
  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
691
692
  vellum/plugins/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
692
693
  vellum/plugins/pydantic.py,sha256=dNtZWHo-IdseG52C2RoTanxyTJg0AhPZrH-9lbNqwYg,2604
693
694
  vellum/plugins/utils.py,sha256=U9ZY9KdE3RRvbcG01hXxu9CvfJD6Fo7nJDgcHjQn0FI,606
694
- vellum/plugins/vellum_mypy.py,sha256=7JmyuTX723-XWnWoE6afiUNOkHyAufCUZwxck9BP2aI,27784
695
+ vellum/plugins/vellum_mypy.py,sha256=QTuMSq6PiZW1dyTUZ5Bf1d4XkgFj0TKAgZLP8f4UgL4,27914
695
696
  vellum/prompts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
696
697
  vellum/prompts/blocks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
697
698
  vellum/prompts/blocks/compilation.py,sha256=ISuvDHaeVCPb1L7l4umORCECkDn0-rvE49hopz6c2gM,9222
@@ -1305,7 +1306,7 @@ vellum/workflows/inputs/base.py,sha256=1kMgr0WqCYdWUqgFvgSoAMw2067FAlgwhGXLgbIOr
1305
1306
  vellum/workflows/logging.py,sha256=_a217XogktV4Ncz6xKFz7WfYmZAzkfVRVuC0rWob8ls,437
1306
1307
  vellum/workflows/nodes/__init__.py,sha256=aVdQVv7Y3Ro3JlqXGpxwaU2zrI06plDHD2aumH5WUIs,1157
1307
1308
  vellum/workflows/nodes/bases/__init__.py,sha256=cniHuz_RXdJ4TQgD8CBzoiKDiPxg62ErdVpCbWICX64,58
1308
- vellum/workflows/nodes/bases/base.py,sha256=CcWQBMR3cx5vftqQp_oJ_GncULJIbOLMyNP4KZNgU10,14487
1309
+ vellum/workflows/nodes/bases/base.py,sha256=P8j70SnO5-xABMbAPCHNKLook5Ip32FpzqfYEUZG8E8,14452
1309
1310
  vellum/workflows/nodes/bases/base_adornment_node.py,sha256=eFTgsPCYb3eyGS0-kw7C6crFnwFx437R5wh9-8bWYts,2905
1310
1311
  vellum/workflows/nodes/bases/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
1311
1312
  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
1385
1386
  vellum/workflows/nodes/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
1386
1387
  vellum/workflows/nodes/experimental/openai_chat_completion_node/__init__.py,sha256=lsyD9laR9p7kx5-BXGH2gUTM242UhKy8SMV0SR6S2iE,90
1387
1388
  vellum/workflows/nodes/experimental/openai_chat_completion_node/node.py,sha256=1EGeiaT-Zoo6pttQFKKBcdf3dmhAbjKGaErYD5FFwlc,10185
1388
- vellum/workflows/nodes/utils.py,sha256=chSsmKe_BsvMIJpzSxO5TWYlr3sAuxiwkfB5azkuN5Q,2715
1389
+ vellum/workflows/nodes/utils.py,sha256=tjBsootwm7vUq9qU4ttDL3CWH0C9Sd6QlX7IhfCCy34,2512
1389
1390
  vellum/workflows/outputs/__init__.py,sha256=AyZ4pRh_ACQIGvkf0byJO46EDnSix1ZCAXfvh-ms1QE,94
1390
1391
  vellum/workflows/outputs/base.py,sha256=a7W6rNSDSawwGAXYjNTF2iHb9lnZu7WFSOagZIyy__k,7976
1391
1392
  vellum/workflows/ports/__init__.py,sha256=bZuMt-R7z5bKwpu4uPW7LlJeePOQWmCcDSXe5frUY5g,101
@@ -1440,8 +1441,8 @@ vellum/workflows/vellum_client.py,sha256=ODrq_TSl-drX2aezXegf7pizpWDVJuTXH-j6528
1440
1441
  vellum/workflows/workflows/__init__.py,sha256=KY45TqvavCCvXIkyCFMEc0dc6jTMOUci93U2DUrlZYc,66
1441
1442
  vellum/workflows/workflows/base.py,sha256=k0kUWWko4fHyCqLSU_1cBK_pXZpl9MXekWiG-bdOAo0,18353
1442
1443
  vellum/workflows/workflows/event_filters.py,sha256=GSxIgwrX26a1Smfd-6yss2abGCnadGsrSZGa7t7LpJA,2008
1443
- vellum_ai-0.13.8.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1444
- vellum_ai-0.13.8.dist-info/METADATA,sha256=DslQiu2NfpU6727w1BqOFezu5H0uKbP0EiBP1ALlB1Y,5334
1445
- vellum_ai-0.13.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1446
- vellum_ai-0.13.8.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1447
- vellum_ai-0.13.8.dist-info/RECORD,,
1444
+ vellum_ai-0.13.9.dist-info/LICENSE,sha256=hOypcdt481qGNISA784bnAGWAE6tyIf9gc2E78mYC3E,1574
1445
+ vellum_ai-0.13.9.dist-info/METADATA,sha256=pIFtYGqbPEQ93I_doRy1gCDYIObzB2saS0Wd_2UNfwY,5334
1446
+ vellum_ai-0.13.9.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
1447
+ vellum_ai-0.13.9.dist-info/entry_points.txt,sha256=HCH4yc_V3J_nDv3qJzZ_nYS8llCHZViCDP1ejgCc5Ak,42
1448
+ vellum_ai-0.13.9.dist-info/RECORD,,
@@ -76,12 +76,12 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
76
76
  )
77
77
  return node_definition
78
78
 
79
- def get_node_output_display(self, output: OutputReference) -> NodeOutputDisplay:
79
+ def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
80
80
  explicit_display = self.output_display.get(output)
81
81
  if explicit_display:
82
- return explicit_display
82
+ return self._node, explicit_display
83
83
 
84
- return NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name)
84
+ return (self._node, NodeOutputDisplay(id=uuid4_from_hash(f"{self.node_id}|{output.name}"), name=output.name))
85
85
 
86
86
  def get_node_port_display(self, port: Port) -> PortDisplay:
87
87
  overrides = self.port_displays.get(port)
@@ -0,0 +1,63 @@
1
+ from typing import Any, Dict, cast
2
+
3
+ from vellum.workflows import BaseWorkflow
4
+ from vellum.workflows.nodes.bases.base import BaseNode
5
+ from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
6
+ from vellum.workflows.nodes.core.try_node.node import TryNode
7
+ from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
8
+ from vellum_ee.workflows.display.workflows.vellum_workflow_display import VellumWorkflowDisplay
9
+
10
+
11
+ def test_try_node_display__serialize_with_error_output() -> None:
12
+ # GIVEN a Base Node wrapped with a TryNode
13
+ @TryNode.wrap()
14
+ class MyNode(BaseNode):
15
+ class Outputs(BaseNode.Outputs):
16
+ hello = "world"
17
+
18
+ # AND a displayable node referencing
19
+ class OtherNode(TemplatingNode):
20
+ template = "{{ hello }} {{ error }}"
21
+ inputs = {
22
+ "hello": MyNode.Outputs.hello,
23
+ "error": MyNode.Outputs.error,
24
+ }
25
+
26
+ # AND a workflow referencing the two nodes
27
+ class MyWorkflow(BaseWorkflow):
28
+ graph = MyNode >> OtherNode
29
+
30
+ # WHEN we serialize the workflow
31
+ workflow_display = get_workflow_display(base_display_class=VellumWorkflowDisplay, workflow_class=MyWorkflow)
32
+ serialized_workflow = cast(Dict[str, Any], workflow_display.serialize())
33
+
34
+ # THEN the correct inputs should be serialized on the node
35
+ serialized_node = next(
36
+ node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(OtherNode.__id__)
37
+ )
38
+ hello_input_value = next(input["value"] for input in serialized_node["inputs"] if input["key"] == "hello")
39
+ error_input_value = next(input["value"] for input in serialized_node["inputs"] if input["key"] == "error")
40
+ assert hello_input_value == {
41
+ "combinator": "OR",
42
+ "rules": [
43
+ {
44
+ "data": {
45
+ "node_id": str(MyNode.__wrapped_node__.__id__),
46
+ "output_id": "c8fbe459-c9ee-4639-a82d-961180cf9411",
47
+ },
48
+ "type": "NODE_OUTPUT",
49
+ }
50
+ ],
51
+ }
52
+ assert error_input_value == {
53
+ "combinator": "OR",
54
+ "rules": [
55
+ {
56
+ "data": {
57
+ "node_id": str(MyNode.__wrapped_node__.__id__),
58
+ "output_id": "efe6e307-3ea4-4862-a26f-4c4416bb4537",
59
+ },
60
+ "type": "NODE_OUTPUT",
61
+ }
62
+ ],
63
+ }
@@ -1,13 +1,12 @@
1
1
  import pytest
2
2
  from uuid import UUID, uuid4
3
- from typing import List, cast
3
+ from typing import List
4
4
 
5
5
  from vellum.client.types.string_vellum_value import StringVellumValue
6
6
  from vellum.workflows.descriptors.base import BaseDescriptor
7
7
  from vellum.workflows.inputs import BaseInputs
8
8
  from vellum.workflows.nodes.bases import BaseNode
9
9
  from vellum.workflows.outputs import BaseOutputs
10
- from vellum.workflows.references import OutputReference, WorkflowInputReference
11
10
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
12
11
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
13
12
  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(
92
91
  entrypoint_node_display=NodeDisplayData(),
93
92
  ),
94
93
  global_workflow_input_displays={
95
- cast(WorkflowInputReference, Inputs.example_workflow_input): WorkflowInputsVellumDisplayOverrides(
94
+ Inputs.example_workflow_input: WorkflowInputsVellumDisplayOverrides(
96
95
  id=UUID("a154c29d-fac0-4cd0-ba88-bc52034f5470"),
97
96
  ),
98
97
  },
99
98
  global_node_output_displays={
100
- cast(OutputReference, MyNodeA.Outputs.output): (
99
+ MyNodeA.Outputs.output: (
101
100
  MyNodeA,
102
101
  NodeOutputDisplay(id=UUID("4b16a629-11a1-4b3f-a965-a57b872d13b8"), name="output"),
103
102
  ),
@@ -1,15 +1,17 @@
1
1
  from uuid import UUID
2
- from typing import Any, Callable, ClassVar, Generic, Optional, Type, TypeVar, cast
2
+ from typing import Any, Callable, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
3
3
 
4
4
  from vellum.workflows.nodes.bases.base import BaseNode
5
5
  from vellum.workflows.nodes.core.try_node.node import TryNode
6
6
  from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME, get_wrapped_node
7
+ from vellum.workflows.references.output import OutputReference
7
8
  from vellum.workflows.types.core import JsonObject
8
9
  from vellum.workflows.types.utils import get_original_base
9
10
  from vellum.workflows.utils.uuids import uuid4_from_hash
10
11
  from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
11
12
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
12
13
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
14
+ from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
13
15
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
14
16
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
15
17
 
@@ -22,9 +24,8 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
22
24
  def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
23
25
  node = self._node
24
26
 
25
- try:
26
- inner_node = get_wrapped_node(node)
27
- except AttributeError:
27
+ inner_node = get_wrapped_node(node)
28
+ if not inner_node:
28
29
  subworkflow = raise_if_descriptor(node.subworkflow)
29
30
  if not isinstance(subworkflow.graph, type) or not issubclass(subworkflow.graph, BaseNode):
30
31
  raise NotImplementedError(
@@ -57,6 +58,22 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
57
58
 
58
59
  return serialized_node
59
60
 
61
+ def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
62
+ inner_node = self._node.__wrapped_node__
63
+ if not inner_node:
64
+ return super().get_node_output_display(output)
65
+
66
+ node_display_class = get_node_display_class(BaseNodeVellumDisplay, inner_node)
67
+ node_display = node_display_class()
68
+ if output.name == "error":
69
+ return inner_node, NodeOutputDisplay(
70
+ id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
71
+ name="error",
72
+ )
73
+
74
+ inner_output = getattr(inner_node.Outputs, output.name)
75
+ return node_display.get_node_output_display(inner_output)
76
+
60
77
  @classmethod
61
78
  def wrap(cls, error_output_id: Optional[UUID] = None) -> Callable[..., Type["BaseTryNodeDisplay"]]:
62
79
  _error_output_id = error_output_id
@@ -28,8 +28,7 @@ from vellum.workflows.expressions.not_between import NotBetweenExpression
28
28
  from vellum.workflows.expressions.not_in import NotInExpression
29
29
  from vellum.workflows.expressions.or_ import OrExpression
30
30
  from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
31
- from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
32
- from vellum.workflows.references import NodeReference, OutputReference
31
+ from vellum.workflows.references import NodeReference
33
32
  from vellum.workflows.utils.uuids import uuid4_from_hash
34
33
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
35
34
  from vellum_ee.workflows.display.utils.vellum import create_node_input_value_pointer_rule
@@ -56,15 +55,6 @@ def create_node_input(
56
55
  pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
57
56
  ) -> NodeInput:
58
57
  input_id = str(input_id) if input_id else str(uuid4_from_hash(f"{node_id}|{input_name}"))
59
- if (
60
- isinstance(value, OutputReference)
61
- and value.outputs_class._node_class
62
- and has_wrapped_node(value.outputs_class._node_class)
63
- ):
64
- wrapped_node = get_wrapped_node(value.outputs_class._node_class)
65
- if wrapped_node._is_wrapped_node:
66
- value = getattr(wrapped_node.Outputs, value.name)
67
-
68
58
  rules = create_node_input_value_pointer_rules(value, display_context, pointer_type=pointer_type)
69
59
  return NodeInput(
70
60
  id=input_id,
@@ -685,8 +685,7 @@ def test_serialize_workflow__try_wrapped():
685
685
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
686
686
  },
687
687
  ],
688
- final_output_nodes,
689
- ignore_order=True,
688
+ sorted(final_output_nodes, key=lambda x: x["id"], reverse=True),
690
689
  )
691
690
 
692
691
  # AND each edge should be serialized correctly
@@ -11,7 +11,7 @@ from vellum.workflows.descriptors.base import BaseDescriptor
11
11
  from vellum.workflows.edges import Edge
12
12
  from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
13
13
  from vellum.workflows.nodes.bases import BaseNode
14
- from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
14
+ from vellum.workflows.nodes.utils import get_wrapped_node
15
15
  from vellum.workflows.ports import Port
16
16
  from vellum.workflows.references import OutputReference, WorkflowInputReference
17
17
  from vellum.workflows.types.core import JsonObject
@@ -138,19 +138,18 @@ class BaseWorkflowDisplay(
138
138
  ):
139
139
  """This method recursively adds nodes wrapped in decorators to the node_output_displays dictionary."""
140
140
 
141
+ inner_node = get_wrapped_node(node)
142
+ if inner_node:
143
+ inner_node_display = self._get_node_display(inner_node)
144
+ self._enrich_global_node_output_displays(inner_node, inner_node_display, node_output_displays)
145
+
141
146
  for node_output in node.Outputs:
142
147
  if node_output in node_output_displays:
143
148
  continue
144
149
 
145
- if has_wrapped_node(node):
146
- inner_node = get_wrapped_node(node)
147
- if inner_node._is_wrapped_node:
148
- inner_node_display = self._get_node_display(inner_node)
149
- self._enrich_global_node_output_displays(inner_node, inner_node_display, node_output_displays)
150
-
151
150
  # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
152
151
  # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
153
- node_output_displays[node_output] = node, node_display.get_node_output_display(node_output)
152
+ node_output_displays[node_output] = node_display.get_node_output_display(node_output)
154
153
 
155
154
  def _enrich_node_port_displays(
156
155
  self,
@@ -160,16 +159,15 @@ class BaseWorkflowDisplay(
160
159
  ):
161
160
  """This method recursively adds nodes wrapped in decorators to the port_displays dictionary."""
162
161
 
162
+ inner_node = get_wrapped_node(node)
163
+ if inner_node:
164
+ inner_node_display = self._get_node_display(inner_node)
165
+ self._enrich_node_port_displays(inner_node, inner_node_display, port_displays)
166
+
163
167
  for port in node.Ports:
164
168
  if port in port_displays:
165
169
  continue
166
170
 
167
- if has_wrapped_node(node):
168
- inner_node = get_wrapped_node(node)
169
- if inner_node._is_wrapped_node:
170
- inner_node_display = self._get_node_display(inner_node)
171
- self._enrich_node_port_displays(inner_node, inner_node_display, port_displays)
172
-
173
171
  port_displays[port] = node_display.get_node_port_display(port)
174
172
 
175
173
  def _get_node_display(self, node: Type[BaseNode]) -> NodeDisplayType:
@@ -216,13 +214,11 @@ class BaseWorkflowDisplay(
216
214
  global_node_displays[node] = node_display
217
215
 
218
216
  # Nodes wrapped in a decorator need to be in our node display dictionary for later retrieval
219
- if has_wrapped_node(node):
220
- inner_node = get_wrapped_node(node)
217
+ inner_node = get_wrapped_node(node)
218
+ if inner_node:
221
219
  inner_node_display = self._get_node_display(inner_node)
222
-
223
- if inner_node._is_wrapped_node:
224
- node_displays[inner_node] = inner_node_display
225
- global_node_displays[inner_node] = inner_node_display
220
+ node_displays[inner_node] = inner_node_display
221
+ global_node_displays[inner_node] = inner_node_display
226
222
 
227
223
  self._enrich_global_node_output_displays(node, node_display, global_node_output_displays)
228
224
  self._enrich_node_port_displays(node, node_display, port_displays)
@@ -7,7 +7,7 @@ from vellum.workflows.edges import Edge
7
7
  from vellum.workflows.nodes.bases import BaseNode
8
8
  from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
9
9
  from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
10
- from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node, has_wrapped_node
10
+ from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port
11
11
  from vellum.workflows.ports import Port
12
12
  from vellum.workflows.references import WorkflowInputReference
13
13
  from vellum.workflows.references.output import OutputReference
@@ -95,11 +95,8 @@ class VellumWorkflowDisplay(
95
95
  )
96
96
 
97
97
  # Add all the nodes in the workflow
98
- for node, node_display in self.display_context.node_displays.items():
99
- if getattr(node, "_is_wrapped_node") is True:
100
- # Nodes that are wrapped or decorated by other nodes are not serialized here
101
- # They are instead serialized by the wrapper node
102
- continue
98
+ for node in self._workflow.get_nodes():
99
+ node_display = self.display_context.node_displays[node]
103
100
 
104
101
  try:
105
102
  serialized_node = node_display.serialize(self.display_context)
@@ -295,13 +292,9 @@ class VellumWorkflowDisplay(
295
292
  else uuid4_from_hash(f"{self.workflow_id}|id|{entrypoint_node_id}")
296
293
  )
297
294
 
298
- if has_wrapped_node(entrypoint):
299
- wrapped_node = get_wrapped_node(entrypoint)
300
- if wrapped_node._is_wrapped_node:
301
- entrypoint = wrapped_node
302
-
303
- target_node_id = node_displays[entrypoint].node_id
304
- target_handle_id = node_displays[entrypoint].get_target_handle_id()
295
+ entrypoint_target = get_unadorned_node(entrypoint)
296
+ target_node_id = node_displays[entrypoint_target].node_id
297
+ target_handle_id = node_displays[entrypoint_target].get_target_handle_id()
305
298
 
306
299
  edge_display = self._generate_edge_display_from_source(
307
300
  entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides