vellum-ai 0.10.4__py3-none-any.whl → 0.10.6__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.
Files changed (56) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/workflows/nodes/__init__.py +6 -7
  3. vellum/workflows/nodes/bases/base.py +0 -1
  4. vellum/workflows/nodes/core/inline_subworkflow_node/node.py +1 -1
  5. vellum/workflows/nodes/core/templating_node/node.py +5 -1
  6. vellum/workflows/nodes/core/try_node/node.py +65 -27
  7. vellum/workflows/nodes/core/try_node/tests/test_node.py +17 -10
  8. vellum/workflows/nodes/displayable/__init__.py +2 -0
  9. vellum/workflows/nodes/displayable/bases/api_node/node.py +3 -3
  10. vellum/workflows/nodes/displayable/code_execution_node/node.py +5 -2
  11. vellum/workflows/nodes/displayable/final_output_node/node.py +6 -2
  12. vellum/workflows/nodes/displayable/note_node/__init__.py +5 -0
  13. vellum/workflows/nodes/displayable/note_node/node.py +10 -0
  14. vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py +10 -11
  15. vellum/workflows/nodes/utils.py +2 -0
  16. vellum/workflows/outputs/base.py +26 -2
  17. vellum/workflows/runner/runner.py +41 -27
  18. vellum/workflows/types/tests/test_utils.py +9 -0
  19. vellum/workflows/types/utils.py +1 -1
  20. vellum/workflows/utils/vellum_variables.py +13 -1
  21. vellum/workflows/workflows/base.py +24 -1
  22. {vellum_ai-0.10.4.dist-info → vellum_ai-0.10.6.dist-info}/METADATA +8 -6
  23. {vellum_ai-0.10.4.dist-info → vellum_ai-0.10.6.dist-info}/RECORD +56 -51
  24. vellum_cli/CONTRIBUTING.md +66 -0
  25. vellum_cli/README.md +3 -0
  26. vellum_ee/workflows/display/base.py +2 -1
  27. vellum_ee/workflows/display/nodes/base_node_display.py +27 -4
  28. vellum_ee/workflows/display/nodes/vellum/__init__.py +2 -0
  29. vellum_ee/workflows/display/nodes/vellum/api_node.py +3 -3
  30. vellum_ee/workflows/display/nodes/vellum/code_execution_node.py +4 -4
  31. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +86 -41
  32. vellum_ee/workflows/display/nodes/vellum/guardrail_node.py +3 -3
  33. vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py +4 -5
  34. vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py +9 -9
  35. vellum_ee/workflows/display/nodes/vellum/map_node.py +5 -5
  36. vellum_ee/workflows/display/nodes/vellum/note_node.py +32 -0
  37. vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py +5 -5
  38. vellum_ee/workflows/display/nodes/vellum/search_node.py +1 -1
  39. vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py +2 -2
  40. vellum_ee/workflows/display/nodes/vellum/templating_node.py +1 -1
  41. vellum_ee/workflows/display/nodes/vellum/try_node.py +16 -4
  42. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py +7 -3
  43. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py +127 -101
  44. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py +6 -5
  45. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +77 -64
  46. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py +4 -3
  47. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py +6 -6
  48. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py +6 -6
  49. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py +4 -3
  50. vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py +7 -6
  51. vellum_ee/workflows/display/workflows/base_workflow_display.py +14 -9
  52. vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py +2 -7
  53. vellum_ee/workflows/display/workflows/vellum_workflow_display.py +18 -16
  54. {vellum_ai-0.10.4.dist-info → vellum_ai-0.10.6.dist-info}/LICENSE +0 -0
  55. {vellum_ai-0.10.4.dist-info → vellum_ai-0.10.6.dist-info}/WHEEL +0 -0
  56. {vellum_ai-0.10.4.dist-info → vellum_ai-0.10.6.dist-info}/entry_points.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  from dataclasses import dataclass
2
2
  from uuid import UUID
3
- from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar
3
+ from typing import Any, ClassVar, Dict, Generic, List, Optional, TypeVar, Union, Tuple
4
4
 
5
5
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
6
6
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -29,7 +29,7 @@ from vellum.workflows.expressions.not_between import NotBetweenExpression
29
29
  from vellum.workflows.expressions.not_in import NotInExpression
30
30
  from vellum.workflows.expressions.or_ import OrExpression
31
31
  from vellum.workflows.nodes.displayable import ConditionalNode
32
- from vellum.workflows.types.core import JsonObject
32
+ from vellum.workflows.types.core import JsonObject, ConditionType
33
33
 
34
34
  _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
35
35
 
@@ -39,11 +39,13 @@ class RuleIdMap:
39
39
  id: UUID
40
40
  lhs: Optional["RuleIdMap"]
41
41
  rhs: Optional["RuleIdMap"]
42
+ field_node_input_id: Optional[UUID]
43
+ value_node_input_id: Optional[UUID]
42
44
 
43
45
  @dataclass
44
46
  class ConditionId:
45
47
  id: UUID
46
- rule_group_id: UUID
48
+ rule_group_id: Optional[UUID]
47
49
 
48
50
 
49
51
  class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Generic[_ConditionalNodeType]):
@@ -73,8 +75,11 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
73
75
  if rule_id_map is not None
74
76
  else str(uuid4_from_hash(f"{node_id}|rule|{','.join(str(p) for p in path)}"))
75
77
  )
76
- field_node_input_id = None
77
- value_node_input_id = None
78
+
79
+ result = self.get_nested_rule_details_by_path(rule_ids, path) if rule_ids else None
80
+ if result is None:
81
+ result = self._generate_hash_for_rule_ids(node_id, rule_id)
82
+ current_id, field_node_input_id, value_node_input_id = result
78
83
 
79
84
  # Recursive step: Keep recursing until we hit the other descriptors
80
85
  if isinstance(descriptor, (AndExpression, OrExpression)):
@@ -97,79 +102,89 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
97
102
  # Base cases for other descriptors
98
103
  elif isinstance(descriptor, (IsNullExpression, IsNotNullExpression)):
99
104
  expression_node_input = create_node_input(
100
- node_id, f"{rule_id}.field", descriptor._expression, display_context, None
105
+ node_id, f"{current_id}.field", descriptor._expression, display_context, field_node_input_id
101
106
  )
102
107
  node_inputs.append(expression_node_input)
103
- field_node_input_id = expression_node_input.id
108
+ field_node_input_id = UUID(expression_node_input.id)
104
109
  operator = self._convert_descriptor_to_operator(descriptor)
105
110
 
106
111
  elif isinstance(descriptor, (BetweenExpression, NotBetweenExpression)):
107
112
  field_node_input = create_node_input(
108
- node_id, f"{rule_id}.field", descriptor._value, display_context, None
113
+ node_id, f"{current_id}.field", descriptor._value, display_context, field_node_input_id
109
114
  )
110
115
  value_node_input = create_node_input(
111
- node_id, f"{rule_id}.value", f"{descriptor._start},{descriptor._end}", display_context, None
116
+ node_id, f"{current_id}.value", f"{descriptor._start},{descriptor._end}", display_context, value_node_input_id
112
117
  )
113
118
  node_inputs.extend([field_node_input, value_node_input])
114
119
  operator = self._convert_descriptor_to_operator(descriptor)
115
- field_node_input_id = field_node_input.id
116
- value_node_input_id = value_node_input.id
120
+ field_node_input_id = UUID(field_node_input.id)
121
+ value_node_input_id = UUID(value_node_input.id)
117
122
 
118
123
  else:
119
124
  lhs = descriptor._lhs # type: ignore[attr-defined]
120
125
  rhs = descriptor._rhs # type: ignore[attr-defined]
121
126
 
122
- lhs_node_input = create_node_input(node_id, f"{rule_id}.field", lhs, display_context, None)
127
+ lhs_node_input = create_node_input(node_id, f"{current_id}.field", lhs, display_context, field_node_input_id)
123
128
  node_inputs.append(lhs_node_input)
124
129
 
125
130
  if descriptor._rhs is not None: # type: ignore[attr-defined]
126
- rhs_node_input = create_node_input(node_id, f"{rule_id}.value", rhs, display_context, None)
131
+ rhs_node_input = create_node_input(node_id, f"{current_id}.value", rhs, display_context, value_node_input_id)
127
132
  node_inputs.append(rhs_node_input)
128
- value_node_input_id = rhs_node_input.id
133
+ value_node_input_id = UUID(rhs_node_input.id)
129
134
 
130
135
  operator = self._convert_descriptor_to_operator(descriptor)
131
- field_node_input_id = lhs_node_input.id
136
+ field_node_input_id = UUID(lhs_node_input.id)
132
137
 
133
138
  return {
134
- "id": rule_id,
139
+ "id": str(current_id),
135
140
  "rules": None,
136
141
  "combinator": None,
137
142
  "negated": False,
138
- "field_node_input_id": field_node_input_id,
143
+ "field_node_input_id": str(field_node_input_id),
139
144
  "operator": operator,
140
- "value_node_input_id": value_node_input_id,
145
+ "value_node_input_id": str(value_node_input_id),
141
146
  }
142
147
 
143
148
  conditions = []
144
149
  for idx, port in enumerate(node.Ports):
145
- if port._condition is None or port._condition_type is None:
146
- continue
147
-
148
150
 
149
151
  condition_id = str(condition_ids[idx].id if condition_ids else uuid4_from_hash(f"{node_id}|conditions|{idx}"))
150
152
  rule_group_id = str(condition_ids[idx].rule_group_id if condition_ids else uuid4_from_hash(f"{condition_id}|rule_group"))
151
153
  source_handle_id = str(source_handle_ids.get(idx) or uuid4_from_hash(f"{node_id}|handles|{idx}"))
152
154
 
153
- rule_ids = self._get_rule_ids()
154
- condition_rule = serialize_rule(port._condition, [idx], rule_ids[idx] if len(rule_ids) > idx else None)
155
- rules = condition_rule["rules"] if condition_rule["rules"] else [condition_rule]
156
-
157
- conditions.append(
158
- {
159
- "id": condition_id,
160
- "type": port._condition_type.value,
161
- "source_handle_id": source_handle_id,
162
- "data": {
163
- "id": rule_group_id,
164
- "rules": rules,
165
- "combinator": "AND",
166
- "negated": False,
167
- "field_node_input_id": None,
168
- "operator": None,
169
- "value_node_input_id": None,
170
- },
171
- }
172
- )
155
+ if port._condition is None:
156
+ if port._condition_type == ConditionType.ELSE:
157
+ conditions.append({
158
+ "id": condition_id,
159
+ "type": ConditionType.ELSE.value,
160
+ "source_handle_id": source_handle_id,
161
+ "data": None
162
+ })
163
+ else:
164
+ continue
165
+
166
+ else:
167
+ rule_ids = self._get_rule_ids()
168
+ condition_rule = serialize_rule(port._condition, [idx], rule_ids[idx] if len(rule_ids) > idx else None)
169
+ rules = condition_rule["rules"] if condition_rule["rules"] else [condition_rule]
170
+ if port._condition_type is None:
171
+ raise ValueError("Condition type is None, but a valid ConditionType is expected.")
172
+ conditions.append(
173
+ {
174
+ "id": condition_id,
175
+ "type": port._condition_type.value,
176
+ "source_handle_id": source_handle_id,
177
+ "data": {
178
+ "id": rule_group_id,
179
+ "rules": rules,
180
+ "combinator": "AND",
181
+ "negated": False,
182
+ "field_node_input_id": None,
183
+ "operator": None,
184
+ "value_node_input_id": None,
185
+ },
186
+ }
187
+ )
173
188
 
174
189
  return {
175
190
  "id": str(node_id),
@@ -225,6 +240,36 @@ class BaseConditionalNodeDisplay(BaseNodeVellumDisplay[_ConditionalNodeType], Ge
225
240
  else:
226
241
  raise ValueError(f"Unsupported descriptor type: {descriptor}")
227
242
 
243
+ def get_nested_rule_details_by_path(
244
+ self, rule_ids: List[RuleIdMap], path: List[int]
245
+ ) -> Union[Tuple[UUID, Optional[UUID], Optional[UUID]], None]:
246
+ current_rule = rule_ids[path[0]]
247
+
248
+ for step in path[1:]:
249
+ if step == 0 and current_rule.lhs:
250
+ current_rule = current_rule.lhs
251
+ elif step == 1 and current_rule.rhs:
252
+ current_rule = current_rule.rhs
253
+ else:
254
+ return None
255
+
256
+ # This is essentially a leaf
257
+ if current_rule.lhs and current_rule.lhs.lhs is None and current_rule.lhs.rhs is None:
258
+ return (
259
+ current_rule.lhs.id,
260
+ current_rule.lhs.field_node_input_id,
261
+ current_rule.lhs.value_node_input_id,
262
+ )
263
+
264
+ return None
265
+
266
+ def _generate_hash_for_rule_ids(self, node_id, rule_id) -> Tuple[UUID, UUID, UUID]:
267
+ return (
268
+ uuid4_from_hash(f"{node_id}|{rule_id}|current"),
269
+ uuid4_from_hash(f"{node_id}|{rule_id}||field"),
270
+ uuid4_from_hash(f"{node_id}|{rule_id}||value")
271
+ )
272
+
228
273
  def _get_source_handle_ids(self) -> Dict[int, UUID]:
229
274
  return self._get_explicit_node_display_attr("source_handle_ids", Dict[int, UUID]) or {}
230
275
 
@@ -1,12 +1,12 @@
1
1
  from uuid import UUID
2
2
  from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar
3
3
 
4
+ from vellum.workflows.nodes import GuardrailNode
5
+ from vellum.workflows.types.core import JsonObject
4
6
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
5
7
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
6
8
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
7
9
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
8
- from vellum.workflows.nodes import GuardrailNode
9
- from vellum.workflows.types.core import JsonObject
10
10
 
11
11
  _GuardrailNodeType = TypeVar("_GuardrailNodeType", bound=GuardrailNode)
12
12
 
@@ -15,7 +15,7 @@ class BaseGuardrailNodeDisplay(BaseNodeVellumDisplay[_GuardrailNodeType], Generi
15
15
  metric_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
16
16
 
17
17
  def serialize(
18
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
18
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
19
19
  ) -> JsonObject:
20
20
  node = self._node
21
21
  node_id = self.node_id
@@ -2,7 +2,9 @@ from uuid import UUID
2
2
  from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, cast
3
3
 
4
4
  from vellum import PromptBlock, RichTextChildBlock, VellumVariable
5
-
5
+ from vellum.workflows.nodes import InlinePromptNode
6
+ from vellum.workflows.references import OutputReference
7
+ from vellum.workflows.types.core import JsonObject
6
8
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
7
9
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
8
10
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -10,9 +12,6 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
10
12
  from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
11
13
  from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
12
14
  from vellum_ee.workflows.display.vellum import NodeInput
13
- from vellum.workflows.nodes import InlinePromptNode
14
- from vellum.workflows.references import OutputReference
15
- from vellum.workflows.types.core import JsonObject
16
15
 
17
16
  _InlinePromptNodeType = TypeVar("_InlinePromptNodeType", bound=InlinePromptNode)
18
17
 
@@ -23,7 +22,7 @@ class BaseInlinePromptNodeDisplay(BaseNodeVellumDisplay[_InlinePromptNodeType],
23
22
  prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
24
23
 
25
24
  def serialize(
26
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
25
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
27
26
  ) -> JsonObject:
28
27
  node = self._node
29
28
  node_id = self.node_id
@@ -1,17 +1,16 @@
1
1
  from uuid import UUID
2
- from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar
2
+ from typing import Any, ClassVar, Dict, Generic, List, Optional, Tuple, Type, TypeVar, cast
3
3
 
4
4
  from vellum import VellumVariable
5
-
5
+ from vellum.workflows.nodes import InlineSubworkflowNode
6
+ from vellum.workflows.types.core import JsonObject
6
7
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
7
8
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
8
9
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
9
10
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
10
11
  from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
11
- from vellum_ee.workflows.display.vellum import NodeInput
12
+ from vellum_ee.workflows.display.vellum import NodeInput, WorkflowOutputVellumDisplay
12
13
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
13
- from vellum.workflows.nodes import InlineSubworkflowNode
14
- from vellum.workflows.types.core import JsonObject
15
14
 
16
15
  _InlineSubworkflowNodeType = TypeVar("_InlineSubworkflowNodeType", bound=InlineSubworkflowNode)
17
16
 
@@ -22,19 +21,19 @@ class BaseInlineSubworkflowNodeDisplay(
22
21
  workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
23
22
 
24
23
  def serialize(
25
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
24
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
26
25
  ) -> JsonObject:
27
26
  node = self._node
28
27
  node_id = self.node_id
29
28
 
30
29
  node_inputs, workflow_inputs = self._generate_node_and_workflow_inputs(node_id, node, display_context)
31
- workflow_outputs = self._generate_workflow_outputs(node)
32
30
 
33
31
  subworkflow_display = get_workflow_display(
34
32
  base_display_class=display_context.workflow_display_class,
35
33
  workflow_class=raise_if_descriptor(node.subworkflow),
36
34
  parent_display_context=display_context,
37
35
  )
36
+ workflow_outputs = self._generate_workflow_outputs(node, subworkflow_display.display_context)
38
37
  serialized_subworkflow = subworkflow_display.serialize()
39
38
 
40
39
  return {
@@ -87,13 +86,14 @@ class BaseInlineSubworkflowNodeDisplay(
87
86
  def _generate_workflow_outputs(
88
87
  self,
89
88
  node: Type[InlineSubworkflowNode],
89
+ display_context: WorkflowDisplayContext,
90
90
  ) -> List[VellumVariable]:
91
91
  workflow_outputs: List[VellumVariable] = []
92
92
  for output_descriptor in raise_if_descriptor(node.subworkflow).Outputs: # type: ignore[union-attr]
93
- node_output_display = self.get_node_output_display(output_descriptor)
93
+ workflow_output_display = cast(WorkflowOutputVellumDisplay, display_context.workflow_output_displays[output_descriptor])
94
94
  output_type = infer_vellum_variable_type(output_descriptor)
95
95
  workflow_outputs.append(
96
- VellumVariable(id=str(node_output_display.id), key=node_output_display.name, type=output_type)
96
+ VellumVariable(id=str(workflow_output_display.id), key=workflow_output_display.name, type=output_type)
97
97
  )
98
98
 
99
99
  return workflow_outputs
@@ -2,7 +2,8 @@ from uuid import UUID
2
2
  from typing import Any, ClassVar, Dict, Generic, List, Optional, Type, TypeVar
3
3
 
4
4
  from vellum import VellumVariable
5
-
5
+ from vellum.workflows.nodes import MapNode
6
+ from vellum.workflows.types.core import JsonObject
6
7
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
7
8
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
8
9
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -10,8 +11,6 @@ from vellum_ee.workflows.display.types import WorkflowDisplayContext
10
11
  from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
11
12
  from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
12
13
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
13
- from vellum.workflows.nodes import MapNode
14
- from vellum.workflows.types.core import JsonObject
15
14
 
16
15
  _MapNodeType = TypeVar("_MapNodeType", bound=MapNode)
17
16
 
@@ -20,7 +19,7 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
20
19
  workflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
21
20
 
22
21
  def serialize(
23
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
22
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
24
23
  ) -> JsonObject:
25
24
  node = self._node
26
25
  node_id = self.node_id
@@ -32,7 +31,8 @@ class BaseMapNodeDisplay(BaseNodeVellumDisplay[_MapNodeType], Generic[_MapNodeTy
32
31
  # In Vellum it's always 'items'
33
32
  variable_name = descriptor.name if descriptor.name != "all_items" else "items"
34
33
  variable_id = str(
35
- self.workflow_input_ids_by_name.get(variable_name) or uuid4_from_hash(f"{self.node_id}|{variable_name}")
34
+ self.workflow_input_ids_by_name.get(variable_name)
35
+ or uuid4_from_hash(f"{self.node_id}|{variable_name}")
36
36
  )
37
37
  workflow_inputs.append(
38
38
  VellumVariable(
@@ -0,0 +1,32 @@
1
+ import json
2
+ from typing import Any, ClassVar, Dict, Generic, TypeVar, Union
3
+
4
+ from vellum.workflows.nodes import NoteNode
5
+ from vellum.workflows.types.core import JsonObject
6
+ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
7
+ from vellum_ee.workflows.display.types import WorkflowDisplayContext
8
+
9
+ _NoteNodeType = TypeVar("_NoteNodeType", bound=NoteNode)
10
+
11
+
12
+ class BaseNoteNodeDisplay(BaseNodeVellumDisplay[_NoteNodeType], Generic[_NoteNodeType]):
13
+ text: ClassVar[str] = ""
14
+ style: ClassVar[Union[Dict[str, Any], None]] = None
15
+
16
+ def serialize(
17
+ self, display_context: WorkflowDisplayContext, **kwargs: Any
18
+ ) -> JsonObject:
19
+ node_id = self.node_id
20
+
21
+ return {
22
+ "id": str(node_id),
23
+ "type": "NOTE",
24
+ "inputs": [],
25
+ "data": {
26
+ "label": self.label,
27
+ "text": self.text,
28
+ "style": json.dumps(self.style) if self.style else None,
29
+ },
30
+ "display_data": self.get_display_data().dict(),
31
+ "definition": self.get_definition().dict(),
32
+ }
@@ -1,14 +1,14 @@
1
1
  from uuid import UUID
2
2
  from typing import Any, ClassVar, Dict, Generic, Optional, TypeVar, cast
3
3
 
4
- from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
5
- from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
6
- from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
7
- from vellum_ee.workflows.display.types import WorkflowDisplayContext
8
4
  from vellum.workflows.nodes.displayable.prompt_deployment_node import PromptDeploymentNode
9
5
  from vellum.workflows.references import OutputReference
10
6
  from vellum.workflows.types.core import JsonObject
11
7
  from vellum.workflows.vellum_client import create_vellum_client
8
+ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
9
+ from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
10
+ from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
11
+ from vellum_ee.workflows.display.types import WorkflowDisplayContext
12
12
 
13
13
  _PromptDeploymentNodeType = TypeVar("_PromptDeploymentNodeType", bound=PromptDeploymentNode)
14
14
 
@@ -21,7 +21,7 @@ class BasePromptDeploymentNodeDisplay(
21
21
  prompt_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
22
22
 
23
23
  def serialize(
24
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
24
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
25
25
  ) -> JsonObject:
26
26
  node = self._node
27
27
  node_id = self.node_id
@@ -31,7 +31,7 @@ class BaseSearchNodeDisplay(BaseNodeVellumDisplay[_SearchNodeType], Generic[_Sea
31
31
  variable_ids: Optional[VariableIdMap] = None
32
32
 
33
33
  def serialize(
34
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
34
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
35
35
  ) -> JsonObject:
36
36
  node = self._node
37
37
  node_id = self.node_id
@@ -16,9 +16,9 @@ class BaseSubworkflowDeploymentNodeDisplay(
16
16
  BaseNodeVellumDisplay[_SubworkflowDeploymentNodeType], Generic[_SubworkflowDeploymentNodeType]
17
17
  ):
18
18
  subworkflow_input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
19
-
19
+
20
20
  def serialize(
21
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
21
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
22
22
  ) -> JsonObject:
23
23
  node = self._node
24
24
  node_id = self.node_id
@@ -17,7 +17,7 @@ class BaseTemplatingNodeDisplay(BaseNodeVellumDisplay[_TemplatingNodeType], Gene
17
17
  input_ids_by_name: ClassVar[Dict[str, UUID]] = {}
18
18
 
19
19
  def serialize(
20
- self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs: Any
20
+ self, display_context: WorkflowDisplayContext, error_output_id: Optional[UUID] = None, **kwargs
21
21
  ) -> JsonObject:
22
22
  node = self._node
23
23
  node_id = self.node_id
@@ -1,13 +1,13 @@
1
1
  from uuid import UUID
2
2
  from typing import Any, ClassVar, Generic, Optional, TypeVar
3
3
 
4
+ from vellum.workflows.nodes.core.try_node.node import TryNode
5
+ from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME, get_wrapped_node
6
+ from vellum.workflows.types.core import JsonObject
4
7
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
5
8
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
6
9
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
7
10
  from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
8
- from vellum.workflows.nodes.core.try_node.node import TryNode
9
- from vellum.workflows.nodes.utils import get_wrapped_node
10
- from vellum.workflows.types.core import JsonObject
11
11
 
12
12
  _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
13
13
 
@@ -32,7 +32,19 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
32
32
 
33
33
  serialized_node = node_display.serialize(
34
34
  display_context,
35
- error_output_id=str(self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id")),
35
+ error_output_id=self.error_output_id or uuid4_from_hash(f"{node_display.node_id}|error_output_id"),
36
36
  )
37
37
 
38
+ serialized_node_definition = serialized_node.get("definition")
39
+ if isinstance(serialized_node_definition, dict):
40
+ serialized_node_definition_module = serialized_node_definition.get("module")
41
+ if isinstance(serialized_node_definition_module, list):
42
+ serialized_node_definition_module.extend(
43
+ [
44
+ serialized_node_definition["name"],
45
+ ADORNMENT_MODULE_NAME,
46
+ ]
47
+ )
48
+ serialized_node_definition["name"] = node.__name__
49
+
38
50
  return serialized_node
@@ -1,10 +1,12 @@
1
1
  from deepdiff import DeepDiff
2
2
 
3
- from tests.workflows.basic_code_execution_node.try_workflow import TrySimpleCodeExecutionWorkflow
4
- from tests.workflows.basic_code_execution_node.workflow import SimpleCodeExecutionWorkflow
3
+ from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
5
4
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
6
5
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
7
6
 
7
+ from tests.workflows.basic_code_execution_node.try_workflow import TrySimpleCodeExecutionWorkflow
8
+ from tests.workflows.basic_code_execution_node.workflow import SimpleCodeExecutionWorkflow
9
+
8
10
 
9
11
  def test_serialize_workflow():
10
12
  # GIVEN a Workflow with a code execution node
@@ -428,8 +430,10 @@ def test_serialize_workflow__try_wrapped():
428
430
  "workflows",
429
431
  "basic_code_execution_node",
430
432
  "try_workflow",
433
+ "SimpleCodeExecutionNode",
434
+ ADORNMENT_MODULE_NAME,
431
435
  ],
432
- "name": "SimpleCodeExecutionNode",
436
+ "name": "TryNode",
433
437
  },
434
438
  }
435
439