vellum-ai 0.13.9__py3-none-any.whl → 0.13.11__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 (60) hide show
  1. vellum/client/core/client_wrapper.py +1 -1
  2. vellum/workflows/descriptors/utils.py +1 -1
  3. vellum/workflows/errors/types.py +21 -0
  4. vellum/workflows/nodes/bases/base.py +1 -1
  5. vellum/workflows/nodes/displayable/api_node/node.py +4 -1
  6. vellum/workflows/nodes/displayable/bases/api_node/node.py +4 -1
  7. vellum/workflows/nodes/displayable/bases/base_prompt_node/node.py +18 -2
  8. vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py +4 -0
  9. vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py +4 -0
  10. vellum/workflows/nodes/displayable/bases/search_node.py +4 -0
  11. vellum/workflows/nodes/displayable/bases/tests/test_utils.py +18 -0
  12. vellum/workflows/nodes/displayable/bases/utils.py +8 -1
  13. vellum/workflows/nodes/displayable/code_execution_node/node.py +4 -1
  14. vellum/workflows/nodes/displayable/conditional_node/node.py +4 -0
  15. vellum/workflows/nodes/displayable/final_output_node/node.py +4 -0
  16. vellum/workflows/nodes/displayable/guardrail_node/node.py +4 -1
  17. vellum/workflows/nodes/displayable/inline_prompt_node/node.py +4 -0
  18. vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py +55 -0
  19. vellum/workflows/nodes/displayable/merge_node/node.py +3 -1
  20. vellum/workflows/nodes/displayable/note_node/node.py +4 -0
  21. vellum/workflows/nodes/displayable/prompt_deployment_node/node.py +4 -0
  22. vellum/workflows/nodes/displayable/search_node/node.py +4 -0
  23. vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py +4 -1
  24. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.11.dist-info}/METADATA +1 -1
  25. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.11.dist-info}/RECORD +59 -59
  26. vellum_cli/__init__.py +9 -1
  27. vellum_cli/config.py +29 -1
  28. vellum_cli/push.py +24 -3
  29. vellum_cli/tests/conftest.py +3 -0
  30. vellum_cli/tests/test_pull.py +6 -0
  31. vellum_cli/tests/test_push.py +88 -1
  32. vellum_ee/workflows/display/nodes/base_node_display.py +207 -3
  33. vellum_ee/workflows/display/nodes/base_node_vellum_display.py +16 -1
  34. vellum_ee/workflows/display/nodes/get_node_display_class.py +6 -4
  35. vellum_ee/workflows/display/nodes/vellum/__init__.py +0 -2
  36. vellum_ee/workflows/display/nodes/vellum/conditional_node.py +2 -1
  37. vellum_ee/workflows/display/nodes/vellum/error_node.py +9 -3
  38. vellum_ee/workflows/display/nodes/vellum/tests/test_error_node.py +44 -0
  39. vellum_ee/workflows/display/nodes/vellum/try_node.py +8 -2
  40. vellum_ee/workflows/display/nodes/vellum/utils.py +0 -69
  41. vellum_ee/workflows/display/tests/test_vellum_workflow_display.py +56 -0
  42. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/conftest.py +4 -3
  43. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_adornments_serialization.py +146 -26
  44. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py +11 -11
  45. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_outputs_serialization.py +7 -7
  46. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_ports_serialization.py +33 -35
  47. vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py +4 -4
  48. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_conditional_node_serialization.py +5 -5
  49. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_error_node_serialization.py +15 -1
  50. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_generic_node_serialization.py +1 -1
  51. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py +6 -3
  52. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py +6 -3
  53. vellum_ee/workflows/display/tests/workflow_serialization/test_basic_merge_node_serialization.py +3 -3
  54. vellum_ee/workflows/display/utils/vellum.py +74 -4
  55. vellum_ee/workflows/display/workflows/base_workflow_display.py +6 -4
  56. vellum_ee/workflows/display/workflows/vellum_workflow_display.py +26 -14
  57. vellum_ee/workflows/display/nodes/vellum/base_node.py +0 -192
  58. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.11.dist-info}/LICENSE +0 -0
  59. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.11.dist-info}/WHEEL +0 -0
  60. {vellum_ai-0.13.9.dist-info → vellum_ai-0.13.11.dist-info}/entry_points.txt +0 -0
@@ -88,7 +88,21 @@ def test_serialize_workflow():
88
88
  {
89
89
  "id": "5cf9c5e3-0eae-4daf-8d73-8b9536258eb9",
90
90
  "type": "ERROR",
91
- "inputs": [],
91
+ "inputs": [
92
+ {
93
+ "id": "690d825f-6ffd-493e-8141-c86d384e6150",
94
+ "key": "error_source_input_id",
95
+ "value": {
96
+ "rules": [
97
+ {
98
+ "type": "CONSTANT_VALUE",
99
+ "data": {"type": "STRING", "value": "Input threshold was too low"},
100
+ }
101
+ ],
102
+ "combinator": "OR",
103
+ },
104
+ }
105
+ ],
92
106
  "data": {
93
107
  "name": "error-node",
94
108
  "label": "Fail Node",
@@ -123,7 +123,7 @@ def test_serialize_workflow(vellum_client):
123
123
  "source_node_id": "f1e4678f-c470-400b-a40e-c8922cc99a86",
124
124
  "source_handle_id": "40201804-8beb-43ad-8873-a027759512f1",
125
125
  "target_node_id": "c2ed23f7-f6cb-4a56-a91c-2e5f9d8fda7f",
126
- "target_handle_id": "b7bfb298-959a-4d2b-8b85-bbd0d2522703",
126
+ "target_handle_id": "9d3a1b3d-4a38-4f2e-bbf1-dd8be152bce8",
127
127
  "type": "DEFAULT",
128
128
  },
129
129
  {
@@ -140,9 +140,12 @@ def test_serialize_workflow():
140
140
  "name": "StartNode",
141
141
  "module": ["tests", "workflows", "basic_inline_subworkflow", "workflow"],
142
142
  },
143
- "trigger": {"id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135", "merge_behavior": "AWAIT_ANY"},
143
+ "trigger": {
144
+ "id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135",
145
+ "merge_behavior": "AWAIT_ATTRIBUTES",
146
+ },
144
147
  "ports": [
145
- {"id": "4a62dea3-6a4a-4390-88d8-4af5d2c81474", "type": "DEFAULT", "name": "default"}
148
+ {"id": "1e739e86-a285-4438-9725-a152c15a63e3", "type": "DEFAULT", "name": "default"}
146
149
  ],
147
150
  "adornments": None,
148
151
  "attributes": [
@@ -271,7 +274,7 @@ def test_serialize_workflow():
271
274
  "source_node_id": "afa49a0f-db35-4552-9217-5b8f237e84bc",
272
275
  "source_handle_id": "9914a6a0-9a99-430d-8ddd-f7c13847fe1a",
273
276
  "target_node_id": "1381c078-efa2-4255-89a1-7b4cb742c7fc",
274
- "target_handle_id": "6492efcf-4437-4af1-9ad7-269795ccb27a",
277
+ "target_handle_id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135",
275
278
  "type": "DEFAULT",
276
279
  },
277
280
  {
@@ -129,9 +129,12 @@ def test_serialize_workflow():
129
129
  "name": "Iteration",
130
130
  "module": ["tests", "workflows", "basic_map_node", "workflow"],
131
131
  },
132
- "trigger": {"id": "01324747-9bc0-4ecd-a8ab-40dca5a94e2e", "merge_behavior": "AWAIT_ANY"},
132
+ "trigger": {
133
+ "id": "01324747-9bc0-4ecd-a8ab-40dca5a94e2e",
134
+ "merge_behavior": "AWAIT_ATTRIBUTES",
135
+ },
133
136
  "ports": [
134
- {"id": "36791877-95b0-4390-a794-10accd8f548e", "type": "DEFAULT", "name": "default"}
137
+ {"id": "71ada606-d791-4a59-a252-0795c5faeeaf", "type": "DEFAULT", "name": "default"}
135
138
  ],
136
139
  "adornments": None,
137
140
  "attributes": [
@@ -211,7 +214,7 @@ def test_serialize_workflow():
211
214
  "source_node_id": "ff9bfe6e-839d-4d40-b8fc-313b3bbd0ab0",
212
215
  "source_handle_id": "520d3616-8369-4e79-9da5-3febae299c2a",
213
216
  "target_node_id": "baf6d316-dc75-41e8-96c0-015aede96309",
214
- "target_handle_id": "551d5528-f4e1-42ea-bde0-9de4b4968253",
217
+ "target_handle_id": "01324747-9bc0-4ecd-a8ab-40dca5a94e2e",
215
218
  "type": "DEFAULT",
216
219
  },
217
220
  {
@@ -145,7 +145,7 @@ def test_serialize_workflow__await_all():
145
145
  "source_node_id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
146
146
  "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
147
147
  "target_node_id": "59243c65-053f-4ea6-9157-3f3edb1477bf",
148
- "target_handle_id": "e622fe61-3bca-4aff-86e1-25dad7bdf9d4",
148
+ "target_handle_id": "75293c10-d6d8-4fc0-afae-ba6437af982e",
149
149
  "type": "DEFAULT",
150
150
  },
151
151
  {
@@ -153,7 +153,7 @@ def test_serialize_workflow__await_all():
153
153
  "source_node_id": "dc8aecd0-49ba-4464-a45f-29d3bfd686e4",
154
154
  "source_handle_id": "017d40f5-8326-4e42-a409-b08995defaa8",
155
155
  "target_node_id": "127ef456-91bc-43c6-bd8b-1772db5e3cb5",
156
- "target_handle_id": "e5cc41cb-71db-43ec-b3f0-c78706af3351",
156
+ "target_handle_id": "10c7392f-7dc0-4273-8380-ddc3497e6072",
157
157
  "type": "DEFAULT",
158
158
  },
159
159
  {
@@ -177,7 +177,7 @@ def test_serialize_workflow__await_all():
177
177
  "source_node_id": "37c10e8a-771b-432b-a767-31f5007851f0",
178
178
  "source_handle_id": "3bbc469f-0fb0-4b3d-a28b-746fefec2818",
179
179
  "target_node_id": "634f0202-9ea9-4c62-b152-1a58c595cffb",
180
- "target_handle_id": "acd48f48-54fb-4b2b-ab37-96d336f6dfb3",
180
+ "target_handle_id": "c0c993cd-1370-480b-b175-59590735aa46",
181
181
  "type": "DEFAULT",
182
182
  },
183
183
  {
@@ -1,7 +1,32 @@
1
- from typing import Any, TypeVar
1
+ from typing import TYPE_CHECKING, Any
2
2
 
3
+ from vellum.client.types.logical_operator import LogicalOperator
3
4
  from vellum.client.types.vellum_variable_type import VellumVariableType
4
5
  from vellum.workflows.descriptors.base import BaseDescriptor
6
+ from vellum.workflows.expressions.and_ import AndExpression
7
+ from vellum.workflows.expressions.begins_with import BeginsWithExpression
8
+ from vellum.workflows.expressions.between import BetweenExpression
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
5
30
  from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
6
31
  from vellum.workflows.references import OutputReference, WorkflowInputReference
7
32
  from vellum.workflows.references.execution_count import ExecutionCountReference
@@ -9,7 +34,6 @@ from vellum.workflows.references.node import NodeReference
9
34
  from vellum.workflows.references.vellum_secret import VellumSecretReference
10
35
  from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
11
36
  from vellum.workflows.vellum_client import create_vellum_client
12
- from vellum_ee.workflows.display.types import WorkflowDisplayContext
13
37
  from vellum_ee.workflows.display.vellum import (
14
38
  ConstantValuePointer,
15
39
  ExecutionCounterData,
@@ -23,7 +47,8 @@ from vellum_ee.workflows.display.vellum import (
23
47
  WorkspaceSecretPointer,
24
48
  )
25
49
 
26
- _T = TypeVar("_T")
50
+ if TYPE_CHECKING:
51
+ from vellum_ee.workflows.display.types import WorkflowDisplayContext
27
52
 
28
53
 
29
54
  def infer_vellum_variable_type(value: Any) -> VellumVariableType:
@@ -47,7 +72,7 @@ def infer_vellum_variable_type(value: Any) -> VellumVariableType:
47
72
 
48
73
 
49
74
  def create_node_input_value_pointer_rule(
50
- value: Any, display_context: WorkflowDisplayContext
75
+ value: Any, display_context: "WorkflowDisplayContext"
51
76
  ) -> NodeInputValuePointerRule:
52
77
  if isinstance(value, OutputReference):
53
78
  upstream_node, output_display = display_context.global_node_output_displays[value]
@@ -82,3 +107,48 @@ def create_node_input_value_pointer_rule(
82
107
  return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
83
108
 
84
109
  raise ValueError(f"Unsupported descriptor type: {value.__class__.__name__}")
110
+
111
+
112
+ def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator:
113
+ if isinstance(descriptor, EqualsExpression):
114
+ return "="
115
+ elif isinstance(descriptor, DoesNotEqualExpression):
116
+ return "!="
117
+ elif isinstance(descriptor, LessThanExpression):
118
+ return "<"
119
+ elif isinstance(descriptor, GreaterThanExpression):
120
+ return ">"
121
+ elif isinstance(descriptor, LessThanOrEqualToExpression):
122
+ return "<="
123
+ elif isinstance(descriptor, GreaterThanOrEqualToExpression):
124
+ return ">="
125
+ elif isinstance(descriptor, ContainsExpression):
126
+ return "contains"
127
+ elif isinstance(descriptor, BeginsWithExpression):
128
+ return "beginsWith"
129
+ elif isinstance(descriptor, EndsWithExpression):
130
+ return "endsWith"
131
+ elif isinstance(descriptor, DoesNotContainExpression):
132
+ return "doesNotContain"
133
+ elif isinstance(descriptor, DoesNotBeginWithExpression):
134
+ return "doesNotBeginWith"
135
+ elif isinstance(descriptor, DoesNotEndWithExpression):
136
+ return "doesNotEndWith"
137
+ elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
138
+ return "null"
139
+ elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
140
+ return "notNull"
141
+ elif isinstance(descriptor, InExpression):
142
+ return "in"
143
+ elif isinstance(descriptor, NotInExpression):
144
+ return "notIn"
145
+ elif isinstance(descriptor, BetweenExpression):
146
+ return "between"
147
+ elif isinstance(descriptor, NotBetweenExpression):
148
+ return "notBetween"
149
+ elif isinstance(descriptor, AndExpression):
150
+ return "and"
151
+ elif isinstance(descriptor, OrExpression):
152
+ return "or"
153
+ else:
154
+ raise ValueError(f"Unsupported descriptor type: {descriptor}")
@@ -207,11 +207,13 @@ class BaseWorkflowDisplay(
207
207
  # TODO: We should still serialize nodes that are in the workflow's directory but aren't used in the graph.
208
208
  # https://app.shortcut.com/vellum/story/5394
209
209
  for node in self._workflow.get_nodes():
210
- if node in global_node_displays:
211
- continue
212
210
  node_display = self._get_node_display(node)
213
- node_displays[node] = node_display
214
- global_node_displays[node] = node_display
211
+
212
+ if node not in node_displays:
213
+ node_displays[node] = node_display
214
+
215
+ if node not in global_node_displays:
216
+ global_node_displays[node] = node_display
215
217
 
216
218
  # Nodes wrapped in a decorator need to be in our node display dictionary for later retrieval
217
219
  inner_node = get_wrapped_node(node)
@@ -14,6 +14,7 @@ from vellum.workflows.references.output import OutputReference
14
14
  from vellum.workflows.types.core import JsonArray, JsonObject
15
15
  from vellum.workflows.types.generics import WorkflowType
16
16
  from vellum.workflows.utils.uuids import uuid4_from_hash
17
+ from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
17
18
  from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
18
19
  from vellum_ee.workflows.display.nodes.types import PortDisplay
19
20
  from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -43,7 +44,7 @@ class VellumWorkflowDisplay(
43
44
  WorkflowMetaVellumDisplayOverrides,
44
45
  WorkflowInputsVellumDisplay,
45
46
  WorkflowInputsVellumDisplayOverrides,
46
- BaseNodeVellumDisplay,
47
+ BaseNodeDisplay,
47
48
  EntrypointVellumDisplay,
48
49
  EntrypointVellumDisplayOverrides,
49
50
  EdgeVellumDisplay,
@@ -52,7 +53,7 @@ class VellumWorkflowDisplay(
52
53
  WorkflowOutputVellumDisplayOverrides,
53
54
  ]
54
55
  ):
55
- node_display_base_class = BaseNodeVellumDisplay
56
+ node_display_base_class = BaseNodeDisplay
56
57
 
57
58
  def serialize(self) -> JsonObject:
58
59
  input_variables: JsonArray = []
@@ -137,7 +138,7 @@ class VellumWorkflowDisplay(
137
138
  workflow_output_display.node_input_id,
138
139
  )
139
140
 
140
- source_node_display: Optional[BaseNodeVellumDisplay]
141
+ source_node_display: Optional[BaseNodeDisplay]
141
142
  first_rule = node_input.value.rules[0]
142
143
  if first_rule.type == "NODE_OUTPUT":
143
144
  source_node_id = UUID(first_rule.data.node_id)
@@ -170,15 +171,20 @@ class VellumWorkflowDisplay(
170
171
  )
171
172
 
172
173
  if source_node_display:
174
+ if isinstance(source_node_display, BaseNodeVellumDisplay):
175
+ source_handle_id = source_node_display.get_source_handle_id(
176
+ port_displays=self.display_context.port_displays
177
+ )
178
+ else:
179
+ source_handle_id = source_node_display.get_node_port_display(
180
+ source_node_display._node.Ports.default
181
+ ).id
182
+
173
183
  synthetic_output_edges.append(
174
184
  {
175
185
  "id": str(workflow_output_display.edge_id),
176
186
  "source_node_id": str(source_node_display.node_id),
177
- "source_handle_id": str(
178
- source_node_display.get_source_handle_id(
179
- port_displays=self.display_context.port_displays
180
- )
181
- ),
187
+ "source_handle_id": str(source_handle_id),
182
188
  "target_node_id": str(workflow_output_display.node_id),
183
189
  "target_handle_id": str(workflow_output_display.target_handle_id),
184
190
  "type": "DEFAULT",
@@ -279,7 +285,7 @@ class VellumWorkflowDisplay(
279
285
  self,
280
286
  entrypoint: Type[BaseNode],
281
287
  workflow_display: WorkflowMetaVellumDisplay,
282
- node_displays: Dict[Type[BaseNode], BaseNodeVellumDisplay],
288
+ node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
283
289
  overrides: Optional[EntrypointVellumDisplayOverrides] = None,
284
290
  ) -> EntrypointVellumDisplay:
285
291
  entrypoint_node_id = workflow_display.entrypoint_node_id
@@ -293,8 +299,12 @@ class VellumWorkflowDisplay(
293
299
  )
294
300
 
295
301
  entrypoint_target = get_unadorned_node(entrypoint)
296
- target_node_id = node_displays[entrypoint_target].node_id
297
- target_handle_id = node_displays[entrypoint_target].get_target_handle_id()
302
+ target_node_display = node_displays[entrypoint_target]
303
+ target_node_id = target_node_display.node_id
304
+ if isinstance(target_node_display, BaseNodeVellumDisplay):
305
+ target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(entrypoint_node_id)
306
+ else:
307
+ target_handle_id = target_node_display.get_trigger_id()
298
308
 
299
309
  edge_display = self._generate_edge_display_from_source(
300
310
  entrypoint_node_id, source_handle_id, target_node_id, target_handle_id, overrides=edge_display_overrides
@@ -339,7 +349,7 @@ class VellumWorkflowDisplay(
339
349
  def _generate_edge_display(
340
350
  self,
341
351
  edge: Edge,
342
- node_displays: Dict[Type[BaseNode], BaseNodeVellumDisplay],
352
+ node_displays: Dict[Type[BaseNode], BaseNodeDisplay],
343
353
  port_displays: Dict[Port, PortDisplay],
344
354
  overrides: Optional[EdgeVellumDisplayOverrides] = None,
345
355
  ) -> EdgeVellumDisplay:
@@ -353,8 +363,10 @@ class VellumWorkflowDisplay(
353
363
  target_node_display = node_displays[target_node]
354
364
  target_node_id = target_node_display.node_id
355
365
 
356
- target_handle_id: UUID
357
- target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
366
+ if isinstance(target_node_display, BaseNodeVellumDisplay):
367
+ target_handle_id = target_node_display.get_target_handle_id_by_source_node_id(source_node_id)
368
+ else:
369
+ target_handle_id = target_node_display.get_trigger_id()
358
370
 
359
371
  return self._generate_edge_display_from_source(
360
372
  source_node_id, source_handle_id, target_node_id, target_handle_id, overrides
@@ -1,192 +0,0 @@
1
- from typing import Any, Generic, TypeVar, cast
2
-
3
- from vellum.workflows.constants import UNDEF
4
- 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
- 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
18
- from vellum.workflows.types.core import JsonArray, JsonObject
19
- from vellum.workflows.utils.uuids import uuid4_from_hash
20
- from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_variable_type
21
- 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
23
- from vellum_ee.workflows.display.types import WorkflowDisplayContext
24
- from vellum_ee.workflows.display.utils.vellum import primitive_to_vellum_value
25
- from vellum_ee.workflows.display.vellum import GenericNodeDisplayData
26
-
27
- _BaseNodeType = TypeVar("_BaseNodeType", bound=BaseNode)
28
-
29
-
30
- class BaseNodeDisplay(BaseNodeVellumDisplay[_BaseNodeType], Generic[_BaseNodeType]):
31
- def serialize(self, display_context: WorkflowDisplayContext, **kwargs: Any) -> JsonObject:
32
- node = self._node
33
- node_id = self.node_id
34
-
35
- attributes: JsonArray = []
36
- for attribute in node:
37
- id = str(uuid4_from_hash(f"{node_id}|{attribute.name}"))
38
-
39
- attributes.append(
40
- {
41
- "id": id,
42
- "name": attribute.name,
43
- "value": self.serialize_value(display_context, cast(BaseDescriptor, attribute.instance)),
44
- }
45
- )
46
-
47
- ports: JsonArray = []
48
- for idx, port in enumerate(node.Ports):
49
- id = str(uuid4_from_hash(f"{node_id}|{idx}"))
50
-
51
- if port._condition_type:
52
- ports.append(
53
- {
54
- "id": id,
55
- "name": port.name,
56
- "type": port._condition_type.value,
57
- "expression": (
58
- self.serialize_condition(display_context, port._condition) if port._condition else None
59
- ),
60
- }
61
- )
62
- else:
63
- ports.append(
64
- {
65
- "id": id,
66
- "name": port.name,
67
- "type": "DEFAULT",
68
- }
69
- )
70
-
71
- outputs: JsonArray = []
72
- for output in node.Outputs:
73
- type = primitive_type_to_vellum_variable_type(output)
74
- value = (
75
- self.serialize_value(display_context, output.instance)
76
- if output.instance is not None and output.instance != UNDEF
77
- else None
78
- )
79
-
80
- outputs.append(
81
- {
82
- "id": str(uuid4_from_hash(f"{node_id}|{output.name}")),
83
- "name": output.name,
84
- "type": type,
85
- "value": value,
86
- }
87
- )
88
-
89
- return {
90
- "id": str(node_id),
91
- "label": node.__qualname__,
92
- "type": "GENERIC",
93
- "display_data": self.get_generic_node_display_data().dict(),
94
- "base": self.get_base().dict(),
95
- "definition": self.get_definition().dict(),
96
- "trigger": {
97
- "id": str(uuid4_from_hash(f"{node_id}|trigger")),
98
- "merge_behavior": node.Trigger.merge_behavior.value,
99
- },
100
- "ports": ports,
101
- "adornments": None,
102
- "attributes": attributes,
103
- "outputs": outputs,
104
- }
105
-
106
- def get_generic_node_display_data(self) -> GenericNodeDisplayData:
107
- explicit_value = self._get_explicit_node_display_attr("display_data", GenericNodeDisplayData)
108
- 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)