vellum-ai 0.11.9__py3-none-any.whl → 0.12.0__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 (62) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/workflows/descriptors/base.py +2 -2
  3. vellum/workflows/descriptors/tests/test_utils.py +4 -4
  4. vellum/workflows/errors/__init__.py +3 -3
  5. vellum/workflows/errors/types.py +46 -3
  6. vellum/workflows/events/node.py +3 -3
  7. vellum/workflows/events/tests/test_event.py +3 -3
  8. vellum/workflows/events/workflow.py +3 -3
  9. vellum/workflows/exceptions.py +8 -4
  10. vellum/workflows/nodes/bases/base.py +9 -2
  11. vellum/workflows/nodes/bases/tests/test_base_node.py +13 -0
  12. vellum/workflows/nodes/core/error_node/node.py +9 -5
  13. vellum/workflows/nodes/core/inline_subworkflow_node/node.py +3 -13
  14. vellum/workflows/nodes/core/map_node/node.py +2 -2
  15. vellum/workflows/nodes/core/retry_node/node.py +5 -5
  16. vellum/workflows/nodes/core/retry_node/tests/test_node.py +6 -6
  17. vellum/workflows/nodes/core/templating_node/node.py +2 -2
  18. vellum/workflows/nodes/core/try_node/node.py +7 -7
  19. vellum/workflows/nodes/core/try_node/tests/test_node.py +9 -7
  20. vellum/workflows/nodes/displayable/bases/api_node/node.py +3 -3
  21. vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py +4 -12
  22. vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py +2 -2
  23. vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py +2 -2
  24. vellum/workflows/nodes/displayable/bases/search_node.py +3 -3
  25. vellum/workflows/nodes/displayable/code_execution_node/node.py +21 -5
  26. vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py +141 -0
  27. vellum/workflows/nodes/displayable/guardrail_node/node.py +3 -3
  28. vellum/workflows/nodes/displayable/inline_prompt_node/node.py +3 -3
  29. vellum/workflows/nodes/displayable/prompt_deployment_node/node.py +3 -3
  30. vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py +7 -14
  31. vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py +4 -4
  32. vellum/workflows/nodes/utils.py +5 -9
  33. vellum/workflows/references/external_input.py +2 -2
  34. vellum/workflows/references/node.py +2 -2
  35. vellum/workflows/references/state_value.py +2 -2
  36. vellum/workflows/references/workflow_input.py +2 -2
  37. vellum/workflows/runner/runner.py +15 -15
  38. {vellum_ee/workflows/display → vellum/workflows}/utils/tests/test_uuids.py +1 -1
  39. vellum/workflows/workflows/base.py +7 -7
  40. {vellum_ai-0.11.9.dist-info → vellum_ai-0.12.0.dist-info}/METADATA +1 -1
  41. {vellum_ai-0.11.9.dist-info → vellum_ai-0.12.0.dist-info}/RECORD +61 -62
  42. {vellum_ai-0.11.9.dist-info → vellum_ai-0.12.0.dist-info}/WHEEL +1 -1
  43. vellum_ee/workflows/display/nodes/base_node_display.py +50 -21
  44. vellum_ee/workflows/display/nodes/base_node_vellum_display.py +10 -1
  45. vellum_ee/workflows/display/nodes/get_node_display_class.py +10 -1
  46. vellum_ee/workflows/display/nodes/tests/test_base_node_display.py +5 -4
  47. vellum_ee/workflows/display/nodes/vellum/code_execution_node.py +13 -3
  48. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +1 -1
  49. vellum_ee/workflows/display/nodes/vellum/final_output_node.py +1 -1
  50. vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py +1 -1
  51. vellum_ee/workflows/display/nodes/vellum/merge_node.py +4 -4
  52. vellum_ee/workflows/display/nodes/vellum/search_node.py +1 -1
  53. vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py +5 -1
  54. vellum_ee/workflows/display/nodes/vellum/try_node.py +12 -6
  55. vellum_ee/workflows/display/nodes/vellum/utils.py +1 -1
  56. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py +269 -85
  57. vellum_ee/workflows/display/workflows/base_workflow_display.py +2 -2
  58. vellum_ee/workflows/display/workflows/vellum_workflow_display.py +2 -6
  59. vellum_ee/workflows/display/utils/tests/__init__.py +0 -0
  60. {vellum_ee/workflows/display → vellum/workflows}/utils/uuids.py +0 -0
  61. {vellum_ai-0.11.9.dist-info → vellum_ai-0.12.0.dist-info}/LICENSE +0 -0
  62. {vellum_ai-0.11.9.dist-info → vellum_ai-0.12.0.dist-info}/entry_points.txt +0 -0
@@ -5,14 +5,15 @@ from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
5
5
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
6
6
 
7
7
  from tests.workflows.basic_code_execution_node.try_workflow import TrySimpleCodeExecutionWorkflow
8
- from tests.workflows.basic_code_execution_node.workflow import SimpleCodeExecutionWorkflow
8
+ from tests.workflows.basic_code_execution_node.workflow import SimpleCodeExecutionWithFilepathWorkflow
9
+ from tests.workflows.basic_code_execution_node.workflow_with_code import SimpleCodeExecutionWithCodeWorkflow
9
10
 
10
11
 
11
- def test_serialize_workflow():
12
+ def test_serialize_workflow_with_filepath():
12
13
  # GIVEN a Workflow with a code execution node
13
14
  # WHEN we serialize it
14
15
  workflow_display = get_workflow_display(
15
- base_display_class=VellumWorkflowDisplay, workflow_class=SimpleCodeExecutionWorkflow
16
+ base_display_class=VellumWorkflowDisplay, workflow_class=SimpleCodeExecutionWithFilepathWorkflow
16
17
  )
17
18
  serialized_workflow: dict = workflow_display.serialize()
18
19
 
@@ -32,8 +33,8 @@ def test_serialize_workflow():
32
33
  assert len(output_variables) == 2
33
34
  assert not DeepDiff(
34
35
  [
35
- {"id": "ecdb8bf3-f456-400a-94af-a8f8338096b8", "key": "log", "type": "STRING"},
36
- {"id": "461585c0-0f65-4d18-ba69-fdfb3874d2b3", "key": "result", "type": "NUMBER"},
36
+ {"id": "1cee930f-342f-421c-89fc-ff212b3764bb", "key": "log", "type": "STRING"},
37
+ {"id": "f6a3e3e0-f83f-4491-8b7a-b20fddd7160c", "key": "result", "type": "NUMBER"},
37
38
  ],
38
39
  output_variables,
39
40
  ignore_order=True,
@@ -48,27 +49,12 @@ def test_serialize_workflow():
48
49
  # AND each node should be serialized correctly
49
50
  entrypoint_node = workflow_raw_data["nodes"][0]
50
51
  assert entrypoint_node == {
51
- "id": "6d94a007-9b25-4b4f-a31d-74e7965b6696",
52
+ "id": "bd18f11c-5f7a-45d5-9970-0b1cf10d3761",
52
53
  "type": "ENTRYPOINT",
53
54
  "inputs": [],
54
- "data": {
55
- "label": "Entrypoint Node",
56
- "source_handle_id": "5474e6da-f14c-458f-a82b-7a05fdfe3e5b",
57
- },
58
- "definition": {
59
- "bases": [],
60
- "module": [
61
- "vellum",
62
- "workflows",
63
- "nodes",
64
- "bases",
65
- "base",
66
- ],
67
- "name": "BaseNode",
68
- },
69
- "display_data": {
70
- "position": {"x": 0.0, "y": 0.0},
71
- },
55
+ "data": {"label": "Entrypoint Node", "source_handle_id": "118e4298-aa79-467c-b8b4-2df540905e86"},
56
+ "display_data": {"position": {"x": 0.0, "y": 0.0}},
57
+ "definition": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []},
72
58
  }
73
59
 
74
60
  code_execution_node = workflow_raw_data["nodes"][1]
@@ -134,23 +120,22 @@ def test_serialize_workflow():
134
120
  "name": "SimpleCodeExecutionNode",
135
121
  },
136
122
  }
137
-
138
123
  assert not DeepDiff(
139
124
  [
140
125
  {
141
- "id": "e81b713c-b356-4654-9715-e8c54e3ca267",
126
+ "id": "994d5c2e-00d2-4dff-9a9d-804766d03698",
142
127
  "type": "TERMINAL",
143
128
  "data": {
144
129
  "label": "Final Output",
145
130
  "name": "result",
146
- "target_handle_id": "c0890e2a-8827-4f99-8e9a-144c7e937f33",
147
- "output_id": "461585c0-0f65-4d18-ba69-fdfb3874d2b3",
131
+ "target_handle_id": "30fb0f4a-61c3-49de-a0aa-7dfdcee6ea07",
132
+ "output_id": "f6a3e3e0-f83f-4491-8b7a-b20fddd7160c",
148
133
  "output_type": "NUMBER",
149
- "node_input_id": "c28cc3c7-1285-4be3-ac51-c73456cee961",
134
+ "node_input_id": "ae302487-ff2a-457a-81ed-9e0348e91833",
150
135
  },
151
136
  "inputs": [
152
137
  {
153
- "id": "c28cc3c7-1285-4be3-ac51-c73456cee961",
138
+ "id": "ae302487-ff2a-457a-81ed-9e0348e91833",
154
139
  "key": "node_input",
155
140
  "value": {
156
141
  "rules": [
@@ -168,44 +153,27 @@ def test_serialize_workflow():
168
153
  ],
169
154
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
170
155
  "definition": {
156
+ "name": "FinalOutputNode",
157
+ "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
171
158
  "bases": [
172
- {
173
- "bases": [],
174
- "module": [
175
- "vellum",
176
- "workflows",
177
- "nodes",
178
- "bases",
179
- "base",
180
- ],
181
- "name": "BaseNode",
182
- },
159
+ {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}
183
160
  ],
184
- "module": [
185
- "vellum",
186
- "workflows",
187
- "nodes",
188
- "displayable",
189
- "final_output_node",
190
- "node",
191
- ],
192
- "name": "FinalOutputNode",
193
161
  },
194
162
  },
195
163
  {
196
- "id": "59e52e86-8059-4f32-b1f0-54f500b167a9",
164
+ "id": "c6e3aced-1fc9-48d2-ae55-d2a880e359cb",
197
165
  "type": "TERMINAL",
198
166
  "data": {
199
167
  "label": "Final Output",
200
168
  "name": "log",
201
- "target_handle_id": "bb56d2ed-f5e2-4c00-a2ef-445a2210d6d1",
202
- "output_id": "ecdb8bf3-f456-400a-94af-a8f8338096b8",
169
+ "target_handle_id": "1e126004-9de7-42c0-b1e1-87f9eb0642e2",
170
+ "output_id": "1cee930f-342f-421c-89fc-ff212b3764bb",
203
171
  "output_type": "STRING",
204
- "node_input_id": "e34369ff-8795-4f4f-8f40-c535a1c368bf",
172
+ "node_input_id": "c6593516-ffc5-49a8-8a65-1038cccec3f8",
205
173
  },
206
174
  "inputs": [
207
175
  {
208
- "id": "e34369ff-8795-4f4f-8f40-c535a1c368bf",
176
+ "id": "c6593516-ffc5-49a8-8a65-1038cccec3f8",
209
177
  "key": "node_input",
210
178
  "value": {
211
179
  "rules": [
@@ -223,28 +191,11 @@ def test_serialize_workflow():
223
191
  ],
224
192
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
225
193
  "definition": {
194
+ "name": "FinalOutputNode",
195
+ "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
226
196
  "bases": [
227
- {
228
- "bases": [],
229
- "module": [
230
- "vellum",
231
- "workflows",
232
- "nodes",
233
- "bases",
234
- "base",
235
- ],
236
- "name": "BaseNode",
237
- },
238
- ],
239
- "module": [
240
- "vellum",
241
- "workflows",
242
- "nodes",
243
- "displayable",
244
- "final_output_node",
245
- "node",
197
+ {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}
246
198
  ],
247
- "name": "FinalOutputNode",
248
199
  },
249
200
  },
250
201
  ],
@@ -257,27 +208,27 @@ def test_serialize_workflow():
257
208
  assert not DeepDiff(
258
209
  [
259
210
  {
260
- "id": "7cb7522b-3d91-474c-a3ad-6170b8d0def1",
261
- "source_node_id": "6d94a007-9b25-4b4f-a31d-74e7965b6696",
262
- "source_handle_id": "5474e6da-f14c-458f-a82b-7a05fdfe3e5b",
211
+ "id": "32673715-d88c-4727-b284-21ae4efe3f85",
212
+ "source_node_id": "bd18f11c-5f7a-45d5-9970-0b1cf10d3761",
213
+ "source_handle_id": "118e4298-aa79-467c-b8b4-2df540905e86",
263
214
  "target_node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
264
215
  "target_handle_id": "e02a2701-22c0-4533-8b00-175998e7350a",
265
216
  "type": "DEFAULT",
266
217
  },
267
218
  {
268
- "id": "1b2ef8d2-9220-4aa7-9f09-fced39002337",
219
+ "id": "d1e66711-75b3-41c3-beb6-424894fdd307",
269
220
  "source_node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
270
221
  "source_handle_id": "832f81ec-427b-42a8-825c-e62c43c1f961",
271
- "target_node_id": "59e52e86-8059-4f32-b1f0-54f500b167a9",
272
- "target_handle_id": "bb56d2ed-f5e2-4c00-a2ef-445a2210d6d1",
222
+ "target_node_id": "994d5c2e-00d2-4dff-9a9d-804766d03698",
223
+ "target_handle_id": "30fb0f4a-61c3-49de-a0aa-7dfdcee6ea07",
273
224
  "type": "DEFAULT",
274
225
  },
275
226
  {
276
- "id": "caa519a9-3c17-45f2-b67f-7dc656780300",
227
+ "id": "67d4c43e-80f9-4875-b6ab-9ecbba19fc7a",
277
228
  "source_node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
278
229
  "source_handle_id": "832f81ec-427b-42a8-825c-e62c43c1f961",
279
- "target_node_id": "e81b713c-b356-4654-9715-e8c54e3ca267",
280
- "target_handle_id": "c0890e2a-8827-4f99-8e9a-144c7e937f33",
230
+ "target_node_id": "c6e3aced-1fc9-48d2-ae55-d2a880e359cb",
231
+ "target_handle_id": "1e126004-9de7-42c0-b1e1-87f9eb0642e2",
281
232
  "type": "DEFAULT",
282
233
  },
283
234
  ],
@@ -298,7 +249,7 @@ def test_serialize_workflow():
298
249
  # AND the definition should be what we expect
299
250
  definition = workflow_raw_data["definition"]
300
251
  assert definition == {
301
- "name": "SimpleCodeExecutionWorkflow",
252
+ "name": "SimpleCodeExecutionWithFilepathWorkflow",
302
253
  "module": [
303
254
  "tests",
304
255
  "workflows",
@@ -308,6 +259,239 @@ def test_serialize_workflow():
308
259
  }
309
260
 
310
261
 
262
+ def test_serialize_workflow_with_code():
263
+ # GIVEN a Workflow with a code execution node
264
+ # WHEN we serialize it
265
+ workflow_display = get_workflow_display(
266
+ base_display_class=VellumWorkflowDisplay, workflow_class=SimpleCodeExecutionWithCodeWorkflow
267
+ )
268
+ serialized_workflow: dict = workflow_display.serialize()
269
+
270
+ # THEN we should get a serialized representation of the Workflow
271
+ assert serialized_workflow.keys() == {
272
+ "workflow_raw_data",
273
+ "input_variables",
274
+ "output_variables",
275
+ }
276
+
277
+ # AND its input variables should be what we expect
278
+ input_variables = serialized_workflow["input_variables"]
279
+ assert len(input_variables) == 0
280
+
281
+ # AND its output variables should be what we expect
282
+ output_variables = serialized_workflow["output_variables"]
283
+ assert len(output_variables) == 2
284
+ assert not DeepDiff(
285
+ [
286
+ {"id": "283d6849-f3ed-4beb-b261-cf70f90e8d10", "key": "result", "type": "NUMBER"},
287
+ {"id": "4c136180-050b-4422-a7a4-2a1c6729042c", "key": "log", "type": "STRING"},
288
+ ],
289
+ output_variables,
290
+ ignore_order=True,
291
+ )
292
+
293
+ # AND its raw data should be what we expect
294
+ workflow_raw_data = serialized_workflow["workflow_raw_data"]
295
+ assert workflow_raw_data.keys() == {"edges", "nodes", "display_data", "definition"}
296
+ assert len(workflow_raw_data["edges"]) == 3
297
+ assert len(workflow_raw_data["nodes"]) == 4
298
+
299
+ # AND each node should be serialized correctly
300
+ entrypoint_node = workflow_raw_data["nodes"][0]
301
+ assert entrypoint_node == {
302
+ "id": "22555158-d8ba-41b4-a6fc-87c3b25bd073",
303
+ "type": "ENTRYPOINT",
304
+ "inputs": [],
305
+ "data": {"label": "Entrypoint Node", "source_handle_id": "e82390bb-c68c-48c1-9f87-7fbfff494c45"},
306
+ "display_data": {"position": {"x": 0.0, "y": 0.0}},
307
+ "definition": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []},
308
+ }
309
+
310
+ code_execution_node = workflow_raw_data["nodes"][1]
311
+ assert code_execution_node == {
312
+ "id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
313
+ "type": "CODE_EXECUTION",
314
+ "inputs": [
315
+ {
316
+ "id": "f2e8a4fa-b54e-41e9-b314-0e5443519ac7",
317
+ "key": "code",
318
+ "value": {
319
+ "rules": [
320
+ {
321
+ "type": "CONSTANT_VALUE",
322
+ "data": {"type": "STRING", "value": 'def main() -> str:\n return "Hello, World!"\n'},
323
+ }
324
+ ],
325
+ "combinator": "OR",
326
+ },
327
+ },
328
+ {
329
+ "id": "19d64948-f22b-4103-a7f5-3add184b31cc",
330
+ "key": "runtime",
331
+ "value": {
332
+ "rules": [{"type": "CONSTANT_VALUE", "data": {"type": "STRING", "value": "PYTHON_3_11_6"}}],
333
+ "combinator": "OR",
334
+ },
335
+ },
336
+ ],
337
+ "data": {
338
+ "label": "Simple Code Execution Node",
339
+ "error_output_id": None,
340
+ "source_handle_id": "832f81ec-427b-42a8-825c-e62c43c1f961",
341
+ "target_handle_id": "e02a2701-22c0-4533-8b00-175998e7350a",
342
+ "code_input_id": "f2e8a4fa-b54e-41e9-b314-0e5443519ac7",
343
+ "runtime_input_id": "19d64948-f22b-4103-a7f5-3add184b31cc",
344
+ "output_type": "NUMBER",
345
+ "packages": [],
346
+ "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
347
+ "log_output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
348
+ },
349
+ "display_data": {"position": {"x": 0.0, "y": 0.0}},
350
+ "definition": {
351
+ "name": "SimpleCodeExecutionNode",
352
+ "module": ["tests", "workflows", "basic_code_execution_node", "workflow_with_code"],
353
+ "bases": [
354
+ {
355
+ "name": "CodeExecutionNode",
356
+ "module": ["vellum", "workflows", "nodes", "displayable", "code_execution_node", "node"],
357
+ }
358
+ ],
359
+ },
360
+ }
361
+ assert not DeepDiff(
362
+ [
363
+ {
364
+ "id": "52f285fe-1f52-4920-b01b-499762b95220",
365
+ "type": "TERMINAL",
366
+ "data": {
367
+ "label": "Final Output",
368
+ "name": "result",
369
+ "target_handle_id": "de8f2cc2-8c32-4782-87d5-4eb5afcd42e3",
370
+ "output_id": "283d6849-f3ed-4beb-b261-cf70f90e8d10",
371
+ "output_type": "NUMBER",
372
+ "node_input_id": "b38ba7a8-0b2a-4146-8d58-9fa0bcba8cd5",
373
+ },
374
+ "inputs": [
375
+ {
376
+ "id": "b38ba7a8-0b2a-4146-8d58-9fa0bcba8cd5",
377
+ "key": "node_input",
378
+ "value": {
379
+ "rules": [
380
+ {
381
+ "type": "NODE_OUTPUT",
382
+ "data": {
383
+ "node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
384
+ "output_id": "0fde9607-353f-42c2-85c4-20f720ebc1ec",
385
+ },
386
+ }
387
+ ],
388
+ "combinator": "OR",
389
+ },
390
+ }
391
+ ],
392
+ "display_data": {"position": {"x": 0.0, "y": 0.0}},
393
+ "definition": {
394
+ "name": "FinalOutputNode",
395
+ "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
396
+ "bases": [
397
+ {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}
398
+ ],
399
+ },
400
+ },
401
+ {
402
+ "id": "eccf97c7-e766-471f-9703-4d2595800e66",
403
+ "type": "TERMINAL",
404
+ "data": {
405
+ "label": "Final Output",
406
+ "name": "log",
407
+ "target_handle_id": "6b7d7f2c-5cc8-4005-9e66-cdb2c97b1998",
408
+ "output_id": "4c136180-050b-4422-a7a4-2a1c6729042c",
409
+ "output_type": "STRING",
410
+ "node_input_id": "76d49710-1ed0-4105-a1d7-9190c0408558",
411
+ },
412
+ "inputs": [
413
+ {
414
+ "id": "76d49710-1ed0-4105-a1d7-9190c0408558",
415
+ "key": "node_input",
416
+ "value": {
417
+ "rules": [
418
+ {
419
+ "type": "NODE_OUTPUT",
420
+ "data": {
421
+ "node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
422
+ "output_id": "7cac05e3-b7c3-475e-8df8-422b496c3398",
423
+ },
424
+ }
425
+ ],
426
+ "combinator": "OR",
427
+ },
428
+ }
429
+ ],
430
+ "display_data": {"position": {"x": 0.0, "y": 0.0}},
431
+ "definition": {
432
+ "name": "FinalOutputNode",
433
+ "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
434
+ "bases": [
435
+ {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}
436
+ ],
437
+ },
438
+ },
439
+ ],
440
+ workflow_raw_data["nodes"][2:],
441
+ ignore_order=True,
442
+ )
443
+
444
+ # AND each edge should be serialized correctly
445
+ serialized_edges = workflow_raw_data["edges"]
446
+ assert not DeepDiff(
447
+ [
448
+ {
449
+ "id": "72f2a432-621f-4a3a-8b41-17a5168cba69",
450
+ "source_node_id": "22555158-d8ba-41b4-a6fc-87c3b25bd073",
451
+ "source_handle_id": "e82390bb-c68c-48c1-9f87-7fbfff494c45",
452
+ "target_node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
453
+ "target_handle_id": "e02a2701-22c0-4533-8b00-175998e7350a",
454
+ "type": "DEFAULT",
455
+ },
456
+ {
457
+ "id": "2ac757e4-87c3-402c-928f-a3845df10c9f",
458
+ "source_node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
459
+ "source_handle_id": "832f81ec-427b-42a8-825c-e62c43c1f961",
460
+ "target_node_id": "eccf97c7-e766-471f-9703-4d2595800e66",
461
+ "target_handle_id": "6b7d7f2c-5cc8-4005-9e66-cdb2c97b1998",
462
+ "type": "DEFAULT",
463
+ },
464
+ {
465
+ "id": "fcc6353a-265c-4a65-9e70-4eb92a04e4e1",
466
+ "source_node_id": "c07155b3-7d99-4d2d-9b29-b5298013aa46",
467
+ "source_handle_id": "832f81ec-427b-42a8-825c-e62c43c1f961",
468
+ "target_node_id": "52f285fe-1f52-4920-b01b-499762b95220",
469
+ "target_handle_id": "de8f2cc2-8c32-4782-87d5-4eb5afcd42e3",
470
+ "type": "DEFAULT",
471
+ },
472
+ ],
473
+ serialized_edges,
474
+ ignore_order=True,
475
+ )
476
+
477
+ # AND the display data should be what we expect
478
+ display_data = workflow_raw_data["display_data"]
479
+ assert display_data == {
480
+ "viewport": {
481
+ "x": 0.0,
482
+ "y": 0.0,
483
+ "zoom": 1.0,
484
+ }
485
+ }
486
+
487
+ # AND the definition should be what we expect
488
+ definition = workflow_raw_data["definition"]
489
+ assert definition == {
490
+ "name": "SimpleCodeExecutionWithCodeWorkflow",
491
+ "module": ["tests", "workflows", "basic_code_execution_node", "workflow_with_code"],
492
+ }
493
+
494
+
311
495
  def test_serialize_workflow__try_wrapped():
312
496
  # GIVEN a Workflow with a code execution node
313
497
  # WHEN we serialize it
@@ -14,6 +14,7 @@ from vellum.workflows.ports import Port
14
14
  from vellum.workflows.references import OutputReference, WorkflowInputReference
15
15
  from vellum.workflows.types.core import JsonObject
16
16
  from vellum.workflows.types.generics import WorkflowType
17
+ from vellum.workflows.utils.uuids import uuid4_from_hash
17
18
  from vellum_ee.workflows.display.base import (
18
19
  EdgeDisplayOverridesType,
19
20
  EdgeDisplayType,
@@ -29,7 +30,6 @@ from vellum_ee.workflows.display.base import (
29
30
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
30
31
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
31
32
  from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
32
- from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
33
33
 
34
34
  logger = logging.getLogger(__name__)
35
35
 
@@ -154,7 +154,7 @@ class BaseWorkflowDisplay(
154
154
 
155
155
  def _get_node_display(self, node: Type[BaseNode]) -> NodeDisplayType:
156
156
  node_display_class = get_node_display_class(self.node_display_base_class, node)
157
- node_display = node_display_class(node)
157
+ node_display = node_display_class()
158
158
 
159
159
  if not isinstance(node_display, self.node_display_base_class):
160
160
  raise ValueError(f"{node.__name__} must be a subclass of {self.node_display_base_class.__name__}")
@@ -12,11 +12,10 @@ from vellum.workflows.references import WorkflowInputReference
12
12
  from vellum.workflows.references.output import OutputReference
13
13
  from vellum.workflows.types.core import JsonArray, JsonObject
14
14
  from vellum.workflows.types.generics import WorkflowType
15
- from vellum_ee.workflows.display.nodes import BaseMergeNodeDisplay
15
+ from vellum.workflows.utils.uuids import uuid4_from_hash
16
16
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
17
17
  from vellum_ee.workflows.display.nodes.types import PortDisplay
18
18
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
19
- from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
20
19
  from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type, primitive_to_vellum_value
21
20
  from vellum_ee.workflows.display.vellum import (
22
21
  EdgeVellumDisplay,
@@ -369,10 +368,7 @@ class VellumWorkflowDisplay(
369
368
  target_node_id = target_node_display.node_id
370
369
 
371
370
  target_handle_id: UUID
372
- if isinstance(target_node_display, BaseMergeNodeDisplay):
373
- target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
374
- else:
375
- target_handle_id = target_node_display.get_target_handle_id()
371
+ target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
376
372
 
377
373
  return self._generate_edge_display_from_source(
378
374
  source_node_id, source_handle_id, target_node_id, target_handle_id, overrides
File without changes