vellum-ai 0.12.16__py3-none-any.whl → 0.13.0__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (93) hide show
  1. vellum/__init__.py +10 -22
  2. vellum/client/__init__.py +8 -0
  3. vellum/client/core/client_wrapper.py +1 -1
  4. vellum/client/resources/__init__.py +4 -0
  5. vellum/client/resources/organizations/__init__.py +2 -0
  6. vellum/client/resources/organizations/client.py +116 -0
  7. vellum/client/resources/workspaces/__init__.py +2 -0
  8. vellum/client/resources/workspaces/client.py +114 -0
  9. vellum/client/types/__init__.py +6 -22
  10. vellum/client/types/new_member_join_behavior_enum.py +8 -0
  11. vellum/client/types/{function_call_variable_value.py → organization_read.py} +6 -4
  12. vellum/client/types/workflow_execution_workflow_result_event.py +0 -2
  13. vellum/client/types/workflow_result_event.py +0 -2
  14. vellum/client/types/workflow_result_event_output_data_array.py +4 -4
  15. vellum/client/types/{string_variable_value.py → workspace_read.py} +12 -5
  16. vellum/plugins/vellum_mypy.py +80 -11
  17. vellum/{types/json_variable_value.py → resources/organizations/__init__.py} +1 -1
  18. vellum/resources/organizations/client.py +3 -0
  19. vellum/{types/image_variable_value.py → resources/workspaces/__init__.py} +1 -1
  20. vellum/{types/array_variable_value.py → resources/workspaces/client.py} +1 -1
  21. vellum/types/{array_variable_value_item.py → new_member_join_behavior_enum.py} +1 -1
  22. vellum/types/{audio_variable_value.py → organization_read.py} +1 -1
  23. vellum/types/{error_variable_value.py → workspace_read.py} +1 -1
  24. vellum/utils/templating/render.py +3 -0
  25. vellum/workflows/nodes/bases/base.py +4 -0
  26. vellum/workflows/nodes/core/retry_node/node.py +24 -3
  27. vellum/workflows/nodes/core/retry_node/tests/test_node.py +40 -0
  28. vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py +20 -1
  29. vellum/workflows/nodes/displayable/api_node/node.py +3 -3
  30. vellum/workflows/workflows/base.py +0 -32
  31. {vellum_ai-0.12.16.dist-info → vellum_ai-0.13.0.dist-info}/METADATA +1 -1
  32. {vellum_ai-0.12.16.dist-info → vellum_ai-0.13.0.dist-info}/RECORD +79 -86
  33. vellum_cli/pull.py +7 -4
  34. vellum_cli/tests/conftest.py +4 -2
  35. vellum_cli/tests/test_push.py +13 -5
  36. vellum_ee/workflows/display/nodes/base_node_display.py +17 -10
  37. vellum_ee/workflows/display/nodes/vellum/api_node.py +1 -0
  38. vellum_ee/workflows/display/nodes/vellum/base_node.py +97 -2
  39. vellum_ee/workflows/display/nodes/vellum/code_execution_node.py +1 -0
  40. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +5 -62
  41. vellum_ee/workflows/display/nodes/vellum/error_node.py +1 -0
  42. vellum_ee/workflows/display/nodes/vellum/final_output_node.py +1 -0
  43. vellum_ee/workflows/display/nodes/vellum/guardrail_node.py +1 -0
  44. vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py +1 -0
  45. vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py +1 -0
  46. vellum_ee/workflows/display/nodes/vellum/map_node.py +1 -0
  47. vellum_ee/workflows/display/nodes/vellum/merge_node.py +1 -0
  48. vellum_ee/workflows/display/nodes/vellum/note_node.py +1 -0
  49. vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py +1 -0
  50. vellum_ee/workflows/display/nodes/vellum/search_node.py +1 -0
  51. vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py +1 -0
  52. vellum_ee/workflows/display/nodes/vellum/templating_node.py +1 -0
  53. vellum_ee/workflows/display/nodes/vellum/utils.py +63 -0
  54. vellum_ee/workflows/display/tests/test_vellum_workflow_display.py +2 -5
  55. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py +18 -2
  56. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py +355 -0
  57. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py +37 -22
  58. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py +12 -56
  59. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py +43 -93
  60. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py +31 -151
  61. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py +8 -26
  62. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py +4 -15
  63. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py +9 -44
  64. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +19 -101
  65. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py +19 -73
  66. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py +9 -44
  67. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py +9 -44
  68. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py +8 -6
  69. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py +11 -58
  70. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py +8 -11
  71. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py +7 -30
  72. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py +2 -11
  73. vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py +9 -44
  74. vellum_ee/workflows/display/vellum.py +2 -7
  75. vellum_ee/workflows/display/workflows/vellum_workflow_display.py +5 -9
  76. vellum_ee/workflows/server/virtual_file_loader.py +3 -3
  77. vellum/client/types/array_variable_value.py +0 -27
  78. vellum/client/types/array_variable_value_item.py +0 -29
  79. vellum/client/types/audio_variable_value.py +0 -25
  80. vellum/client/types/chat_history_variable_value.py +0 -21
  81. vellum/client/types/error_variable_value.py +0 -21
  82. vellum/client/types/image_variable_value.py +0 -25
  83. vellum/client/types/json_variable_value.py +0 -20
  84. vellum/client/types/number_variable_value.py +0 -20
  85. vellum/client/types/search_results_variable_value.py +0 -21
  86. vellum/types/chat_history_variable_value.py +0 -3
  87. vellum/types/function_call_variable_value.py +0 -3
  88. vellum/types/number_variable_value.py +0 -3
  89. vellum/types/search_results_variable_value.py +0 -3
  90. vellum/types/string_variable_value.py +0 -3
  91. {vellum_ai-0.12.16.dist-info → vellum_ai-0.13.0.dist-info}/LICENSE +0 -0
  92. {vellum_ai-0.12.16.dist-info → vellum_ai-0.13.0.dist-info}/WHEEL +0 -0
  93. {vellum_ai-0.12.16.dist-info → vellum_ai-0.13.0.dist-info}/entry_points.txt +0 -0
vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py CHANGED
@@ -60,17 +60,8 @@ def test_serialize_workflow():
60
60
  assert entrypoint_node == {
61
61
  "id": "54c5c7d0-ab86-4ae9-b0b8-ea9ca7b87c14",
62
62
  "type": "ENTRYPOINT",
63
- "definition": {
64
- "bases": [],
65
- "module": [
66
- "vellum",
67
- "workflows",
68
- "nodes",
69
- "bases",
70
- "base",
71
- ],
72
- "name": "BaseNode",
73
- },
63
+ "base": None,
64
+ "definition": None,
74
65
  "inputs": [],
75
66
  "data": {"label": "Entrypoint Node", "source_handle_id": "41840690-8d85-486e-a864-b0661ccf0f2e"},
76
67
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
@@ -117,26 +108,12 @@ def test_serialize_workflow():
117
108
  "release_tag": "LATEST",
118
109
  },
119
110
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
111
+ "base": {
112
+ "module": ["vellum", "workflows", "nodes", "displayable", "guardrail_node", "node"],
113
+ "name": "GuardrailNode",
114
+ },
120
115
  "definition": {
121
- "bases": [
122
- {
123
- "module": [
124
- "vellum",
125
- "workflows",
126
- "nodes",
127
- "displayable",
128
- "guardrail_node",
129
- "node",
130
- ],
131
- "name": "GuardrailNode",
132
- }
133
- ],
134
- "module": [
135
- "tests",
136
- "workflows",
137
- "basic_guardrail_node",
138
- "workflow",
139
- ],
116
+ "module": ["tests", "workflows", "basic_guardrail_node", "workflow"],
140
117
  "name": "ExampleGuardrailNode",
141
118
  },
142
119
  }
@@ -145,20 +122,7 @@ def test_serialize_workflow():
145
122
  assert final_output_node == {
146
123
  "id": "cbc7197e-67c9-4af5-b781-879c8fd3e4c9",
147
124
  "type": "TERMINAL",
148
- "definition": {
149
- "bases": [
150
- {
151
- "bases": [],
152
- "module": [
153
- "vellum",
154
- "workflows",
155
- "nodes",
156
- "bases",
157
- "base",
158
- ],
159
- "name": "BaseNode",
160
- },
161
- ],
125
+ "base": {
162
126
  "module": [
163
127
  "vellum",
164
128
  "workflows",
@@ -169,6 +133,7 @@ def test_serialize_workflow():
169
133
  ],
170
134
  "name": "FinalOutputNode",
171
135
  },
136
+ "definition": None,
172
137
  "data": {
173
138
  "label": "Final Output",
174
139
  "name": "score",
vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py CHANGED
@@ -78,17 +78,8 @@ def test_serialize_workflow():
78
78
  assert entrypoint_node == {
79
79
  "id": "6358dcfe-b162-4e19-99ca-401d1ada9bdc",
80
80
  "type": "ENTRYPOINT",
81
- "definition": {
82
- "bases": [],
83
- "module": [
84
- "vellum",
85
- "workflows",
86
- "nodes",
87
- "bases",
88
- "base",
89
- ],
90
- "name": "BaseNode",
91
- },
81
+ "base": None,
82
+ "definition": None,
92
83
  "inputs": [],
93
84
  "data": {
94
85
  "label": "Entrypoint Node",
@@ -136,32 +127,21 @@ def test_serialize_workflow():
136
127
  "source_handle_id": "9914a6a0-9a99-430d-8ddd-f7c13847fe1a",
137
128
  },
138
129
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
139
- "definition": {
140
- "name": "BaseNode",
141
- "module": [
142
- "vellum",
143
- "workflows",
144
- "nodes",
145
- "bases",
146
- "base",
147
- ],
148
- "bases": [],
149
- },
130
+ "base": None,
131
+ "definition": None,
150
132
  },
151
133
  {
152
134
  "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
153
135
  "label": "StartNode",
154
136
  "type": "GENERIC",
155
137
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
138
+ "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
156
139
  "definition": {
157
140
  "name": "StartNode",
158
141
  "module": ["tests", "workflows", "basic_inline_subworkflow", "workflow"],
159
- "bases": [
160
- {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}
161
- ],
162
142
  },
163
143
  "trigger": {"id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135", "merge_behavior": "AWAIT_ANY"},
164
- "ports": [],
144
+ "ports": [{"id": "4a62dea3-6a4a-4390-88d8-4af5d2c81474", "type": "DEFAULT"}],
165
145
  "adornments": None,
166
146
  "attributes": [],
167
147
  },
@@ -195,7 +175,7 @@ def test_serialize_workflow():
195
175
  }
196
176
  ],
197
177
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
198
- "definition": {
178
+ "base": {
199
179
  "name": "FinalOutputNode",
200
180
  "module": [
201
181
  "vellum",
@@ -205,20 +185,8 @@ def test_serialize_workflow():
205
185
  "final_output_node",
206
186
  "node",
207
187
  ],
208
- "bases": [
209
- {
210
- "name": "BaseNode",
211
- "module": [
212
- "vellum",
213
- "workflows",
214
- "nodes",
215
- "bases",
216
- "base",
217
- ],
218
- "bases": [],
219
- }
220
- ],
221
188
  },
189
+ "definition": None,
222
190
  },
223
191
  {
224
192
  "id": "570f4d12-69ff-49f1-ba98-ade6283dd7c2",
@@ -250,7 +218,7 @@ def test_serialize_workflow():
250
218
  }
251
219
  ],
252
220
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
253
- "definition": {
221
+ "base": {
254
222
  "name": "FinalOutputNode",
255
223
  "module": [
256
224
  "vellum",
@@ -260,20 +228,8 @@ def test_serialize_workflow():
260
228
  "final_output_node",
261
229
  "node",
262
230
  ],
263
- "bases": [
264
- {
265
- "name": "BaseNode",
266
- "module": [
267
- "vellum",
268
- "workflows",
269
- "nodes",
270
- "bases",
271
- "base",
272
- ],
273
- "bases": [],
274
- }
275
- ],
276
231
  },
232
+ "definition": None,
277
233
  },
278
234
  ],
279
235
  "edges": [
@@ -320,27 +276,13 @@ def test_serialize_workflow():
320
276
  ],
321
277
  },
322
278
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
279
+ "base": {
280
+ "name": "InlineSubworkflowNode",
281
+ "module": ["vellum", "workflows", "nodes", "core", "inline_subworkflow_node", "node"],
282
+ },
323
283
  "definition": {
324
284
  "name": "ExampleInlineSubworkflowNode",
325
- "module": [
326
- "tests",
327
- "workflows",
328
- "basic_inline_subworkflow",
329
- "workflow",
330
- ],
331
- "bases": [
332
- {
333
- "name": "InlineSubworkflowNode",
334
- "module": [
335
- "vellum",
336
- "workflows",
337
- "nodes",
338
- "core",
339
- "inline_subworkflow_node",
340
- "node",
341
- ],
342
- }
343
- ],
285
+ "module": ["tests", "workflows", "basic_inline_subworkflow", "workflow"],
344
286
  },
345
287
  },
346
288
  subworkflow_node,
@@ -358,20 +300,7 @@ def test_serialize_workflow():
358
300
  {
359
301
  "id": "31b74695-3f1c-47cf-8be8-a4d86cc589e8",
360
302
  "type": "TERMINAL",
361
- "definition": {
362
- "bases": [
363
- {
364
- "bases": [],
365
- "module": [
366
- "vellum",
367
- "workflows",
368
- "nodes",
369
- "bases",
370
- "base",
371
- ],
372
- "name": "BaseNode",
373
- },
374
- ],
303
+ "base": {
375
304
  "module": [
376
305
  "vellum",
377
306
  "workflows",
@@ -382,6 +311,7 @@ def test_serialize_workflow():
382
311
  ],
383
312
  "name": "FinalOutputNode",
384
313
  },
314
+ "definition": None,
385
315
  "data": {
386
316
  "label": "Final Output",
387
317
  "name": "reasoning",
@@ -421,20 +351,7 @@ def test_serialize_workflow():
421
351
  {
422
352
  "id": "0779b232-82ab-4dbe-a340-6a85e6ab3368",
423
353
  "type": "TERMINAL",
424
- "definition": {
425
- "bases": [
426
- {
427
- "bases": [],
428
- "module": [
429
- "vellum",
430
- "workflows",
431
- "nodes",
432
- "bases",
433
- "base",
434
- ],
435
- "name": "BaseNode",
436
- },
437
- ],
354
+ "base": {
438
355
  "module": [
439
356
  "vellum",
440
357
  "workflows",
@@ -445,6 +362,7 @@ def test_serialize_workflow():
445
362
  ],
446
363
  "name": "FinalOutputNode",
447
364
  },
365
+ "definition": None,
448
366
  "data": {
449
367
  "label": "Final Output",
450
368
  "name": "temperature",
vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py CHANGED
@@ -68,17 +68,9 @@ def test_serialize_workflow():
68
68
  "label": "Entrypoint Node",
69
69
  "source_handle_id": "844d992e-60ab-4af2-a8ff-52cd858386f7",
70
70
  },
71
- "definition": {
72
- "bases": [],
73
- "module": [
74
- "vellum",
75
- "workflows",
76
- "nodes",
77
- "bases",
78
- "base",
79
- ],
80
- "name": "BaseNode",
81
- },
71
+ "base": None,
72
+ "base": None,
73
+ "definition": None,
82
74
  "display_data": {
83
75
  "position": {"x": 0.0, "y": 0.0},
84
76
  },
@@ -121,32 +113,24 @@ def test_serialize_workflow():
121
113
  "source_handle_id": "520d3616-8369-4e79-9da5-3febae299c2a",
122
114
  },
123
115
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
124
- "definition": {
125
- "name": "BaseNode",
126
- "module": [
127
- "vellum",
128
- "workflows",
129
- "nodes",
130
- "bases",
131
- "base",
132
- ],
133
- "bases": [],
134
- },
116
+ "base": None,
117
+ "definition": None,
135
118
  },
136
119
  {
137
120
  "id": "baf6d316-dc75-41e8-96c0-015aede96309",
138
121
  "label": "Iteration",
139
122
  "type": "GENERIC",
140
123
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
124
+ "base": {
125
+ "name": "BaseNode",
126
+ "module": ["vellum", "workflows", "nodes", "bases", "base"],
127
+ },
141
128
  "definition": {
142
129
  "name": "Iteration",
143
130
  "module": ["tests", "workflows", "basic_map_node", "workflow"],
144
- "bases": [
145
- {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]}
146
- ],
147
131
  },
148
132
  "trigger": {"id": "01324747-9bc0-4ecd-a8ab-40dca5a94e2e", "merge_behavior": "AWAIT_ANY"},
149
- "ports": [],
133
+ "ports": [{"id": "36791877-95b0-4390-a794-10accd8f548e", "type": "DEFAULT"}],
150
134
  "adornments": None,
151
135
  "attributes": [],
152
136
  },
@@ -180,7 +164,7 @@ def test_serialize_workflow():
180
164
  }
181
165
  ],
182
166
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
183
- "definition": {
167
+ "base": {
184
168
  "name": "FinalOutputNode",
185
169
  "module": [
186
170
  "vellum",
@@ -190,20 +174,8 @@ def test_serialize_workflow():
190
174
  "final_output_node",
191
175
  "node",
192
176
  ],
193
- "bases": [
194
- {
195
- "name": "BaseNode",
196
- "module": [
197
- "vellum",
198
- "workflows",
199
- "nodes",
200
- "bases",
201
- "base",
202
- ],
203
- "bases": [],
204
- }
205
- ],
206
177
  },
178
+ "definition": None,
207
179
  },
208
180
  ],
209
181
  "edges": [
@@ -268,27 +240,13 @@ def test_serialize_workflow():
268
240
  "index_input_id": "17e7ca49-668f-450d-a792-e1f97d13db67",
269
241
  },
270
242
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
243
+ "base": {
244
+ "name": "MapNode",
245
+ "module": ["vellum", "workflows", "nodes", "core", "map_node", "node"],
246
+ },
271
247
  "definition": {
272
248
  "name": "MapFruitsNode",
273
- "module": [
274
- "tests",
275
- "workflows",
276
- "basic_map_node",
277
- "workflow",
278
- ],
279
- "bases": [
280
- {
281
- "name": "MapNode",
282
- "module": [
283
- "vellum",
284
- "workflows",
285
- "nodes",
286
- "core",
287
- "map_node",
288
- "node",
289
- ],
290
- }
291
- ],
249
+ "module": ["tests", "workflows", "basic_map_node", "workflow"],
292
250
  },
293
251
  },
294
252
  map_node,
@@ -326,20 +284,7 @@ def test_serialize_workflow():
326
284
  }
327
285
  ],
328
286
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
329
- "definition": {
330
- "bases": [
331
- {
332
- "bases": [],
333
- "module": [
334
- "vellum",
335
- "workflows",
336
- "nodes",
337
- "bases",
338
- "base",
339
- ],
340
- "name": "BaseNode",
341
- },
342
- ],
287
+ "base": {
343
288
  "module": [
344
289
  "vellum",
345
290
  "workflows",
@@ -350,6 +295,7 @@ def test_serialize_workflow():
350
295
  ],
351
296
  "name": "FinalOutputNode",
352
297
  },
298
+ "definition": None,
353
299
  },
354
300
  workflow_raw_data["nodes"][2],
355
301
  # TODO: Fix output ID not referencing map node workflow output
vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py CHANGED
@@ -47,17 +47,8 @@ def test_serialize_workflow__await_all():
47
47
  assert entrypoint_node == {
48
48
  "id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
49
49
  "type": "ENTRYPOINT",
50
- "definition": {
51
- "bases": [],
52
- "module": [
53
- "vellum",
54
- "workflows",
55
- "nodes",
56
- "bases",
57
- "base",
58
- ],
59
- "name": "BaseNode",
60
- },
50
+ "base": None,
51
+ "definition": None,
61
52
  "inputs": [],
62
53
  "data": {
63
54
  "label": "Entrypoint Node",
@@ -87,26 +78,12 @@ def test_serialize_workflow__await_all():
87
78
  "source_handle_id": "3bbc469f-0fb0-4b3d-a28b-746fefec2818",
88
79
  },
89
80
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
81
+ "base": {
82
+ "module": ["vellum", "workflows", "nodes", "displayable", "merge_node", "node"],
83
+ "name": "MergeNode",
84
+ },
90
85
  "definition": {
91
- "bases": [
92
- {
93
- "module": [
94
- "vellum",
95
- "workflows",
96
- "nodes",
97
- "displayable",
98
- "merge_node",
99
- "node",
100
- ],
101
- "name": "MergeNode",
102
- }
103
- ],
104
- "module": [
105
- "tests",
106
- "workflows",
107
- "basic_merge_node",
108
- "await_all_workflow",
109
- ],
86
+ "module": ["tests", "workflows", "basic_merge_node", "await_all_workflow"],
110
87
  "name": "AwaitAllMergeNode",
111
88
  },
112
89
  },
@@ -126,20 +103,7 @@ def test_serialize_workflow__await_all():
126
103
  "output_type": "STRING",
127
104
  "node_input_id": "7f950be4-2fab-44e0-87a3-b1631aadd0e3",
128
105
  },
129
- "definition": {
130
- "bases": [
131
- {
132
- "bases": [],
133
- "module": [
134
- "vellum",
135
- "workflows",
136
- "nodes",
137
- "bases",
138
- "base",
139
- ],
140
- "name": "BaseNode",
141
- },
142
- ],
106
+ "base": {
143
107
  "module": [
144
108
  "vellum",
145
109
  "workflows",
@@ -150,6 +114,7 @@ def test_serialize_workflow__await_all():
150
114
  ],
151
115
  "name": "FinalOutputNode",
152
116
  },
117
+ "definition": None,
153
118
  "inputs": [
154
119
  {
155
120
  "id": "7f950be4-2fab-44e0-87a3-b1631aadd0e3",
vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py CHANGED
@@ -86,17 +86,8 @@ def test_serialize_workflow(vellum_client):
86
86
  "id": "d680afbd-de64-4cf6-aa50-912686c48c64",
87
87
  "type": "ENTRYPOINT",
88
88
  "inputs": [],
89
- "definition": {
90
- "bases": [],
91
- "module": [
92
- "vellum",
93
- "workflows",
94
- "nodes",
95
- "bases",
96
- "base",
97
- ],
98
- "name": "BaseNode",
99
- },
89
+ "base": None,
90
+ "definition": None,
100
91
  "data": {
101
92
  "label": "Entrypoint Node",
102
93
  "source_handle_id": "7065a943-1cab-4afd-9690-e678c5b74a2f",
@@ -159,26 +150,12 @@ def test_serialize_workflow(vellum_client):
159
150
  "y": 0.0,
160
151
  },
161
152
  },
153
+ "base": {
154
+ "module": ["vellum", "workflows", "nodes", "displayable", "prompt_deployment_node", "node"],
155
+ "name": "PromptDeploymentNode",
156
+ },
162
157
  "definition": {
163
- "bases": [
164
- {
165
- "module": [
166
- "vellum",
167
- "workflows",
168
- "nodes",
169
- "displayable",
170
- "prompt_deployment_node",
171
- "node",
172
- ],
173
- "name": "PromptDeploymentNode",
174
- }
175
- ],
176
- "module": [
177
- "tests",
178
- "workflows",
179
- "basic_text_prompt_deployment",
180
- "workflow",
181
- ],
158
+ "module": ["tests", "workflows", "basic_text_prompt_deployment", "workflow"],
182
159
  "name": "ExamplePromptDeploymentNode",
183
160
  },
184
161
  }
@@ -187,20 +164,7 @@ def test_serialize_workflow(vellum_client):
187
164
  assert final_output_node == {
188
165
  "id": "64ff72c7-8ffc-4e1f-b7a7-e7cd0697f576",
189
166
  "type": "TERMINAL",
190
- "definition": {
191
- "bases": [
192
- {
193
- "bases": [],
194
- "module": [
195
- "vellum",
196
- "workflows",
197
- "nodes",
198
- "bases",
199
- "base",
200
- ],
201
- "name": "BaseNode",
202
- },
203
- ],
167
+ "base": {
204
168
  "module": [
205
169
  "vellum",
206
170
  "workflows",
@@ -211,6 +175,7 @@ def test_serialize_workflow(vellum_client):
211
175
  ],
212
176
  "name": "FinalOutputNode",
213
177
  },
178
+ "definition": None,
214
179
  "inputs": [
215
180
  {
216
181
  "id": "5f0fdd0f-63fe-4437-bc1b-cd6a84bb84c5",
vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py CHANGED
@@ -50,7 +50,8 @@ def test_serialize_workflow():
50
50
  "inputs": [],
51
51
  "data": {"label": "Entrypoint Node", "source_handle_id": "df80b4aa-2ba1-49a2-8375-fb1f78eee31f"},
52
52
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
53
- "definition": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []},
53
+ "base": None,
54
+ "definition": None,
54
55
  }
55
56
 
56
57
  search_node = workflow_raw_data["nodes"][1]
@@ -223,12 +224,13 @@ def test_serialize_workflow():
223
224
  "metadata_filters_node_input_id": "855d3f57-e633-467e-a348-a394360247df",
224
225
  },
225
226
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
227
+ "base": {
228
+ "name": "SearchNode",
229
+ "module": ["vellum", "workflows", "nodes", "displayable", "search_node", "node"],
230
+ },
226
231
  "definition": {
227
232
  "name": "SimpleSearchNode",
228
233
  "module": ["tests", "workflows", "basic_search_node", "workflow"],
229
- "bases": [
230
- {"name": "SearchNode", "module": ["vellum", "workflows", "nodes", "displayable", "search_node", "node"]}
231
- ],
232
234
  },
233
235
  }
234
236
 
@@ -263,11 +265,11 @@ def test_serialize_workflow():
263
265
  }
264
266
  ],
265
267
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
266
- "definition": {
268
+ "base": {
267
269
  "name": "FinalOutputNode",
268
270
  "module": ["vellum", "workflows", "nodes", "displayable", "final_output_node", "node"],
269
- "bases": [{"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"], "bases": []}],
270
271
  },
272
+ "definition": None,
271
273
  }
272
274
 
273
275
  # AND each edge should be serialized correctly