vellum-ai 0.14.44__py3-none-any.whl → 0.14.46__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (56) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/client/core/pydantic_utilities.py +7 -1
  3. vellum/workflows/nodes/bases/base.py +1 -0
  4. vellum/workflows/nodes/bases/tests/test_base_node.py +20 -0
  5. vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py +8 -14
  6. vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py +62 -0
  7. vellum/workflows/nodes/displayable/code_execution_node/utils.py +3 -54
  8. vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py +5 -6
  9. vellum/workflows/nodes/utils.py +4 -0
  10. vellum/workflows/ports/port.py +13 -3
  11. vellum/workflows/types/code_execution_node_wrappers.py +64 -0
  12. vellum/workflows/types/tests/test_utils.py +3 -3
  13. vellum/workflows/types/utils.py +31 -10
  14. vellum/workflows/vellum_client.py +19 -7
  15. {vellum_ai-0.14.44.dist-info → vellum_ai-0.14.46.dist-info}/METADATA +1 -1
  16. {vellum_ai-0.14.44.dist-info → vellum_ai-0.14.46.dist-info}/RECORD +56 -53
  17. vellum_cli/config.py +7 -2
  18. vellum_cli/push.py +5 -1
  19. vellum_cli/tests/test_push.py +192 -8
  20. vellum_ee/workflows/display/nodes/base_node_display.py +4 -173
  21. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +1 -1
  22. vellum_ee/workflows/display/nodes/vellum/final_output_node.py +2 -1
  23. vellum_ee/workflows/display/nodes/vellum/prompt_deployment_node.py +5 -6
  24. vellum_ee/workflows/display/nodes/vellum/retry_node.py +3 -3
  25. vellum_ee/workflows/display/nodes/vellum/subworkflow_deployment_node.py +5 -6
  26. vellum_ee/workflows/display/nodes/vellum/tests/test_prompt_deployment_node.py +106 -0
  27. vellum_ee/workflows/display/nodes/vellum/tests/test_subworkflow_deployment_node.py +109 -0
  28. vellum_ee/workflows/display/nodes/vellum/try_node.py +3 -3
  29. vellum_ee/workflows/display/tests/test_base_workflow_display.py +1 -0
  30. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py +73 -111
  31. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_api_node_serialization.py +0 -1
  32. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_code_execution_node_serialization.py +0 -3
  33. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py +0 -4
  34. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_default_state_serialization.py +0 -1
  35. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py +0 -1
  36. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py +0 -1
  37. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_guardrail_node_serialization.py +0 -1
  38. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +18 -2
  39. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py +10 -1
  40. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py +0 -1
  41. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_prompt_deployment_serialization.py +2 -3
  42. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_search_node_serialization.py +0 -1
  43. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_subworkflow_deployment_serialization.py +2 -3
  44. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_templating_node_serialization.py +0 -1
  45. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_terminal_node_serialization.py +1 -2
  46. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_try_node_serialization.py +0 -1
  47. vellum_ee/workflows/display/tests/workflow_serialization/test_complex_terminal_node_serialization.py +5 -55
  48. vellum_ee/workflows/display/types.py +3 -0
  49. vellum_ee/workflows/display/utils/expressions.py +222 -2
  50. vellum_ee/workflows/display/utils/vellum.py +1 -79
  51. vellum_ee/workflows/display/workflows/base_workflow_display.py +59 -37
  52. vellum_ee/workflows/display/workflows/get_vellum_workflow_display_class.py +3 -0
  53. vellum_ee/workflows/display/workflows/tests/test_workflow_display.py +98 -0
  54. {vellum_ai-0.14.44.dist-info → vellum_ai-0.14.46.dist-info}/LICENSE +0 -0
  55. {vellum_ai-0.14.44.dist-info → vellum_ai-0.14.46.dist-info}/WHEEL +0 -0
  56. {vellum_ai-0.14.44.dist-info → vellum_ai-0.14.46.dist-info}/entry_points.txt +0 -0
@@ -20,45 +20,10 @@ from typing import (
20
20
  from vellum.client.types.code_resource_definition import CodeResourceDefinition
21
21
  from vellum.workflows import BaseWorkflow
22
22
  from vellum.workflows.constants import undefined
23
- from vellum.workflows.descriptors.base import BaseDescriptor
24
- from vellum.workflows.expressions.accessor import AccessorExpression
25
- from vellum.workflows.expressions.and_ import AndExpression
26
- from vellum.workflows.expressions.begins_with import BeginsWithExpression
27
- from vellum.workflows.expressions.between import BetweenExpression
28
- from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
29
- from vellum.workflows.expressions.contains import ContainsExpression
30
- from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
31
- from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
32
- from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
33
- from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
34
- from vellum.workflows.expressions.ends_with import EndsWithExpression
35
- from vellum.workflows.expressions.equals import EqualsExpression
36
- from vellum.workflows.expressions.greater_than import GreaterThanExpression
37
- from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
38
- from vellum.workflows.expressions.in_ import InExpression
39
- from vellum.workflows.expressions.is_nil import IsNilExpression
40
- from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
41
- from vellum.workflows.expressions.is_not_null import IsNotNullExpression
42
- from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
43
- from vellum.workflows.expressions.is_null import IsNullExpression
44
- from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
45
- from vellum.workflows.expressions.less_than import LessThanExpression
46
- from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
47
- from vellum.workflows.expressions.not_between import NotBetweenExpression
48
- from vellum.workflows.expressions.not_in import NotInExpression
49
- from vellum.workflows.expressions.or_ import OrExpression
50
- from vellum.workflows.expressions.parse_json import ParseJsonExpression
51
23
  from vellum.workflows.nodes.bases.base import BaseNode
52
- from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
53
24
  from vellum.workflows.nodes.utils import get_unadorned_node, get_wrapped_node
54
25
  from vellum.workflows.ports import Port
55
26
  from vellum.workflows.references import OutputReference
56
- from vellum.workflows.references.constant import ConstantValueReference
57
- from vellum.workflows.references.execution_count import ExecutionCountReference
58
- from vellum.workflows.references.lazy import LazyReference
59
- from vellum.workflows.references.state_value import StateValueReference
60
- from vellum.workflows.references.vellum_secret import VellumSecretReference
61
- from vellum.workflows.references.workflow_input import WorkflowInputReference
62
27
  from vellum.workflows.types.core import JsonArray, JsonObject
63
28
  from vellum.workflows.types.generics import NodeType
64
29
  from vellum.workflows.types.utils import get_original_base
@@ -68,10 +33,8 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
68
33
  from vellum_ee.workflows.display.editor.types import NodeDisplayComment, NodeDisplayData
69
34
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
70
35
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
71
- from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
72
- from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
36
+ from vellum_ee.workflows.display.utils.expressions import serialize_condition, serialize_value
73
37
  from vellum_ee.workflows.display.utils.registry import register_node_display_class
74
- from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator
75
38
 
76
39
  if TYPE_CHECKING:
77
40
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
@@ -157,7 +120,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
157
120
  {
158
121
  "id": id,
159
122
  "name": attribute.name,
160
- "value": self.serialize_value(display_context, attribute.instance),
123
+ "value": serialize_value(display_context, attribute.instance),
161
124
  }
162
125
  )
163
126
  except ValueError as e:
@@ -184,7 +147,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
184
147
  for output in node.Outputs:
185
148
  type = primitive_type_to_vellum_variable_type(output)
186
149
  value = (
187
- self.serialize_value(display_context, output.instance)
150
+ serialize_value(display_context, output.instance)
188
151
  if output.instance is not None and output.instance != undefined
189
152
  else None
190
153
  )
@@ -229,7 +192,7 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
229
192
  "name": port.name,
230
193
  "type": port._condition_type.value,
231
194
  "expression": (
232
- self.serialize_condition(display_context, port._condition) if port._condition else None
195
+ serialize_condition(display_context, port._condition) if port._condition else None
233
196
  ),
234
197
  }
235
198
  )
@@ -405,137 +368,5 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
405
368
 
406
369
  return NodeDisplayData()
407
370
 
408
- def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
409
- if isinstance(
410
- condition,
411
- (
412
- IsNullExpression,
413
- IsNotNullExpression,
414
- IsNilExpression,
415
- IsNotNilExpression,
416
- IsUndefinedExpression,
417
- IsNotUndefinedExpression,
418
- ParseJsonExpression,
419
- ),
420
- ):
421
- lhs = self.serialize_value(display_context, condition._expression)
422
- return {
423
- "type": "UNARY_EXPRESSION",
424
- "lhs": lhs,
425
- "operator": convert_descriptor_to_operator(condition),
426
- }
427
- elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
428
- base = self.serialize_value(display_context, condition._value)
429
- lhs = self.serialize_value(display_context, condition._start)
430
- rhs = self.serialize_value(display_context, condition._end)
431
-
432
- return {
433
- "type": "TERNARY_EXPRESSION",
434
- "base": base,
435
- "operator": convert_descriptor_to_operator(condition),
436
- "lhs": lhs,
437
- "rhs": rhs,
438
- }
439
- elif isinstance(
440
- condition,
441
- (
442
- AndExpression,
443
- BeginsWithExpression,
444
- CoalesceExpression,
445
- ContainsExpression,
446
- DoesNotBeginWithExpression,
447
- DoesNotContainExpression,
448
- DoesNotEndWithExpression,
449
- DoesNotEqualExpression,
450
- EndsWithExpression,
451
- EqualsExpression,
452
- GreaterThanExpression,
453
- GreaterThanOrEqualToExpression,
454
- InExpression,
455
- LessThanExpression,
456
- LessThanOrEqualToExpression,
457
- NotInExpression,
458
- OrExpression,
459
- ),
460
- ):
461
- lhs = self.serialize_value(display_context, condition._lhs)
462
- rhs = self.serialize_value(display_context, condition._rhs)
463
-
464
- return {
465
- "type": "BINARY_EXPRESSION",
466
- "lhs": lhs,
467
- "operator": convert_descriptor_to_operator(condition),
468
- "rhs": rhs,
469
- }
470
- elif isinstance(condition, AccessorExpression):
471
- return {
472
- "type": "BINARY_EXPRESSION",
473
- "lhs": self.serialize_value(display_context, condition._base),
474
- "operator": "accessField",
475
- "rhs": self.serialize_value(display_context, condition._field),
476
- }
477
-
478
- raise UnsupportedSerializationException(f"Unsupported condition type: {condition.__class__.__name__}")
479
-
480
- def serialize_value(self, display_context: "WorkflowDisplayContext", value: Any) -> JsonObject:
481
- if isinstance(value, ConstantValueReference):
482
- return self.serialize_value(display_context, value._value)
483
-
484
- if isinstance(value, LazyReference):
485
- child_descriptor = get_child_descriptor(value, display_context)
486
- return self.serialize_value(display_context, child_descriptor)
487
-
488
- if isinstance(value, WorkflowInputReference):
489
- workflow_input_display = display_context.global_workflow_input_displays[value]
490
- return {
491
- "type": "WORKFLOW_INPUT",
492
- "input_variable_id": str(workflow_input_display.id),
493
- }
494
-
495
- if isinstance(value, StateValueReference):
496
- state_value_display = display_context.global_state_value_displays[value]
497
- return {
498
- "type": "STATE_VALUE",
499
- "state_variable_id": str(state_value_display.id),
500
- }
501
-
502
- if isinstance(value, OutputReference):
503
- upstream_node, output_display = display_context.global_node_output_displays[value]
504
- upstream_node_display = display_context.global_node_displays[upstream_node]
505
-
506
- return {
507
- "type": "NODE_OUTPUT",
508
- "node_id": str(upstream_node_display.node_id),
509
- "node_output_id": str(output_display.id),
510
- }
511
-
512
- if isinstance(value, VellumSecretReference):
513
- return {
514
- "type": "VELLUM_SECRET",
515
- "vellum_secret_name": value.name,
516
- }
517
-
518
- if isinstance(value, ExecutionCountReference):
519
- node_class_display = display_context.global_node_displays[value.node_class]
520
-
521
- return {
522
- "type": "EXECUTION_COUNTER",
523
- "node_id": str(node_class_display.node_id),
524
- }
525
-
526
- if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
527
- raise ValueError("Nested references are not supported.")
528
-
529
- if not isinstance(value, BaseDescriptor):
530
- vellum_value = primitive_to_vellum_value(value)
531
- return {
532
- "type": "CONSTANT_VALUE",
533
- "value": vellum_value.dict(),
534
- }
535
-
536
- # If it's not any of the references we know about,
537
- # then try to serialize it as a nested value
538
- return self.serialize_condition(display_context, value)
539
-
540
371
 
541
372
  register_node_display_class(node_class=BaseNode, node_display_class=BaseNodeDisplay)
@@ -19,7 +19,7 @@ from vellum.workflows.utils.uuids import uuid4_from_hash
19
19
  from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
20
20
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
21
21
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
22
- from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator
22
+ from vellum_ee.workflows.display.utils.expressions import convert_descriptor_to_operator
23
23
  from vellum_ee.workflows.display.vellum import NodeInput
24
24
 
25
25
  _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
@@ -8,6 +8,7 @@ from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
8
8
  from vellum_ee.workflows.display.nodes.utils import to_kebab_case
9
9
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
10
10
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
11
+ from vellum_ee.workflows.display.utils.expressions import serialize_value
11
12
  from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
12
13
 
13
14
  _FinalOutputNodeType = TypeVar("_FinalOutputNodeType", bound=FinalOutputNode)
@@ -52,7 +53,7 @@ class BaseFinalOutputNodeDisplay(BaseNodeDisplay[_FinalOutputNodeType], Generic[
52
53
  "id": str(self._get_output_id()),
53
54
  "name": node.Outputs.value.name,
54
55
  "type": inferred_type,
55
- "value": self.serialize_value(display_context, node.Outputs.value.instance),
56
+ "value": serialize_value(display_context, node.Outputs.value.instance),
56
57
  }
57
58
  ],
58
59
  }
@@ -4,7 +4,6 @@ from typing import Generic, Optional, TypeVar, cast
4
4
  from vellum.workflows.nodes.displayable.prompt_deployment_node import PromptDeploymentNode
5
5
  from vellum.workflows.references import OutputReference
6
6
  from vellum.workflows.types.core import JsonObject
7
- from vellum.workflows.vellum_client import create_vellum_client
8
7
  from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
9
8
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
10
9
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -28,7 +27,10 @@ class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType]
28
27
  input_name=variable_name,
29
28
  value=variable_value,
30
29
  display_context=display_context,
31
- input_id=self.node_input_ids_by_name.get(variable_name),
30
+ input_id=self.node_input_ids_by_name.get(
31
+ f"{PromptDeploymentNode.prompt_inputs.name}.{variable_name}"
32
+ )
33
+ or self.node_input_ids_by_name.get(variable_name),
32
34
  )
33
35
  for variable_name, variable_value in prompt_inputs.items()
34
36
  ]
@@ -39,10 +41,7 @@ class BasePromptDeploymentNodeDisplay(BaseNodeDisplay[_PromptDeploymentNodeType]
39
41
  _, output_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.text)]
40
42
  _, array_display = display_context.global_node_output_displays[cast(OutputReference, node.Outputs.results)]
41
43
 
42
- # TODO: Pass through the name instead of retrieving the ID
43
- # https://app.shortcut.com/vellum/story/4702
44
- vellum_client = create_vellum_client()
45
- deployment = vellum_client.deployments.retrieve(
44
+ deployment = display_context.client.deployments.retrieve(
46
45
  id=str(raise_if_descriptor(node.deployment)),
47
46
  )
48
47
  ml_model_fallbacks = raise_if_descriptor(node.ml_model_fallbacks)
@@ -1,7 +1,6 @@
1
1
  import inspect
2
- from typing import Any, Generic, Tuple, Type, TypeVar, cast
2
+ from typing import Any, Generic, Tuple, Type, TypeVar
3
3
 
4
- from vellum.workflows.descriptors.base import BaseDescriptor
5
4
  from vellum.workflows.nodes.bases.base import BaseNode
6
5
  from vellum.workflows.nodes.core.retry_node.node import RetryNode
7
6
  from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
@@ -13,6 +12,7 @@ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_di
13
12
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
14
13
  from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
15
14
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
15
+ from vellum_ee.workflows.display.utils.expressions import serialize_value
16
16
 
17
17
  _RetryNodeType = TypeVar("_RetryNodeType", bound=RetryNode)
18
18
 
@@ -33,7 +33,7 @@ class BaseRetryNodeDisplay(BaseAdornmentNodeDisplay[_RetryNodeType], Generic[_Re
33
33
  {
34
34
  "id": id,
35
35
  "name": attribute.name,
36
- "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
36
+ "value": serialize_value(display_context, attribute.instance),
37
37
  }
38
38
  )
39
39
 
@@ -3,7 +3,6 @@ from typing import Generic, Optional, TypeVar
3
3
 
4
4
  from vellum.workflows.nodes import SubworkflowDeploymentNode
5
5
  from vellum.workflows.types.core import JsonObject
6
- from vellum.workflows.vellum_client import create_vellum_client
7
6
  from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
8
7
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
9
8
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -29,15 +28,15 @@ class BaseSubworkflowDeploymentNodeDisplay(
29
28
  input_name=variable_name,
30
29
  value=variable_value,
31
30
  display_context=display_context,
32
- input_id=self.node_input_ids_by_name.get(variable_name),
31
+ input_id=self.node_input_ids_by_name.get(
32
+ f"{SubworkflowDeploymentNode.subworkflow_inputs.name}.{variable_name}"
33
+ )
34
+ or self.node_input_ids_by_name.get(variable_name),
33
35
  )
34
36
  for variable_name, variable_value in subworkflow_inputs.items()
35
37
  ]
36
38
 
37
- # TODO: Pass through the name instead of retrieving the ID
38
- # https://app.shortcut.com/vellum/story/4702
39
- vellum_client = create_vellum_client()
40
- deployment = vellum_client.workflow_deployments.retrieve(
39
+ deployment = display_context.client.workflow_deployments.retrieve(
41
40
  id=str(raise_if_descriptor(node.deployment)),
42
41
  )
43
42
 
@@ -0,0 +1,106 @@
1
+ import pytest
2
+ from datetime import datetime
3
+ from uuid import UUID, uuid4
4
+ from typing import Type
5
+
6
+ from vellum.workflows import BaseWorkflow
7
+ from vellum.workflows.nodes import PromptDeploymentNode
8
+ from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
9
+
10
+
11
+ def _no_display_class(Node: Type[PromptDeploymentNode]): # type: ignore
12
+ return None
13
+
14
+
15
+ def _display_class_with_node_input_ids_by_name(Node: Type[PromptDeploymentNode]):
16
+ class PromptDeploymentNodeDisplay(PromptDeploymentNode[Node]): # type: ignore[valid-type]
17
+ node_input_ids_by_name = {"foo": UUID("6037747a-1d35-4094-b363-4369fc92c5d4")}
18
+
19
+ return PromptDeploymentNodeDisplay
20
+
21
+
22
+ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[PromptDeploymentNode]):
23
+ class PromptDeploymentNodeDisplay(PromptDeploymentNode[Node]): # type: ignore[valid-type]
24
+ node_input_ids_by_name = {"prompt_inputs.foo": UUID("6037747a-1d35-4094-b363-4369fc92c5d4")}
25
+
26
+ return PromptDeploymentNodeDisplay
27
+
28
+
29
+ @pytest.fixture
30
+ def mock_fetch_deployment(mocker):
31
+ # Create a mock deployment response
32
+ mock_deployment = mocker.Mock(
33
+ id="test-id",
34
+ name="test-deployment",
35
+ label="Test Deployment",
36
+ status="ACTIVE",
37
+ environment="DEVELOPMENT",
38
+ created=datetime.now(),
39
+ last_deployed_on=datetime.now(),
40
+ last_deployed_history_item_id=str(uuid4()),
41
+ input_variables=[],
42
+ output_variables=[],
43
+ description="Test deployment description",
44
+ )
45
+
46
+ # Patch the create_vellum_client function to return our mock client
47
+ mocker.patch("vellum.client.resources.deployments.client.DeploymentsClient.retrieve", return_value=mock_deployment)
48
+ return mock_deployment
49
+
50
+
51
+ @pytest.mark.parametrize(
52
+ ["GetDisplayClass", "expected_input_id"],
53
+ [
54
+ (_no_display_class, "6037747a-1d35-4094-b363-4369fc92c5d4"),
55
+ (_display_class_with_node_input_ids_by_name, "6037747a-1d35-4094-b363-4369fc92c5d4"),
56
+ (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "6037747a-1d35-4094-b363-4369fc92c5d4"),
57
+ ],
58
+ ids=[
59
+ "no_display_class",
60
+ "display_class_with_node_input_ids_by_name",
61
+ "display_class_with_node_input_ids_by_name_with_inputs_prefix",
62
+ ],
63
+ )
64
+ def test_serialize_node__prompt_inputs(GetDisplayClass, expected_input_id, mock_fetch_deployment):
65
+ # GIVEN a prompt node with inputs
66
+ class MyPromptDeploymentNode(PromptDeploymentNode):
67
+ deployment = "DEPLOYMENT"
68
+ prompt_inputs = {"foo": "bar"}
69
+ ml_model_fallbacks = None
70
+
71
+ # AND a workflow with the prompt node
72
+ class Workflow(BaseWorkflow):
73
+ graph = MyPromptDeploymentNode
74
+
75
+ # AND a display class
76
+ GetDisplayClass(MyPromptDeploymentNode)
77
+
78
+ # WHEN the workflow is serialized
79
+ workflow_display = get_workflow_display(workflow_class=Workflow)
80
+ serialized_workflow: dict = workflow_display.serialize()
81
+
82
+ # THEN the node should properly serialize the inputs
83
+ my_prompt_node = next(
84
+ node
85
+ for node in serialized_workflow["workflow_raw_data"]["nodes"]
86
+ if node["id"] == str(MyPromptDeploymentNode.__id__)
87
+ )
88
+
89
+ assert my_prompt_node["inputs"] == [
90
+ {
91
+ "id": expected_input_id,
92
+ "key": "foo",
93
+ "value": {
94
+ "rules": [
95
+ {
96
+ "type": "CONSTANT_VALUE",
97
+ "data": {
98
+ "type": "STRING",
99
+ "value": "bar",
100
+ },
101
+ }
102
+ ],
103
+ "combinator": "OR",
104
+ },
105
+ }
106
+ ]
@@ -0,0 +1,109 @@
1
+ import pytest
2
+ from datetime import datetime
3
+ from uuid import UUID, uuid4
4
+ from typing import Type
5
+
6
+ from vellum.workflows import BaseWorkflow
7
+ from vellum.workflows.nodes import SubworkflowDeploymentNode
8
+ from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
9
+
10
+
11
+ def _no_display_class(Node: Type[SubworkflowDeploymentNode]): # type: ignore
12
+ return None
13
+
14
+
15
+ def _display_class_with_node_input_ids_by_name(Node: Type[SubworkflowDeploymentNode]):
16
+ class SubworkflowNodeDisplay(SubworkflowDeploymentNode[Node]): # type: ignore[valid-type]
17
+ node_input_ids_by_name = {"foo": UUID("aff4f838-577e-44b9-ac5c-6d8213abbb9c")}
18
+
19
+ return SubworkflowNodeDisplay
20
+
21
+
22
+ def _display_class_with_node_input_ids_by_name_with_inputs_prefix(Node: Type[SubworkflowDeploymentNode]):
23
+ class SubworkflowNodeDisplay(SubworkflowDeploymentNode[Node]): # type: ignore[valid-type]
24
+ node_input_ids_by_name = {"subworkflow_inputs.foo": UUID("aff4f838-577e-44b9-ac5c-6d8213abbb9c")}
25
+
26
+ return SubworkflowNodeDisplay
27
+
28
+
29
+ @pytest.fixture
30
+ def mock_fetch_deployment(mocker):
31
+ # Create a mock deployment response
32
+ mock_deployment = mocker.Mock(
33
+ id="test-id",
34
+ name="test-deployment",
35
+ label="Test Deployment",
36
+ status="ACTIVE",
37
+ environment="DEVELOPMENT",
38
+ created=datetime.now(),
39
+ last_deployed_on=datetime.now(),
40
+ last_deployed_history_item_id=str(uuid4()),
41
+ input_variables=[],
42
+ output_variables=[],
43
+ description="Test deployment description",
44
+ )
45
+
46
+ # Patch the create_vellum_client function to return our mock client
47
+ mocker.patch(
48
+ "vellum.client.resources.workflow_deployments.client.WorkflowDeploymentsClient.retrieve",
49
+ return_value=mock_deployment,
50
+ )
51
+
52
+ return mock_deployment
53
+
54
+
55
+ @pytest.mark.parametrize(
56
+ ["GetDisplayClass", "expected_input_id"],
57
+ [
58
+ (_no_display_class, "aff4f838-577e-44b9-ac5c-6d8213abbb9c"),
59
+ (_display_class_with_node_input_ids_by_name, "aff4f838-577e-44b9-ac5c-6d8213abbb9c"),
60
+ (_display_class_with_node_input_ids_by_name_with_inputs_prefix, "aff4f838-577e-44b9-ac5c-6d8213abbb9c"),
61
+ ],
62
+ ids=[
63
+ "no_display_class",
64
+ "display_class_with_node_input_ids_by_name",
65
+ "display_class_with_node_input_ids_by_name_with_inputs_prefix",
66
+ ],
67
+ )
68
+ def test_serialize_node__subworkflow_inputs(GetDisplayClass, expected_input_id, mock_fetch_deployment):
69
+ # GIVEN a deployment subworkflow node with inputs
70
+ class MySubworkflowDeploymentNode(SubworkflowDeploymentNode):
71
+ deployment = "DEPLOYMENT"
72
+ subworkflow_inputs = {"foo": "bar"}
73
+
74
+ # AND a workflow with the subworkflow node
75
+ class Workflow(BaseWorkflow):
76
+ graph = MySubworkflowDeploymentNode
77
+
78
+ # AND a display class
79
+ GetDisplayClass(MySubworkflowDeploymentNode)
80
+
81
+ # WHEN the workflow is serialized
82
+ workflow_display = get_workflow_display(workflow_class=Workflow)
83
+ serialized_workflow: dict = workflow_display.serialize()
84
+
85
+ # THEN the node should properly serialize the inputs
86
+ my_subworkflow_node = next(
87
+ node
88
+ for node in serialized_workflow["workflow_raw_data"]["nodes"]
89
+ if node["id"] == str(MySubworkflowDeploymentNode.__id__)
90
+ )
91
+
92
+ assert my_subworkflow_node["inputs"] == [
93
+ {
94
+ "id": expected_input_id,
95
+ "key": "foo",
96
+ "value": {
97
+ "rules": [
98
+ {
99
+ "type": "CONSTANT_VALUE",
100
+ "data": {
101
+ "type": "STRING",
102
+ "value": "bar",
103
+ },
104
+ }
105
+ ],
106
+ "combinator": "OR",
107
+ },
108
+ }
109
+ ]
@@ -1,8 +1,7 @@
1
1
  import inspect
2
2
  from uuid import UUID
3
- from typing import Any, ClassVar, Generic, Optional, Tuple, Type, TypeVar, cast
3
+ from typing import Any, ClassVar, Generic, Optional, Tuple, Type, TypeVar
4
4
 
5
- from vellum.workflows.descriptors.base import BaseDescriptor
6
5
  from vellum.workflows.nodes.bases.base import BaseNode
7
6
  from vellum.workflows.nodes.core.try_node.node import TryNode
8
7
  from vellum.workflows.nodes.utils import ADORNMENT_MODULE_NAME
@@ -14,6 +13,7 @@ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_di
14
13
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
15
14
  from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
16
15
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
16
+ from vellum_ee.workflows.display.utils.expressions import serialize_value
17
17
 
18
18
  _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
19
19
 
@@ -37,7 +37,7 @@ class BaseTryNodeDisplay(BaseAdornmentNodeDisplay[_TryNodeType], Generic[_TryNod
37
37
  {
38
38
  "id": id,
39
39
  "name": attribute.name,
40
- "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
40
+ "value": serialize_value(display_context, attribute.instance),
41
41
  }
42
42
  )
43
43
 
@@ -43,6 +43,7 @@ def test_base_workflow_display__serialize_empty_workflow():
43
43
  "type": "ENTRYPOINT",
44
44
  }
45
45
  ],
46
+ "output_values": [],
46
47
  },
47
48
  }
48
49