vellum-ai 0.13.9__py3-none-any.whl → 0.13.10__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 (44) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/workflows/descriptors/utils.py +1 -1
  3. vellum/workflows/nodes/bases/base.py +1 -1
  4. vellum/workflows/nodes/displayable/api_node/node.py +4 -1
  5. vellum/workflows/nodes/displayable/bases/api_node/node.py +4 -1
  6. vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py +4 -1
  7. vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py +4 -0
  8. vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py +4 -0
  9. vellum/workflows/nodes/displayable/bases/search_node.py +4 -0
  10. vellum/workflows/nodes/displayable/code_execution_node/node.py +4 -1
  11. vellum/workflows/nodes/displayable/conditional_node/node.py +4 -0
  12. vellum/workflows/nodes/displayable/final_output_node/node.py +4 -0
  13. vellum/workflows/nodes/displayable/guardrail_node/node.py +4 -1
  14. vellum/workflows/nodes/displayable/inline_prompt_node/node.py +4 -0
  15. vellum/workflows/nodes/displayable/merge_node/node.py +3 -1
  16. vellum/workflows/nodes/displayable/note_node/node.py +4 -0
  17. vellum/workflows/nodes/displayable/prompt_deployment_node/node.py +4 -0
  18. vellum/workflows/nodes/displayable/search_node/node.py +4 -0
  19. vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py +4 -1
  20. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.10.dist-info}/METADATA +1 -1
  21. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.10.dist-info}/RECORD +44 -44
  22. vellum_ee/workflows/display/nodes/base_node_display.py +100 -0
  23. vellum_ee/workflows/display/nodes/vellum/base_node.py +34 -105
  24. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +2 -1
  25. vellum_ee/workflows/display/nodes/vellum/try_node.py +7 -0
  26. vellum_ee/workflows/display/nodes/vellum/utils.py +0 -69
  27. vellum_ee/workflows/display/tests/test_vellum_workflow_display.py +56 -0
  28. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py +3 -2
  29. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py +143 -26
  30. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py +10 -10
  31. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py +6 -6
  32. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py +32 -34
  33. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py +4 -4
  34. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py +5 -5
  35. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py +1 -1
  36. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +6 -3
  37. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py +6 -3
  38. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py +3 -3
  39. vellum_ee/workflows/display/utils/vellum.py +74 -4
  40. vellum_ee/workflows/display/workflows/base_workflow_display.py +6 -4
  41. vellum_ee/workflows/display/workflows/vellum_workflow_display.py +0 -2
  42. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.10.dist-info}/LICENSE +0 -0
  43. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.10.dist-info}/WHEEL +0 -0
  44. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.10.dist-info}/entry_points.txt +0 -0
@@ -16,15 +16,28 @@ from typing import (
16
16
  get_origin,
17
17
  )
18
18
 
19
+ from vellum.workflows.descriptors.base import BaseDescriptor
20
+ from vellum.workflows.expressions.between import BetweenExpression
21
+ from vellum.workflows.expressions.is_nil import IsNilExpression
22
+ from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
23
+ from vellum.workflows.expressions.is_not_null import IsNotNullExpression
24
+ from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
25
+ from vellum.workflows.expressions.is_null import IsNullExpression
26
+ from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
27
+ from vellum.workflows.expressions.not_between import NotBetweenExpression
19
28
  from vellum.workflows.nodes.bases.base import BaseNode
20
29
  from vellum.workflows.ports import Port
21
30
  from vellum.workflows.references import OutputReference
31
+ from vellum.workflows.references.execution_count import ExecutionCountReference
32
+ from vellum.workflows.references.vellum_secret import VellumSecretReference
33
+ from vellum.workflows.references.workflow_input import WorkflowInputReference
22
34
  from vellum.workflows.types.core import JsonObject
23
35
  from vellum.workflows.types.generics import NodeType
24
36
  from vellum.workflows.types.utils import get_original_base
25
37
  from vellum.workflows.utils.names import pascal_to_title_case
26
38
  from vellum.workflows.utils.uuids import uuid4_from_hash
27
39
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay, PortDisplay, PortDisplayOverrides
40
+ from vellum_ee.workflows.display.utils.vellum import convert_descriptor_to_operator, primitive_to_vellum_value
28
41
  from vellum_ee.workflows.display.vellum import CodeResourceDefinition
29
42
 
30
43
  if TYPE_CHECKING:
@@ -76,6 +89,9 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
76
89
  )
77
90
  return node_definition
78
91
 
92
+ def get_trigger_id(self) -> UUID:
93
+ return uuid4_from_hash(f"{self.node_id}|trigger")
94
+
79
95
  def get_node_output_display(self, output: OutputReference) -> Tuple[Type[BaseNode], NodeOutputDisplay]:
80
96
  explicit_display = self.output_display.get(output)
81
97
  if explicit_display:
@@ -174,3 +190,87 @@ class BaseNodeDisplay(Generic[NodeType], metaclass=BaseNodeDisplayMeta):
174
190
 
175
191
  node_class = cls.infer_node_class()
176
192
  cls._node_display_registry[node_class] = cls
193
+
194
+ def serialize_condition(self, display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
195
+ if isinstance(
196
+ condition,
197
+ (
198
+ IsNullExpression,
199
+ IsNotNullExpression,
200
+ IsNilExpression,
201
+ IsNotNilExpression,
202
+ IsUndefinedExpression,
203
+ IsNotUndefinedExpression,
204
+ ),
205
+ ):
206
+ lhs = self.serialize_value(display_context, condition._expression)
207
+ return {
208
+ "type": "UNARY_EXPRESSION",
209
+ "lhs": lhs,
210
+ "operator": convert_descriptor_to_operator(condition),
211
+ }
212
+ elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
213
+ base = self.serialize_value(display_context, condition._value)
214
+ lhs = self.serialize_value(display_context, condition._start)
215
+ rhs = self.serialize_value(display_context, condition._end)
216
+
217
+ return {
218
+ "type": "TERNARY_EXPRESSION",
219
+ "base": base,
220
+ "operator": convert_descriptor_to_operator(condition),
221
+ "lhs": lhs,
222
+ "rhs": rhs,
223
+ }
224
+ else:
225
+ lhs = self.serialize_value(display_context, condition._lhs) # type: ignore[attr-defined]
226
+ rhs = self.serialize_value(display_context, condition._rhs) # type: ignore[attr-defined]
227
+
228
+ return {
229
+ "type": "BINARY_EXPRESSION",
230
+ "lhs": lhs,
231
+ "operator": convert_descriptor_to_operator(condition),
232
+ "rhs": rhs,
233
+ }
234
+
235
+ def serialize_value(self, display_context: "WorkflowDisplayContext", value: BaseDescriptor) -> JsonObject:
236
+ if isinstance(value, WorkflowInputReference):
237
+ workflow_input_display = display_context.global_workflow_input_displays[value]
238
+ return {
239
+ "type": "WORKFLOW_INPUT",
240
+ "input_variable_id": str(workflow_input_display.id),
241
+ }
242
+
243
+ if isinstance(value, OutputReference):
244
+ upstream_node, output_display = display_context.global_node_output_displays[value]
245
+ upstream_node_display = display_context.global_node_displays[upstream_node]
246
+
247
+ return {
248
+ "type": "NODE_OUTPUT",
249
+ "node_id": str(upstream_node_display.node_id),
250
+ "node_output_id": str(output_display.id),
251
+ }
252
+
253
+ if isinstance(value, VellumSecretReference):
254
+ return {
255
+ "type": "VELLUM_SECRET",
256
+ "vellum_secret_name": value.name,
257
+ }
258
+
259
+ if isinstance(value, ExecutionCountReference):
260
+ node_class_display = display_context.global_node_displays[value.node_class]
261
+
262
+ return {
263
+ "type": "EXECUTION_COUNTER",
264
+ "node_id": str(node_class_display.node_id),
265
+ }
266
+
267
+ if not isinstance(value, BaseDescriptor):
268
+ vellum_value = primitive_to_vellum_value(value)
269
+ return {
270
+ "type": "CONSTANT_VALUE",
271
+ "value": vellum_value.dict(),
272
+ }
273
+
274
+ # If it's not any of the references we know about,
275
+ # then try to serialize it as a nested value
276
+ return self.serialize_condition(display_context, value)
@@ -1,41 +1,37 @@
1
- from typing import Any, Generic, TypeVar, cast
1
+ import inspect
2
+ from uuid import UUID
3
+ from typing import Any, Generic, Optional, TypeVar, cast
2
4
 
3
5
  from vellum.workflows.constants import UNDEF
4
6
  from vellum.workflows.descriptors.base import BaseDescriptor
5
- from vellum.workflows.expressions.between import BetweenExpression
6
- from vellum.workflows.expressions.is_nil import IsNilExpression
7
- from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
8
- from vellum.workflows.expressions.is_not_null import IsNotNullExpression
9
- from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
10
- from vellum.workflows.expressions.is_null import IsNullExpression
11
- from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
12
- from vellum.workflows.expressions.not_between import NotBetweenExpression
13
7
  from vellum.workflows.nodes.bases.base import BaseNode
14
- from vellum.workflows.references.execution_count import ExecutionCountReference
15
- from vellum.workflows.references.output import OutputReference
16
- from vellum.workflows.references.vellum_secret import VellumSecretReference
17
- from vellum.workflows.references.workflow_input import WorkflowInputReference
8
+ from vellum.workflows.nodes.utils import get_wrapped_node
18
9
  from vellum.workflows.types.core import JsonArray, JsonObject
19
10
  from vellum.workflows.utils.uuids import uuid4_from_hash
20
11
  from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
12
+ from vellum.workflows.workflows.base import BaseWorkflow
21
13
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
22
- from vellum_ee.workflows.display.nodes.vellum.utils import convert_descriptor_to_operator
14
+ from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
23
15
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
24
- from vellum_ee.workflows.display.utils.vellum import primitive_to_vellum_value
25
16
  from vellum_ee.workflows.display.vellum import GenericNodeDisplayData
26
17
 
27
18
  _BaseNodeType = TypeVar("_BaseNodeType", bound=BaseNode)
28
19
 
29
20
 
30
21
  class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeType]):
31
- def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
22
+ def serialize(
23
+ self, display_context: WorkflowDisplayContext, adornments: Optional[JsonArray] = None, **kwargs: Any
24
+ ) -> JsonObject:
32
25
  node = self._node
33
26
  node_id = self.node_id
34
27
 
35
28
  attributes: JsonArray = []
36
29
  for attribute in node:
37
- id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
30
+ if inspect.isclass(attribute.instance) and issubclass(attribute.instance, BaseWorkflow):
31
+ # We don't need to serialize generic node attributes containing a subworkflow
32
+ continue
38
33
 
34
+ id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
39
35
  attributes.append(
40
36
  {
41
37
  "id": id,
@@ -44,9 +40,23 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
44
40
  }
45
41
  )
46
42
 
43
+ wrapped_node = get_wrapped_node(node)
44
+ if wrapped_node is not None:
45
+ display_class = get_node_display_class(BaseNodeDisplay, wrapped_node)
46
+
47
+ adornment: JsonObject = {
48
+ "id": str(node_id),
49
+ "label": node.__qualname__,
50
+ "base": self.get_base().dict(),
51
+ "attributes": attributes,
52
+ }
53
+
54
+ existing_adornments = adornments if adornments is not None else []
55
+ return display_class().serialize(display_context, adornments=existing_adornments + [adornment])
56
+
47
57
  ports: JsonArray = []
48
- for idx, port in enumerate(node.Ports):
49
- id = str(uuid4_from_hash(f"{node_id}|{idx}"))
58
+ for port in node.Ports:
59
+ id = str(self.get_node_port_display(port).id)
50
60
 
51
61
  if port._condition_type:
52
62
  ports.append(
@@ -94,99 +104,18 @@ class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeTyp
94
104
  "base": self.get_base().dict(),
95
105
  "definition": self.get_definition().dict(),
96
106
  "trigger": {
97
- "id": str(uuid4_from_hash(f"{node_id}|trigger")),
107
+ "id": str(self.get_trigger_id()),
98
108
  "merge_behavior": node.Trigger.merge_behavior.value,
99
109
  },
100
110
  "ports": ports,
101
- "adornments": None,
111
+ "adornments": adornments,
102
112
  "attributes": attributes,
103
113
  "outputs": outputs,
104
114
  }
105
115
 
116
+ def get_target_handle_id(self) -> UUID:
117
+ return self.get_trigger_id()
118
+
106
119
  def get_generic_node_display_data(self) -> GenericNodeDisplayData:
107
120
  explicit_value = self._get_explicit_node_display_attr("display_data", GenericNodeDisplayData)
108
121
  return explicit_value if explicit_value else GenericNodeDisplayData()
109
-
110
- def serialize_condition(self, display_context: WorkflowDisplayContext, condition: BaseDescriptor) -> JsonObject:
111
- if isinstance(
112
- condition,
113
- (
114
- IsNullExpression,
115
- IsNotNullExpression,
116
- IsNilExpression,
117
- IsNotNilExpression,
118
- IsUndefinedExpression,
119
- IsNotUndefinedExpression,
120
- ),
121
- ):
122
- lhs = self.serialize_value(display_context, condition._expression)
123
- return {
124
- "type": "UNARY_EXPRESSION",
125
- "lhs": lhs,
126
- "operator": convert_descriptor_to_operator(condition),
127
- }
128
- elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
129
- base = self.serialize_value(display_context, condition._value)
130
- lhs = self.serialize_value(display_context, condition._start)
131
- rhs = self.serialize_value(display_context, condition._end)
132
-
133
- return {
134
- "type": "TERNARY_EXPRESSION",
135
- "base": base,
136
- "operator": convert_descriptor_to_operator(condition),
137
- "lhs": lhs,
138
- "rhs": rhs,
139
- }
140
- else:
141
- lhs = self.serialize_value(display_context, condition._lhs) # type: ignore[attr-defined]
142
- rhs = self.serialize_value(display_context, condition._rhs) # type: ignore[attr-defined]
143
-
144
- return {
145
- "type": "BINARY_EXPRESSION",
146
- "lhs": lhs,
147
- "operator": convert_descriptor_to_operator(condition),
148
- "rhs": rhs,
149
- }
150
-
151
- def serialize_value(self, display_context: WorkflowDisplayContext, value: BaseDescriptor) -> JsonObject:
152
- if isinstance(value, WorkflowInputReference):
153
- workflow_input_display = display_context.global_workflow_input_displays[value]
154
- return {
155
- "type": "WORKFLOW_INPUT",
156
- "input_variable_id": str(workflow_input_display.id),
157
- }
158
-
159
- if isinstance(value, OutputReference):
160
- upstream_node, output_display = display_context.global_node_output_displays[value]
161
- upstream_node_display = display_context.global_node_displays[upstream_node]
162
-
163
- return {
164
- "type": "NODE_OUTPUT",
165
- "node_id": str(upstream_node_display.node_id),
166
- "node_output_id": str(output_display.id),
167
- }
168
-
169
- if isinstance(value, VellumSecretReference):
170
- return {
171
- "type": "VELLUM_SECRET",
172
- "vellum_secret_name": value.name,
173
- }
174
-
175
- if isinstance(value, ExecutionCountReference):
176
- node_class_display = display_context.global_node_displays[value.node_class]
177
-
178
- return {
179
- "type": "EXECUTION_COUNTER",
180
- "node_id": str(node_class_display.node_id),
181
- }
182
-
183
- if not isinstance(value, BaseDescriptor):
184
- vellum_value = primitive_to_vellum_value(value)
185
- return {
186
- "type": "CONSTANT_VALUE",
187
- "value": vellum_value.dict(),
188
- }
189
-
190
- # If it's not any of the references we know about,
191
- # then try to serialize it as a nested value
192
- return self.serialize_condition(display_context, value)
@@ -17,8 +17,9 @@ from vellum.workflows.nodes.displayable import ConditionalNode
17
17
  from vellum.workflows.types.core import ConditionType, JsonObject
18
18
  from vellum.workflows.utils.uuids import uuid4_from_hash
19
19
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
20
- from vellum_ee.workflows.display.nodes.vellum.utils import convert_descriptor_to_operator, create_node_input
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
23
  from vellum_ee.workflows.display.vellum import NodeInput
23
24
 
24
25
  _ConditionalNodeType = TypeVar("_ConditionalNodeType", bound=ConditionalNode)
@@ -13,6 +13,7 @@ from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeV
13
13
  from vellum_ee.workflows.display.nodes.get_node_display_class import get_node_display_class
14
14
  from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
15
15
  from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
16
+ from vellum_ee.workflows.display.nodes.vellum.base_node import BaseNodeDisplay as GenericBaseNodeDisplay
16
17
  from vellum_ee.workflows.display.types import WorkflowDisplayContext
17
18
 
18
19
  _TryNodeType = TypeVar("_TryNodeType", bound=TryNode)
@@ -33,6 +34,12 @@ class BaseTryNodeDisplay(BaseNodeVellumDisplay[_TryNodeType], Generic[_TryNodeTy
33
34
  )
34
35
 
35
36
  inner_node = subworkflow.graph
37
+ elif inner_node.__bases__[0] is BaseNode:
38
+ # If the wrapped node is a generic node, we let generic node do adornment handling
39
+ class TryBaseNodeDisplay(GenericBaseNodeDisplay[node]): # type: ignore[valid-type]
40
+ pass
41
+
42
+ return TryBaseNodeDisplay().serialize(display_context)
36
43
 
37
44
  # We need the node display class of the underlying node because
38
45
  # it contains the logic for serializing the node and potential display overrides
@@ -2,31 +2,7 @@ from uuid import UUID
2
2
  from typing import Any, List, Optional, Type, Union, cast
3
3
 
4
4
  from vellum.workflows.descriptors.base import BaseDescriptor
5
- from vellum.workflows.expressions.and_ import AndExpression
6
- from vellum.workflows.expressions.begins_with import BeginsWithExpression
7
- from vellum.workflows.expressions.between import BetweenExpression
8
5
  from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
9
- from vellum.workflows.expressions.contains import ContainsExpression
10
- from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
11
- from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
12
- from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
13
- from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
14
- from vellum.workflows.expressions.ends_with import EndsWithExpression
15
- from vellum.workflows.expressions.equals import EqualsExpression
16
- from vellum.workflows.expressions.greater_than import GreaterThanExpression
17
- from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
18
- from vellum.workflows.expressions.in_ import InExpression
19
- from vellum.workflows.expressions.is_nil import IsNilExpression
20
- from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
21
- from vellum.workflows.expressions.is_not_null import IsNotNullExpression
22
- from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
23
- from vellum.workflows.expressions.is_null import IsNullExpression
24
- from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
25
- from vellum.workflows.expressions.less_than import LessThanExpression
26
- from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
27
- from vellum.workflows.expressions.not_between import NotBetweenExpression
28
- from vellum.workflows.expressions.not_in import NotInExpression
29
- from vellum.workflows.expressions.or_ import OrExpression
30
6
  from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
31
7
  from vellum.workflows.references import NodeReference
32
8
  from vellum.workflows.utils.uuids import uuid4_from_hash
@@ -124,48 +100,3 @@ def create_pointer(
124
100
  return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_variable_value)
125
101
  else:
126
102
  raise ValueError(f"Pointer type {pointer_type} not supported")
127
-
128
-
129
- def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> str:
130
- if isinstance(descriptor, EqualsExpression):
131
- return "="
132
- elif isinstance(descriptor, DoesNotEqualExpression):
133
- return "!="
134
- elif isinstance(descriptor, LessThanExpression):
135
- return "<"
136
- elif isinstance(descriptor, GreaterThanExpression):
137
- return ">"
138
- elif isinstance(descriptor, LessThanOrEqualToExpression):
139
- return "<="
140
- elif isinstance(descriptor, GreaterThanOrEqualToExpression):
141
- return ">="
142
- elif isinstance(descriptor, ContainsExpression):
143
- return "contains"
144
- elif isinstance(descriptor, BeginsWithExpression):
145
- return "beginsWith"
146
- elif isinstance(descriptor, EndsWithExpression):
147
- return "endsWith"
148
- elif isinstance(descriptor, DoesNotContainExpression):
149
- return "doesNotContain"
150
- elif isinstance(descriptor, DoesNotBeginWithExpression):
151
- return "doesNotBeginWith"
152
- elif isinstance(descriptor, DoesNotEndWithExpression):
153
- return "doesNotEndWith"
154
- elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
155
- return "null"
156
- elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
157
- return "notNull"
158
- elif isinstance(descriptor, InExpression):
159
- return "in"
160
- elif isinstance(descriptor, NotInExpression):
161
- return "notIn"
162
- elif isinstance(descriptor, BetweenExpression):
163
- return "between"
164
- elif isinstance(descriptor, NotBetweenExpression):
165
- return "notBetween"
166
- elif isinstance(descriptor, AndExpression):
167
- return "and"
168
- elif isinstance(descriptor, OrExpression):
169
- return "or"
170
- else:
171
- raise ValueError(f"Unsupported descriptor type: {descriptor}")
@@ -88,3 +88,59 @@ def test_vellum_workflow_display__serialize_input_variables_with_capitalized_var
88
88
  "extensions": {"color": None},
89
89
  }
90
90
  ]
91
+
92
+
93
+ def test_vellum_workflow_display_serialize_valid_handle_ids_for_base_nodes():
94
+ # GIVEN a workflow between two base nodes
95
+ class StartNode(BaseNode):
96
+ pass
97
+
98
+ class EndNode(BaseNode):
99
+ class Outputs(BaseNode.Outputs):
100
+ hello = "world"
101
+
102
+ class Workflow(BaseWorkflow):
103
+ graph = StartNode >> EndNode
104
+
105
+ class Outputs(BaseWorkflow.Outputs):
106
+ final_value = EndNode.Outputs.hello
107
+
108
+ # AND a display class for this workflow
109
+ workflow_display = get_workflow_display(
110
+ base_display_class=VellumWorkflowDisplay,
111
+ workflow_class=Workflow,
112
+ )
113
+
114
+ # WHEN we serialize the workflow
115
+ exec_config = workflow_display.serialize()
116
+
117
+ # THEN the serialized workflow handle ids are valid
118
+ raw_data = exec_config.get("workflow_raw_data")
119
+ assert isinstance(raw_data, dict)
120
+ nodes = raw_data.get("nodes")
121
+ edges = raw_data.get("edges")
122
+
123
+ assert isinstance(nodes, list)
124
+ assert isinstance(edges, list)
125
+
126
+ edge_source_handle_ids = {edge.get("source_handle_id") for edge in edges if isinstance(edge, dict)}
127
+ edge_target_handle_ids = {edge.get("target_handle_id") for edge in edges if isinstance(edge, dict)}
128
+
129
+ for node in nodes:
130
+ assert isinstance(node, dict)
131
+
132
+ if node["type"] in {"ENTRYPOINT", "TERMINAL"}:
133
+ continue
134
+
135
+ ports = node.get("ports")
136
+ assert isinstance(ports, list)
137
+ for port in ports:
138
+ assert isinstance(port, dict)
139
+ assert (
140
+ port["id"] in edge_source_handle_ids
141
+ ), f"Port {port['id']} from node {node['label']} not found in edge source handle ids"
142
+
143
+ assert isinstance(node["trigger"], dict)
144
+ assert (
145
+ node["trigger"]["id"] in edge_target_handle_ids
146
+ ), f"Trigger {node['trigger']['id']} from node {node['label']} not found in edge target handle ids"
@@ -1,6 +1,6 @@
1
1
  import pytest
2
2
  from uuid import uuid4
3
- from typing import Dict, Tuple, Type
3
+ from typing import Any, Dict, Tuple, Type
4
4
 
5
5
  from vellum.workflows.nodes.bases.base import BaseNode
6
6
  from vellum.workflows.references.output import OutputReference
@@ -20,11 +20,12 @@ from vellum_ee.workflows.display.workflows.vellum_workflow_display import Vellum
20
20
  def serialize_node():
21
21
  def _serialize_node(
22
22
  node_class: Type[NodeType],
23
+ base_class: type[BaseNodeDisplay[Any]] = BaseNodeDisplay,
23
24
  global_workflow_input_displays: Dict[WorkflowInputReference, WorkflowInputsDisplayType] = {},
24
25
  global_node_displays: Dict[Type[BaseNode], NodeDisplayType] = {},
25
26
  global_node_output_displays: Dict[OutputReference, Tuple[Type[BaseNode], NodeOutputDisplay]] = {},
26
27
  ) -> JsonObject:
27
- node_display_class = get_node_display_class(BaseNodeDisplay, node_class)
28
+ node_display_class = get_node_display_class(base_class, node_class)
28
29
  node_display = node_display_class()
29
30
 
30
31
  context: WorkflowDisplayContext = WorkflowDisplayContext(