vellum-ai 0.14.68__py3-none-any.whl → 0.14.70__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 (54) hide show
  1. vellum/__init__.py +4 -0
  2. vellum/client/core/client_wrapper.py +1 -1
  3. vellum/client/types/__init__.py +4 -0
  4. vellum/client/types/fast_embed_vectorizer_baai_bge_small_en_v_15.py +23 -0
  5. vellum/client/types/fast_embed_vectorizer_baai_bge_small_en_v_15_request.py +23 -0
  6. vellum/client/types/folder_entity_document_index_data.py +2 -0
  7. vellum/client/types/indexing_config_vectorizer.py +2 -0
  8. vellum/client/types/indexing_config_vectorizer_request.py +2 -0
  9. vellum/types/fast_embed_vectorizer_baai_bge_small_en_v_15.py +3 -0
  10. vellum/types/fast_embed_vectorizer_baai_bge_small_en_v_15_request.py +3 -0
  11. vellum/workflows/environment/__init__.py +2 -1
  12. vellum/workflows/environment/environment.py +5 -1
  13. vellum/workflows/nodes/displayable/bases/search_node.py +15 -3
  14. vellum/workflows/nodes/displayable/tests/test_search_node_error_handling.py +215 -0
  15. vellum/workflows/nodes/experimental/tool_calling_node/tests/test_node.py +77 -1
  16. vellum/workflows/nodes/experimental/tool_calling_node/utils.py +2 -2
  17. vellum/workflows/references/environment_variable.py +2 -3
  18. {vellum_ai-0.14.68.dist-info → vellum_ai-0.14.70.dist-info}/METADATA +1 -1
  19. {vellum_ai-0.14.68.dist-info → vellum_ai-0.14.70.dist-info}/RECORD +54 -45
  20. vellum_cli/__init__.py +5 -2
  21. vellum_cli/image_push.py +24 -1
  22. vellum_cli/tests/test_image_push.py +103 -12
  23. vellum_ee/workflows/display/nodes/base_node_display.py +1 -1
  24. vellum_ee/workflows/display/nodes/utils.py +2 -2
  25. vellum_ee/workflows/display/nodes/vellum/api_node.py +2 -2
  26. vellum_ee/workflows/display/nodes/vellum/code_execution_node.py +1 -1
  27. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +1 -1
  28. vellum_ee/workflows/display/nodes/vellum/error_node.py +1 -1
  29. vellum_ee/workflows/display/nodes/vellum/final_output_node.py +2 -2
  30. vellum_ee/workflows/display/nodes/vellum/guardrail_node.py +1 -1
  31. vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py +4 -4
  32. vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py +9 -1
  33. vellum_ee/workflows/display/nodes/vellum/map_node.py +1 -1
  34. vellum_ee/workflows/display/nodes/vellum/merge_node.py +1 -1
  35. vellum_ee/workflows/display/nodes/vellum/note_node.py +1 -0
  36. vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py +1 -1
  37. vellum_ee/workflows/display/nodes/vellum/retry_node.py +1 -1
  38. vellum_ee/workflows/display/nodes/vellum/search_node.py +1 -1
  39. vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py +1 -1
  40. vellum_ee/workflows/display/nodes/vellum/templating_node.py +1 -1
  41. vellum_ee/workflows/display/nodes/vellum/tests/test_inline_subworkflow_node.py +88 -0
  42. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py +16 -0
  43. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_prompt_node_serialization.py +81 -0
  44. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +9 -1
  45. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_inline_workflow_serialization.py +59 -297
  46. vellum_ee/workflows/display/utils/auto_layout.py +130 -0
  47. vellum_ee/workflows/display/utils/expressions.py +7 -0
  48. vellum_ee/workflows/display/utils/tests/__init__.py +0 -0
  49. vellum_ee/workflows/display/utils/tests/test_auto_layout.py +56 -0
  50. vellum_ee/workflows/display/workflows/base_workflow_display.py +15 -10
  51. vellum_ee/workflows/display/workflows/tests/test_workflow_display.py +41 -0
  52. {vellum_ai-0.14.68.dist-info → vellum_ai-0.14.70.dist-info}/LICENSE +0 -0
  53. {vellum_ai-0.14.68.dist-info → vellum_ai-0.14.70.dist-info}/WHEEL +0 -0
  54. {vellum_ai-0.14.68.dist-info → vellum_ai-0.14.70.dist-info}/entry_points.txt +0 -0
@@ -1,4 +1,3 @@
1
- # type: ignore # subworkflow is causing mypy to hang indefinitely
2
1
  from deepdiff import DeepDiff
3
2
 
4
3
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -148,290 +147,16 @@ def test_serialize_workflow():
148
147
  "definition": None,
149
148
  },
150
149
  {
151
- "id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
152
- "type": "SUBWORKFLOW",
153
- "inputs": [
154
- {
155
- "id": "704c4640-bfda-44f0-8da3-e9cfc4f21cf2",
156
- "key": "metro",
157
- "value": {
158
- "rules": [
159
- {
160
- "type": "INPUT_VARIABLE",
161
- "data": {
162
- "input_variable_id": "fa73da37-34c3-47a9-be58-69cc6cdbfca5" # noqa: E501
163
- },
164
- }
165
- ],
166
- "combinator": "OR",
167
- },
168
- }
169
- ],
170
- "data": {
171
- "label": "Example Inline Subworkflow Node",
172
- "error_output_id": None,
173
- "source_handle_id": "cfd831bc-ee7f-44d0-8d76-0ba0cd0277dc",
174
- "target_handle_id": "859a75a6-1bd2-4350-9509-4af66245e8e4",
175
- "variant": "INLINE",
176
- "workflow_raw_data": {
177
- "nodes": [
178
- {
179
- "id": "afa49a0f-db35-4552-9217-5b8f237e84bc",
180
- "type": "ENTRYPOINT",
181
- "inputs": [],
182
- "data": {
183
- "label": "Entrypoint Node",
184
- "source_handle_id": "9914a6a0-9a99-430d-8ddd-f7c13847fe1a", # noqa: E501
185
- },
186
- "display_data": {"position": {"x": 0.0, "y": 0.0}},
187
- "base": None,
188
- "definition": None,
189
- },
190
- {
191
- "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
192
- "label": "StartNode",
193
- "type": "GENERIC",
194
- "display_data": {"position": {"x": 0.0, "y": 0.0}},
195
- "base": {
196
- "name": "BaseNode",
197
- "module": [
198
- "vellum",
199
- "workflows",
200
- "nodes",
201
- "bases",
202
- "base",
203
- ],
204
- },
205
- "definition": {
206
- "name": "StartNode",
207
- "module": [
208
- "tests",
209
- "workflows",
210
- "basic_tool_calling_node_inline_workflow",
211
- "workflow",
212
- ],
213
- },
214
- "trigger": {
215
- "id": "6492efcf-4437-4af1-9ad7-269795ccb27a",
216
- "merge_behavior": "AWAIT_ATTRIBUTES",
217
- },
218
- "ports": [
219
- {
220
- "id": "1e739e86-a285-4438-9725-a152c15a63e3",
221
- "name": "default",
222
- "type": "DEFAULT",
223
- }
224
- ],
225
- "adornments": None,
226
- "attributes": [
227
- {
228
- "id": "b0ac6b50-22a8-42ba-a707-1aa09a653205",
229
- "name": "metro",
230
- "value": {
231
- "type": "WORKFLOW_INPUT",
232
- "input_variable_id": "f2f5da15-026d-4905-bfe7-7d16bda20eed", # noqa: E501
233
- },
234
- },
235
- {
236
- "id": "c5f2d66c-5bb6-4d2a-8e4d-5356318cd3ba",
237
- "name": "date",
238
- "value": {
239
- "type": "WORKFLOW_INPUT",
240
- "input_variable_id": "aba1e6e0-dfa7-4c15-a4e6-aec6feebfaca", # noqa: E501
241
- },
242
- },
243
- ],
244
- "outputs": [
245
- {
246
- "id": "3f4c753e-f057-47bb-9748-7968283cc8aa",
247
- "name": "temperature",
248
- "type": "NUMBER",
249
- "value": None,
250
- },
251
- {
252
- "id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd",
253
- "name": "reasoning",
254
- "type": "STRING",
255
- "value": None,
256
- },
257
- ],
258
- },
259
- {
260
- "id": "a773c3a5-78cb-4250-8d29-7282e8a579d3",
261
- "type": "TERMINAL",
262
- "data": {
263
- "label": "Final Output",
264
- "name": "temperature",
265
- "target_handle_id": "804bb543-9cf4-457f-acf1-fb4b8b7d9259", # noqa: E501
266
- "output_id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622",
267
- "output_type": "NUMBER",
268
- "node_input_id": "712eaeec-9e1e-41bd-9217-9caec8b6ade7", # noqa: E501
269
- },
270
- "inputs": [
271
- {
272
- "id": "712eaeec-9e1e-41bd-9217-9caec8b6ade7",
273
- "key": "node_input",
274
- "value": {
275
- "rules": [
276
- {
277
- "type": "NODE_OUTPUT",
278
- "data": {
279
- "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
280
- "output_id": "3f4c753e-f057-47bb-9748-7968283cc8aa", # noqa: E501
281
- },
282
- }
283
- ],
284
- "combinator": "OR",
285
- },
286
- }
287
- ],
288
- "display_data": {"position": {"x": 0.0, "y": 0.0}},
289
- "base": {
290
- "name": "FinalOutputNode",
291
- "module": [
292
- "vellum",
293
- "workflows",
294
- "nodes",
295
- "displayable",
296
- "final_output_node",
297
- "node",
298
- ],
299
- },
300
- "definition": None,
301
- },
302
- {
303
- "id": "570f4d12-69ff-49f1-ba98-ade6283dd7c2",
304
- "type": "TERMINAL",
305
- "data": {
306
- "label": "Final Output",
307
- "name": "reasoning",
308
- "target_handle_id": "6d4c4a14-c388-4c7a-b223-eb39baf5c080", # noqa: E501
309
- "output_id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82",
310
- "output_type": "STRING",
311
- "node_input_id": "8fd4279a-4f13-4257-9577-1b55e964cdf1", # noqa: E501
312
- },
313
- "inputs": [
314
- {
315
- "id": "8fd4279a-4f13-4257-9577-1b55e964cdf1",
316
- "key": "node_input",
317
- "value": {
318
- "rules": [
319
- {
320
- "type": "NODE_OUTPUT",
321
- "data": {
322
- "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
323
- "output_id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd", # noqa: E501
324
- },
325
- }
326
- ],
327
- "combinator": "OR",
328
- },
329
- }
330
- ],
331
- "display_data": {"position": {"x": 0.0, "y": 0.0}},
332
- "base": {
333
- "name": "FinalOutputNode",
334
- "module": [
335
- "vellum",
336
- "workflows",
337
- "nodes",
338
- "displayable",
339
- "final_output_node",
340
- "node",
341
- ],
342
- },
343
- "definition": None,
344
- },
345
- ],
346
- "edges": [
347
- {
348
- "id": "fb2f58f0-9d49-4658-af78-afa9b94091a6",
349
- "source_node_id": "afa49a0f-db35-4552-9217-5b8f237e84bc", # noqa: E501
350
- "source_handle_id": "9914a6a0-9a99-430d-8ddd-f7c13847fe1a", # noqa: E501
351
- "target_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
352
- "target_handle_id": "6492efcf-4437-4af1-9ad7-269795ccb27a", # noqa: E501
353
- "type": "DEFAULT",
354
- },
355
- {
356
- "id": "6f16dfb8-d794-4be8-8860-6ea34f0b9e7c",
357
- "source_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
358
- "source_handle_id": "1e739e86-a285-4438-9725-a152c15a63e3", # noqa: E501
359
- "target_node_id": "a773c3a5-78cb-4250-8d29-7282e8a579d3", # noqa: E501
360
- "target_handle_id": "804bb543-9cf4-457f-acf1-fb4b8b7d9259", # noqa: E501
361
- "type": "DEFAULT",
362
- },
363
- {
364
- "id": "63b77ff0-5282-46ce-8da9-37ced05ac61c",
365
- "source_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
366
- "source_handle_id": "1e739e86-a285-4438-9725-a152c15a63e3", # noqa: E501
367
- "target_node_id": "570f4d12-69ff-49f1-ba98-ade6283dd7c2", # noqa: E501
368
- "target_handle_id": "6d4c4a14-c388-4c7a-b223-eb39baf5c080", # noqa: E501
369
- "type": "DEFAULT",
370
- },
371
- ],
372
- "display_data": {"viewport": {"x": 0.0, "y": 0.0, "zoom": 1.0}},
373
- "definition": {
374
- "name": "NestedWorkflow",
375
- "module": [
376
- "tests",
377
- "workflows",
378
- "basic_tool_calling_node_inline_workflow",
379
- "workflow",
380
- ],
381
- },
382
- "output_values": [
383
- {
384
- "output_variable_id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622", # noqa: E501
385
- "value": {
386
- "type": "NODE_OUTPUT",
387
- "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
388
- "node_output_id": "3f4c753e-f057-47bb-9748-7968283cc8aa", # noqa: E501
389
- },
390
- },
391
- {
392
- "output_variable_id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82", # noqa: E501
393
- "value": {
394
- "type": "NODE_OUTPUT",
395
- "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
396
- "node_output_id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd", # noqa: E501
397
- },
398
- },
399
- ],
400
- },
401
- "input_variables": [
402
- {
403
- "id": "704c4640-bfda-44f0-8da3-e9cfc4f21cf2",
404
- "key": "metro",
405
- "type": "STRING",
406
- }
407
- ],
408
- "output_variables": [
409
- {
410
- "id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622",
411
- "key": "temperature",
412
- "type": "NUMBER",
413
- },
414
- {
415
- "id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82",
416
- "key": "reasoning",
417
- "type": "STRING",
418
- },
419
- ],
420
- },
150
+ "id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
151
+ "label": "StartNode",
152
+ "type": "GENERIC",
421
153
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
422
154
  "base": {
423
- "name": "InlineSubworkflowNode",
424
- "module": [
425
- "vellum",
426
- "workflows",
427
- "nodes",
428
- "core",
429
- "inline_subworkflow_node",
430
- "node",
431
- ],
155
+ "name": "BaseNode",
156
+ "module": ["vellum", "workflows", "nodes", "bases", "base"],
432
157
  },
433
158
  "definition": {
434
- "name": "ExampleInlineSubworkflowNode",
159
+ "name": "StartNode",
435
160
  "module": [
436
161
  "tests",
437
162
  "workflows",
@@ -439,13 +164,50 @@ def test_serialize_workflow():
439
164
  "workflow",
440
165
  ],
441
166
  },
167
+ "trigger": {
168
+ "id": "6492efcf-4437-4af1-9ad7-269795ccb27a",
169
+ "merge_behavior": "AWAIT_ATTRIBUTES",
170
+ },
442
171
  "ports": [
443
172
  {
444
- "id": "cfd831bc-ee7f-44d0-8d76-0ba0cd0277dc",
173
+ "id": "1e739e86-a285-4438-9725-a152c15a63e3",
445
174
  "name": "default",
446
175
  "type": "DEFAULT",
447
176
  }
448
177
  ],
178
+ "adornments": None,
179
+ "attributes": [
180
+ {
181
+ "id": "60ad78cd-fc78-4e08-926d-5a095b34d4f5",
182
+ "name": "city",
183
+ "value": {
184
+ "type": "WORKFLOW_INPUT",
185
+ "input_variable_id": "fa73da37-34c3-47a9-be58-69cc6cdbfca5",
186
+ },
187
+ },
188
+ {
189
+ "id": "c5f2d66c-5bb6-4d2a-8e4d-5356318cd3ba",
190
+ "name": "date",
191
+ "value": {
192
+ "type": "WORKFLOW_INPUT",
193
+ "input_variable_id": "aba1e6e0-dfa7-4c15-a4e6-aec6feebfaca",
194
+ },
195
+ },
196
+ ],
197
+ "outputs": [
198
+ {
199
+ "id": "3f4c753e-f057-47bb-9748-7968283cc8aa",
200
+ "name": "temperature",
201
+ "type": "NUMBER",
202
+ "value": None,
203
+ },
204
+ {
205
+ "id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd",
206
+ "name": "reasoning",
207
+ "type": "STRING",
208
+ "value": None,
209
+ },
210
+ ],
449
211
  },
450
212
  {
451
213
  "id": "0779b232-82ab-4dbe-a340-6a85e6ab3368",
@@ -467,8 +229,8 @@ def test_serialize_workflow():
467
229
  {
468
230
  "type": "NODE_OUTPUT",
469
231
  "data": {
470
- "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239", # noqa: E501
471
- "output_id": "86dd0202-c141-48a3-8382-2da60372e77c", # noqa: E501
232
+ "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
233
+ "output_id": "3f4c753e-f057-47bb-9748-7968283cc8aa", # noqa: E501
472
234
  },
473
235
  }
474
236
  ],
@@ -510,8 +272,8 @@ def test_serialize_workflow():
510
272
  {
511
273
  "type": "NODE_OUTPUT",
512
274
  "data": {
513
- "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239", # noqa: E501
514
- "output_id": "0a7192da-5576-4933-bba4-de8adf5d7996", # noqa: E501
275
+ "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc", # noqa: E501
276
+ "output_id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd", # noqa: E501
515
277
  },
516
278
  }
517
279
  ],
@@ -536,25 +298,25 @@ def test_serialize_workflow():
536
298
  ],
537
299
  "edges": [
538
300
  {
539
- "id": "71dd3569-ccf8-4352-ad42-3594be3a6c16",
301
+ "id": "a37781d1-f7a5-4386-a67d-0c3d5c929602",
540
302
  "source_node_id": "6358dcfe-b162-4e19-99ca-401d1ada9bdc",
541
303
  "source_handle_id": "c344fdee-282b-40c9-8c97-6dd08830948c",
542
- "target_node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
543
- "target_handle_id": "859a75a6-1bd2-4350-9509-4af66245e8e4",
304
+ "target_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
305
+ "target_handle_id": "6492efcf-4437-4af1-9ad7-269795ccb27a",
544
306
  "type": "DEFAULT",
545
307
  },
546
308
  {
547
309
  "id": "3c5d8990-48f5-42e1-893e-bc8308d2110a",
548
- "source_node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
549
- "source_handle_id": "cfd831bc-ee7f-44d0-8d76-0ba0cd0277dc",
310
+ "source_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
311
+ "source_handle_id": "1e739e86-a285-4438-9725-a152c15a63e3",
550
312
  "target_node_id": "0779b232-82ab-4dbe-a340-6a85e6ab3368",
551
313
  "target_handle_id": "9e077063-c394-4c7b-b0c6-e6686df67984",
552
314
  "type": "DEFAULT",
553
315
  },
554
316
  {
555
317
  "id": "de0b8090-a26e-4e09-9173-9f7400a5be4c",
556
- "source_node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
557
- "source_handle_id": "cfd831bc-ee7f-44d0-8d76-0ba0cd0277dc",
318
+ "source_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
319
+ "source_handle_id": "1e739e86-a285-4438-9725-a152c15a63e3",
558
320
  "target_node_id": "31b74695-3f1c-47cf-8be8-a4d86cc589e8",
559
321
  "target_handle_id": "8b525943-6c27-414b-a329-e29c0b217f72",
560
322
  "type": "DEFAULT",
@@ -575,16 +337,16 @@ def test_serialize_workflow():
575
337
  "output_variable_id": "99afb757-2782-465d-ab55-80ccf50552b9",
576
338
  "value": {
577
339
  "type": "NODE_OUTPUT",
578
- "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
579
- "node_output_id": "86dd0202-c141-48a3-8382-2da60372e77c",
340
+ "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
341
+ "node_output_id": "3f4c753e-f057-47bb-9748-7968283cc8aa",
580
342
  },
581
343
  },
582
344
  {
583
345
  "output_variable_id": "7444a019-081a-4e10-a528-3249299159f7",
584
346
  "value": {
585
347
  "type": "NODE_OUTPUT",
586
- "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
587
- "node_output_id": "0a7192da-5576-4933-bba4-de8adf5d7996",
348
+ "node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
349
+ "node_output_id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd",
588
350
  },
589
351
  },
590
352
  ],
@@ -0,0 +1,130 @@
1
+ from collections import defaultdict
2
+ from typing import Dict, List, Set, Tuple
3
+
4
+ from vellum_ee.workflows.display.base import EdgeDisplay
5
+ from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
6
+
7
+
8
+ def auto_layout_nodes(
9
+ nodes: List[Tuple[str, NodeDisplayData]],
10
+ edges: List[Tuple[str, str, EdgeDisplay]],
11
+ node_spacing: float = 150.0,
12
+ layer_spacing: float = 200.0,
13
+ ) -> List[Tuple[str, NodeDisplayData]]:
14
+ """
15
+ Auto-layout nodes in a hierarchical left-to-right arrangement.
16
+
17
+ Args:
18
+ nodes: List of (node_id, NodeDisplayData) tuples
19
+ edges: List of (source_node_id, target_node_id, EdgeDisplay) tuples
20
+ node_spacing: Vertical spacing between nodes in the same layer
21
+ layer_spacing: Horizontal spacing between layers
22
+
23
+ Returns:
24
+ List of (node_id, NodeDisplayData) tuples with updated positions
25
+ """
26
+ if not nodes:
27
+ return []
28
+
29
+ node_dict = {node_id: data for node_id, data in nodes}
30
+
31
+ graph: Dict[str, List[str]] = defaultdict(list)
32
+ in_degree: Dict[str, int] = defaultdict(int)
33
+ all_nodes = set(node_dict.keys())
34
+
35
+ for source, target, _ in edges:
36
+ if source in all_nodes and target in all_nodes:
37
+ graph[source].append(target)
38
+ in_degree[target] += 1
39
+
40
+ for node_id in all_nodes:
41
+ if node_id not in in_degree:
42
+ in_degree[node_id] = 0
43
+
44
+ layers = _topological_sort_layers(graph, in_degree, all_nodes)
45
+
46
+ positioned_nodes = []
47
+ current_x = 0.0
48
+
49
+ for layer in layers:
50
+ if not layer:
51
+ continue
52
+
53
+ layer_nodes = [(node_id, node_dict[node_id]) for node_id in layer]
54
+ total_height = _calculate_layer_height(layer_nodes, node_spacing)
55
+
56
+ current_y = -total_height / 2.0
57
+
58
+ for node_id in layer:
59
+ node_data = node_dict[node_id]
60
+ node_height = node_data.height or 100.0
61
+
62
+ new_position = NodeDisplayPosition(x=current_x, y=current_y)
63
+
64
+ updated_data = NodeDisplayData(
65
+ position=new_position, width=node_data.width, height=node_data.height, comment=node_data.comment
66
+ )
67
+
68
+ positioned_nodes.append((node_id, updated_data))
69
+ current_y += node_height + node_spacing
70
+
71
+ current_x += layer_spacing
72
+
73
+ return positioned_nodes
74
+
75
+
76
+ def _topological_sort_layers(
77
+ graph: Dict[str, List[str]], in_degree: Dict[str, int], all_nodes: Set[str]
78
+ ) -> List[List[str]]:
79
+ """
80
+ Perform topological sorting and group nodes into layers.
81
+
82
+ Returns:
83
+ List of layers, where each layer is a list of node IDs
84
+ """
85
+ layers = []
86
+ remaining_nodes = set(all_nodes)
87
+
88
+ while remaining_nodes:
89
+ current_layer = []
90
+ for node in remaining_nodes:
91
+ if in_degree[node] == 0:
92
+ current_layer.append(node)
93
+
94
+ if not current_layer:
95
+ current_layer = [next(iter(remaining_nodes))]
96
+
97
+ layers.append(current_layer)
98
+
99
+ for node in current_layer:
100
+ remaining_nodes.remove(node)
101
+ for neighbor in graph[node]:
102
+ if neighbor in remaining_nodes:
103
+ in_degree[neighbor] -= 1
104
+
105
+ return layers
106
+
107
+
108
+ def _calculate_layer_height(layer_nodes: List[Tuple[str, NodeDisplayData]], node_spacing: float) -> float:
109
+ """
110
+ Calculate the total height needed for a layer of nodes.
111
+
112
+ Args:
113
+ layer_nodes: List of (node_id, NodeDisplayData) tuples in the layer
114
+ node_spacing: Spacing between nodes
115
+
116
+ Returns:
117
+ Total height needed for the layer
118
+ """
119
+ if not layer_nodes:
120
+ return 0.0
121
+
122
+ total_height = 0.0
123
+ for i, (_, node_data) in enumerate(layer_nodes):
124
+ node_height = node_data.height or 100.0
125
+ total_height += node_height
126
+
127
+ if i < len(layer_nodes) - 1:
128
+ total_height += node_spacing
129
+
130
+ return total_height
@@ -31,6 +31,7 @@ from vellum.workflows.expressions.or_ import OrExpression
31
31
  from vellum.workflows.expressions.parse_json import ParseJsonExpression
32
32
  from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
33
33
  from vellum.workflows.references.constant import ConstantValueReference
34
+ from vellum.workflows.references.environment_variable import EnvironmentVariableReference
34
35
  from vellum.workflows.references.execution_count import ExecutionCountReference
35
36
  from vellum.workflows.references.lazy import LazyReference
36
37
  from vellum.workflows.references.output import OutputReference
@@ -239,6 +240,12 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
239
240
  "vellum_secret_name": value.name,
240
241
  }
241
242
 
243
+ if isinstance(value, EnvironmentVariableReference):
244
+ return {
245
+ "type": "ENVIRONMENT_VARIABLE",
246
+ "environment_variable": value.name,
247
+ }
248
+
242
249
  if isinstance(value, ExecutionCountReference):
243
250
  node_class_display = display_context.global_node_displays[value.node_class]
244
251
 
File without changes