vellum-ai 0.10.3__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 (61) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/workflows/events/tests/test_event.py +30 -0
  3. vellum/workflows/events/types.py +57 -3
  4. vellum/workflows/nodes/__init__.py +6 -7
  5. vellum/workflows/nodes/bases/base.py +0 -1
  6. vellum/workflows/nodes/core/inline_subworkflow_node/node.py +2 -1
  7. vellum/workflows/nodes/core/map_node/node.py +1 -1
  8. vellum/workflows/nodes/core/retry_node/node.py +1 -0
  9. vellum/workflows/nodes/core/templating_node/node.py +5 -1
  10. vellum/workflows/nodes/core/try_node/node.py +66 -27
  11. vellum/workflows/nodes/core/try_node/tests/test_node.py +39 -8
  12. vellum/workflows/nodes/displayable/__init__.py +2 -0
  13. vellum/workflows/nodes/displayable/bases/api_node/node.py +3 -3
  14. vellum/workflows/nodes/displayable/code_execution_node/node.py +5 -2
  15. vellum/workflows/nodes/displayable/final_output_node/node.py +6 -2
  16. vellum/workflows/nodes/displayable/note_node/__init__.py +5 -0
  17. vellum/workflows/nodes/displayable/note_node/node.py +10 -0
  18. vellum/workflows/nodes/displayable/tests/test_inline_text_prompt_node.py +10 -11
  19. vellum/workflows/nodes/utils.py +2 -0
  20. vellum/workflows/outputs/base.py +26 -2
  21. vellum/workflows/runner/runner.py +41 -27
  22. vellum/workflows/state/tests/test_state.py +2 -0
  23. vellum/workflows/types/tests/test_utils.py +9 -0
  24. vellum/workflows/types/utils.py +1 -1
  25. vellum/workflows/utils/vellum_variables.py +13 -1
  26. vellum/workflows/workflows/base.py +24 -1
  27. {vellum_ai-0.10.3.dist-info → vellum_ai-0.10.6.dist-info}/METADATA +8 -6
  28. {vellum_ai-0.10.3.dist-info → vellum_ai-0.10.6.dist-info}/RECORD +61 -56
  29. vellum_cli/CONTRIBUTING.md +66 -0
  30. vellum_cli/README.md +3 -0
  31. vellum_ee/workflows/display/base.py +2 -1
  32. vellum_ee/workflows/display/nodes/base_node_display.py +27 -4
  33. vellum_ee/workflows/display/nodes/vellum/__init__.py +2 -0
  34. vellum_ee/workflows/display/nodes/vellum/api_node.py +3 -3
  35. vellum_ee/workflows/display/nodes/vellum/code_execution_node.py +4 -4
  36. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +86 -41
  37. vellum_ee/workflows/display/nodes/vellum/guardrail_node.py +3 -3
  38. vellum_ee/workflows/display/nodes/vellum/inline_prompt_node.py +4 -5
  39. vellum_ee/workflows/display/nodes/vellum/inline_subworkflow_node.py +9 -9
  40. vellum_ee/workflows/display/nodes/vellum/map_node.py +5 -5
  41. vellum_ee/workflows/display/nodes/vellum/note_node.py +32 -0
  42. vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py +5 -5
  43. vellum_ee/workflows/display/nodes/vellum/search_node.py +6 -10
  44. vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py +2 -2
  45. vellum_ee/workflows/display/nodes/vellum/templating_node.py +4 -5
  46. vellum_ee/workflows/display/nodes/vellum/try_node.py +16 -4
  47. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py +7 -3
  48. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py +127 -101
  49. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py +6 -5
  50. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +77 -64
  51. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py +4 -3
  52. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py +6 -6
  53. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py +6 -6
  54. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py +4 -3
  55. vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py +7 -6
  56. vellum_ee/workflows/display/workflows/base_workflow_display.py +14 -9
  57. vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py +2 -7
  58. vellum_ee/workflows/display/workflows/vellum_workflow_display.py +18 -16
  59. {vellum_ai-0.10.3.dist-info → vellum_ai-0.10.6.dist-info}/LICENSE +0 -0
  60. {vellum_ai-0.10.3.dist-info → vellum_ai-0.10.6.dist-info}/WHEEL +0 -0
  61. {vellum_ai-0.10.3.dist-info → vellum_ai-0.10.6.dist-info}/entry_points.txt +0 -0
@@ -2,11 +2,12 @@ from unittest import mock
2
2
 
3
3
  from deepdiff import DeepDiff
4
4
 
5
- from tests.workflows.basic_inline_subworkflow.workflow import BasicInlineSubworkflowWorkflow
6
5
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
7
6
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
8
7
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
9
8
 
9
+ from tests.workflows.basic_inline_subworkflow.workflow import BasicInlineSubworkflowWorkflow
10
+
10
11
 
11
12
  def test_serialize_workflow():
12
13
  # GIVEN a Workflow
@@ -37,17 +38,17 @@ def test_serialize_workflow():
37
38
  "id": "fa73da37-34c3-47a9-be58-69cc6cdbfca5",
38
39
  "key": "city",
39
40
  "type": "STRING",
40
- "required": None,
41
+ "required": True,
41
42
  "default": None,
42
- "extensions": None,
43
+ "extensions": { "color": None },
43
44
  },
44
45
  {
45
46
  "id": "aba1e6e0-dfa7-4c15-a4e6-aec6feebfaca",
46
47
  "key": "date",
47
48
  "type": "STRING",
48
- "required": None,
49
+ "required": True,
49
50
  "default": None,
50
- "extensions": None,
51
+ "extensions": { "color": None },
51
52
  },
52
53
  ],
53
54
  input_variables,
@@ -306,8 +307,8 @@ def test_serialize_workflow():
306
307
  },
307
308
  "input_variables": [{"id": "704c4640-bfda-44f0-8da3-e9cfc4f21cf2", "key": "metro", "type": "STRING"}],
308
309
  "output_variables": [
309
- {"id": "86dd0202-c141-48a3-8382-2da60372e77c", "key": "temperature", "type": "NUMBER"},
310
- {"id": "0a7192da-5576-4933-bba4-de8adf5d7996", "key": "reasoning", "type": "STRING"},
310
+ {"id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622", "key": "temperature", "type": "NUMBER"},
311
+ {"id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82", "key": "reasoning", "type": "STRING"},
311
312
  ],
312
313
  },
313
314
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
@@ -338,64 +339,74 @@ def test_serialize_workflow():
338
339
  ignore_order=True,
339
340
  )
340
341
 
342
+ temperature_terminal_node = next(node for node in workflow_raw_data["nodes"][2:] if node['data']['name'] == 'temperature')
343
+ reasoning_terminal_node = next(node for node in workflow_raw_data["nodes"][2:] if node['data']['name'] == 'reasoning')
344
+
341
345
  assert not DeepDiff(
342
- [
343
- {
344
- "id": "31b74695-3f1c-47cf-8be8-a4d86cc589e8",
345
- "type": "TERMINAL",
346
- "definition": {
347
- "bases": [
348
- {
349
- "bases": [],
350
- "module": [
351
- "vellum",
352
- "workflows",
353
- "nodes",
354
- "bases",
355
- "base",
356
- ],
357
- "name": "BaseNode",
358
- },
359
- ],
360
- "module": [
361
- "vellum",
362
- "workflows",
363
- "nodes",
364
- "displayable",
365
- "final_output_node",
366
- "node",
367
- ],
368
- "name": "FinalOutputNode",
369
- },
370
- "data": {
371
- "label": "Final Output",
372
- "name": "reasoning",
373
- "target_handle_id": "8b525943-6c27-414b-a329-e29c0b217f72",
374
- "output_id": "7444a019-081a-4e10-a528-3249299159f7",
375
- "output_type": "STRING",
376
- "node_input_id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
377
- },
378
- "inputs": [
346
+ {
347
+ "id": "31b74695-3f1c-47cf-8be8-a4d86cc589e8",
348
+ "type": "TERMINAL",
349
+ "definition": {
350
+ "bases": [
379
351
  {
380
- "id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
381
- "key": "node_input",
382
- "value": {
383
- "rules": [
384
- {
385
- "type": "NODE_OUTPUT",
386
- "data": {
387
- "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
388
- "output_id": "0a7192da-5576-4933-bba4-de8adf5d7996",
389
- },
390
- }
391
- ],
392
- "combinator": "OR",
393
- },
394
- }
352
+ "bases": [],
353
+ "module": [
354
+ "vellum",
355
+ "workflows",
356
+ "nodes",
357
+ "bases",
358
+ "base",
359
+ ],
360
+ "name": "BaseNode",
361
+ },
395
362
  ],
396
- "display_data": {"position": {"x": 0.0, "y": 0.0}},
363
+ "module": [
364
+ "vellum",
365
+ "workflows",
366
+ "nodes",
367
+ "displayable",
368
+ "final_output_node",
369
+ "node",
370
+ ],
371
+ "name": "FinalOutputNode",
397
372
  },
398
- {
373
+ "data": {
374
+ "label": "Final Output",
375
+ "name": "reasoning",
376
+ "target_handle_id": "8b525943-6c27-414b-a329-e29c0b217f72",
377
+ "output_id": "7444a019-081a-4e10-a528-3249299159f7",
378
+ "output_type": "STRING",
379
+ "node_input_id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
380
+ },
381
+ "inputs": [
382
+ {
383
+ "id": "736473c8-b0b4-4cdd-b743-6453dd5306fc",
384
+ "key": "node_input",
385
+ "value": {
386
+ "rules": [
387
+ {
388
+ "type": "NODE_OUTPUT",
389
+ "data": {
390
+ "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
391
+ "output_id": "fad5dd9f-3328-4e70-ad55-65a5325a4a82",
392
+ },
393
+ }
394
+ ],
395
+ "combinator": "OR",
396
+ },
397
+ }
398
+ ],
399
+ "display_data": {"position": {"x": 0.0, "y": 0.0}},
400
+ },
401
+ reasoning_terminal_node,
402
+ ignore_order=True,
403
+ # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
404
+ # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
405
+ exclude_regex_paths=r"root\['inputs'\]\[0\]\['value'\]\['rules'\]\[0\]\['data'\]\['output_id'\]"
406
+ )
407
+
408
+ assert not DeepDiff(
409
+ {
399
410
  "id": "0779b232-82ab-4dbe-a340-6a85e6ab3368",
400
411
  "type": "TERMINAL",
401
412
  "definition": {
@@ -440,7 +451,7 @@ def test_serialize_workflow():
440
451
  "type": "NODE_OUTPUT",
441
452
  "data": {
442
453
  "node_id": "080e4343-c7ce-4f82-b9dd-e94c8cc92239",
443
- "output_id": "86dd0202-c141-48a3-8382-2da60372e77c",
454
+ "output_id": "2fc57139-7420-49e5-96a6-dcbb3ff5d622",
444
455
  },
445
456
  }
446
457
  ],
@@ -450,9 +461,11 @@ def test_serialize_workflow():
450
461
  ],
451
462
  "display_data": {"position": {"x": 0.0, "y": 0.0}},
452
463
  },
453
- ],
454
- workflow_raw_data["nodes"][2:],
464
+ temperature_terminal_node,
455
465
  ignore_order=True,
466
+ # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
467
+ # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
468
+ exclude_regex_paths=r"root\['inputs'\]\[0\]\['value'\]\['rules'\]\[0\]\['data'\]\['output_id'\]"
456
469
  )
457
470
 
458
471
  # AND each edge should be serialized correctly
@@ -2,11 +2,12 @@ from unittest import mock
2
2
 
3
3
  from deepdiff import DeepDiff
4
4
 
5
- from tests.workflows.basic_map_node.workflow import SimpleMapExample
6
5
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
7
6
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
8
7
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
9
8
 
9
+ from tests.workflows.basic_map_node.workflow import SimpleMapExample
10
+
10
11
 
11
12
  def test_serialize_workflow():
12
13
  # GIVEN a Workflow that uses a MapNode
@@ -35,9 +36,9 @@ def test_serialize_workflow():
35
36
  "id": "db2eb237-38e4-417a-8bfc-5bda0f3165ca",
36
37
  "key": "fruits",
37
38
  "type": "JSON",
38
- "required": None,
39
+ "required": True,
39
40
  "default": None,
40
- "extensions": None,
41
+ "extensions": { "color": None },
41
42
  },
42
43
  ],
43
44
  input_variables,
@@ -4,11 +4,11 @@ from uuid import uuid4
4
4
  from deepdiff import DeepDiff
5
5
 
6
6
  from vellum import DeploymentRead
7
-
8
- from tests.workflows.basic_text_prompt_deployment.workflow import BasicTextPromptDeployment
9
7
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
10
8
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
11
9
 
10
+ from tests.workflows.basic_text_prompt_deployment.workflow import BasicTextPromptDeployment
11
+
12
12
 
13
13
  def test_serialize_workflow(vellum_client):
14
14
  # GIVEN a Workflow with stubbed out API calls
@@ -46,17 +46,17 @@ def test_serialize_workflow(vellum_client):
46
46
  "id": "52995b50-84c9-465f-8a4b-a4ee2a92e388",
47
47
  "key": "city",
48
48
  "type": "STRING",
49
- "required": None,
49
+ "required": True,
50
50
  "default": None,
51
- "extensions": None,
51
+ "extensions": { "color": None },
52
52
  },
53
53
  {
54
54
  "id": "aa3ca842-250c-4a3f-853f-23928c28d0f8",
55
55
  "key": "date",
56
56
  "type": "STRING",
57
- "required": None,
57
+ "required": True,
58
58
  "default": None,
59
- "extensions": None,
59
+ "extensions": { "color": None },
60
60
  },
61
61
  ],
62
62
  input_variables,
@@ -4,11 +4,11 @@ from uuid import uuid4
4
4
  from deepdiff import DeepDiff
5
5
 
6
6
  from vellum import WorkflowDeploymentRead
7
-
8
- from tests.workflows.basic_subworkflow_deployment.workflow import BasicSubworkflowDeploymentWorkflow
9
7
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
10
8
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
11
9
 
10
+ from tests.workflows.basic_subworkflow_deployment.workflow import BasicSubworkflowDeploymentWorkflow
11
+
12
12
 
13
13
  def test_serialize_workflow(vellum_client):
14
14
  # GIVEN a Workflow with stubbed out API calls
@@ -46,17 +46,17 @@ def test_serialize_workflow(vellum_client):
46
46
  "id": "693cc9a5-8d74-4a58-bdcf-2b4989cdf250",
47
47
  "key": "city",
48
48
  "type": "STRING",
49
- "required": None,
49
+ "required": True,
50
50
  "default": None,
51
- "extensions": None,
51
+ "extensions": { "color": None },
52
52
  },
53
53
  {
54
54
  "id": "19a78824-9a98-4ae8-a1fc-61f81a422a17",
55
55
  "key": "date",
56
56
  "type": "STRING",
57
- "required": None,
57
+ "required": True,
58
58
  "default": None,
59
- "extensions": None,
59
+ "extensions": { "color": None },
60
60
  },
61
61
  ],
62
62
  input_variables,
@@ -1,7 +1,8 @@
1
- from tests.workflows.basic_final_output_node.workflow import BasicFinalOutputNodeWorkflow
2
1
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
3
2
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
4
3
 
4
+ from tests.workflows.basic_final_output_node.workflow import BasicFinalOutputNodeWorkflow
5
+
5
6
 
6
7
  def test_serialize_workflow():
7
8
  # GIVEN a Workflow that uses a Final Output Node
@@ -26,9 +27,9 @@ def test_serialize_workflow():
26
27
  "id": "e39a7b63-de15-490a-ae9b-8112c767aea0",
27
28
  "key": "input",
28
29
  "type": "STRING",
29
- "required": None,
30
+ "required": True,
30
31
  "default": None,
31
- "extensions": None,
32
+ "extensions": { "color": None },
32
33
  }
33
34
  ]
34
35
 
@@ -3,12 +3,13 @@ from unittest import mock
3
3
 
4
4
  from deepdiff import DeepDiff
5
5
 
6
- from tests.workflows.complex_final_output_node.missing_final_output_node import MissingFinalOutputNodeWorkflow
7
- from tests.workflows.complex_final_output_node.missing_workflow_output import MissingWorkflowOutputWorkflow
8
6
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
9
7
  from vellum_ee.workflows.display.workflows import VellumWorkflowDisplay
10
8
  from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
11
9
 
10
+ from tests.workflows.complex_final_output_node.missing_final_output_node import MissingFinalOutputNodeWorkflow
11
+ from tests.workflows.complex_final_output_node.missing_workflow_output import MissingWorkflowOutputWorkflow
12
+
12
13
 
13
14
  def test_serialize_workflow__missing_final_output_node():
14
15
  # GIVEN a Workflow that is missing a Terminal Node
@@ -40,17 +41,17 @@ def test_serialize_workflow__missing_final_output_node():
40
41
  "id": "da086239-d743-4246-b666-5c91e22fb88c",
41
42
  "key": "alpha",
42
43
  "type": "STRING",
44
+ "required": True,
43
45
  "default": None,
44
- "required": None,
45
- "extensions": None,
46
+ "extensions": { "color": None },
46
47
  },
47
48
  {
48
49
  "id": "a8b6c5d4-a0e9-4457-834b-46b633c466a6",
49
50
  "key": "beta",
50
51
  "type": "STRING",
52
+ "required": True,
51
53
  "default": None,
52
- "required": None,
53
- "extensions": None,
54
+ "extensions": { "color": None },
54
55
  },
55
56
  ],
56
57
  input_variables,
@@ -5,6 +5,15 @@ import logging
5
5
  from uuid import UUID
6
6
  from typing import Any, Dict, Generic, Optional, Tuple, Type, get_args
7
7
 
8
+ from vellum.workflows.descriptors.base import BaseDescriptor
9
+ from vellum.workflows.edges import Edge
10
+ from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
11
+ from vellum.workflows.nodes.bases import BaseNode
12
+ from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
13
+ from vellum.workflows.ports import Port
14
+ from vellum.workflows.references import OutputReference, WorkflowInputReference
15
+ from vellum.workflows.types.core import JsonObject
16
+ from vellum.workflows.types.generics import WorkflowType
8
17
  from vellum_ee.workflows.display.base import (
9
18
  EdgeDisplayOverridesType,
10
19
  EdgeDisplayType,
@@ -21,15 +30,6 @@ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_di
21
30
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
22
31
  from vellum_ee.workflows.display.types import NodeDisplayType, WorkflowDisplayContext
23
32
  from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
24
- from vellum.workflows.descriptors.base import BaseDescriptor
25
- from vellum.workflows.edges import Edge
26
- from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
27
- from vellum.workflows.nodes.bases import BaseNode
28
- from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
29
- from vellum.workflows.ports import Port
30
- from vellum.workflows.references import OutputReference, WorkflowInputReference
31
- from vellum.workflows.types.core import JsonObject
32
- from vellum.workflows.types.generics import WorkflowType
33
33
 
34
34
  logger = logging.getLogger(__name__)
35
35
 
@@ -128,6 +128,8 @@ class BaseWorkflowDisplay(
128
128
  inner_node_display = self._get_node_display(inner_node)
129
129
  self._enrich_node_output_displays(inner_node, inner_node_display, node_output_displays)
130
130
 
131
+ # TODO: Make sure this output ID matches the workflow output ID of the subworkflow node's workflow
132
+ # https://app.shortcut.com/vellum/story/5660/fix-output-id-in-subworkflow-nodes
131
133
  node_output_displays[node_output] = node, node_display.get_node_output_display(node_output)
132
134
 
133
135
  def _enrich_node_port_displays(
@@ -179,6 +181,9 @@ class BaseWorkflowDisplay(
179
181
 
180
182
  node_displays: Dict[Type[BaseNode], NodeDisplayType] = {}
181
183
  port_displays: Dict[Port, PortDisplay] = {}
184
+
185
+ # TODO: We should still serialize nodes that are in the workflow's directory but aren't used in the graph.
186
+ # https://app.shortcut.com/vellum/story/5394
182
187
  for node in self._workflow.get_nodes():
183
188
  node_display = self._get_node_display(node)
184
189
  node_displays[node] = node_display
@@ -1,7 +1,7 @@
1
1
  from typing import Optional, Type
2
2
 
3
- from vellum_ee.workflows.display.types import WorkflowDisplayContext, WorkflowDisplayType
4
3
  from vellum.workflows.types.generics import WorkflowType
4
+ from vellum_ee.workflows.display.types import WorkflowDisplayContext, WorkflowDisplayType
5
5
 
6
6
 
7
7
  def get_workflow_display(
@@ -24,9 +24,4 @@ def get_workflow_display(
24
24
  except IndexError:
25
25
  return base_display_class(workflow_class)
26
26
 
27
- if not issubclass(workflow_display_class, base_display_class):
28
- raise TypeError(
29
- f"Expected to find a subclass of '{base_display_class.__name__}' for workflow class '{workflow_class.__name__}'" # noqa: E501
30
- )
31
-
32
- return workflow_display_class(workflow_class, parent_display_context=parent_display_context)
27
+ return workflow_display_class(workflow_class, parent_display_context=parent_display_context) # type: ignore[return-value]
@@ -2,11 +2,22 @@ import logging
2
2
  from uuid import UUID
3
3
  from typing import Dict, List, Optional, Type, cast
4
4
 
5
+ from vellum.workflows.descriptors.base import BaseDescriptor
6
+ from vellum.workflows.edges import Edge
7
+ from vellum.workflows.nodes.bases import BaseNode
8
+ from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
9
+ from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
10
+ from vellum.workflows.ports import Port
11
+ from vellum.workflows.references import WorkflowInputReference
12
+ from vellum.workflows.references.output import OutputReference
13
+ from vellum.workflows.types.core import JsonArray, JsonObject
14
+ from vellum.workflows.types.generics import WorkflowType
5
15
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
6
16
  from vellum_ee.workflows.display.nodes.types import PortDisplay
17
+ from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
7
18
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
8
19
  from vellum_ee.workflows.display.utils.uuids import uuid4_from_hash
9
- from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
20
+ from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type, primitive_to_vellum_value
10
21
  from vellum_ee.workflows.display.vellum import (
11
22
  EdgeVellumDisplay,
12
23
  EdgeVellumDisplayOverrides,
@@ -21,16 +32,6 @@ from vellum_ee.workflows.display.vellum import (
21
32
  WorkflowOutputVellumDisplayOverrides,
22
33
  )
23
34
  from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
24
- from vellum.workflows.descriptors.base import BaseDescriptor
25
- from vellum.workflows.edges import Edge
26
- from vellum.workflows.nodes.bases import BaseNode
27
- from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
28
- from vellum.workflows.nodes.utils import get_wrapped_node, has_wrapped_node
29
- from vellum.workflows.ports import Port
30
- from vellum.workflows.references import WorkflowInputReference
31
- from vellum.workflows.references.output import OutputReference
32
- from vellum.workflows.types.core import JsonArray, JsonObject
33
- from vellum.workflows.types.generics import WorkflowType
34
35
 
35
36
  logger = logging.getLogger(__name__)
36
37
 
@@ -56,16 +57,17 @@ class VellumWorkflowDisplay(
56
57
  def serialize(self, raise_errors: bool = True) -> JsonObject:
57
58
  input_variables: JsonArray = []
58
59
  for workflow_input, workflow_input_display in self.display_context.workflow_input_displays.items():
60
+ default = primitive_to_vellum_value(raise_if_descriptor(workflow_input.instance)) if workflow_input.instance else None
61
+ required = type(None) not in workflow_input.types
62
+
59
63
  input_variables.append(
60
64
  {
61
65
  "id": str(workflow_input_display.id),
62
66
  "key": workflow_input.name,
63
67
  "type": infer_vellum_variable_type(workflow_input),
64
- # TODO: Add support for serializing default, required, and extensions
65
- # https://app.shortcut.com/vellum/story/5429
66
- "default": None,
67
- "required": None,
68
- "extensions": None,
68
+ "default": default.dict() if default else None,
69
+ "required": required,
70
+ "extensions": { "color" : workflow_input_display.color },
69
71
  }
70
72
  )
71
73