Skip to content

Commit dde8d29

Browse files
[APO-2825] Add codegen test for state variables in inline subworkflow nodes
- Add stateVariables field to InlineSubworkflowNodeData type - Update InlineSubworkflowNodeDataSerializer to handle stateVariables - Update inline-subworkflow-node.ts to use state variables from node data - Create test fixture that verifies both parent and inner state files are generated Co-Authored-By: vargas@vellum.ai <vargas@vellum.ai>
1 parent f2f3672 commit dde8d29

File tree

5 files changed

+343
-1
lines changed

5 files changed

+343
-1
lines changed

ee/codegen/src/__test__/__snapshots__/generate-code.test.ts.snap

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -153,6 +153,28 @@ class WorkflowDisplay(BaseWorkflowDisplay[Workflow]):
153153
"
154154
`;
155155

156+
exports[`generateCode > should generate code for %1 inline-subworkflow-with-state-variables.ts > nodes/subworkflow_with_state/state.py 1`] = `
157+
"from typing import Optional, Union
158+
159+
from vellum.workflows.state import BaseState
160+
161+
162+
class State(BaseState):
163+
inner_counter: Optional[Union[float, int]] = 0
164+
"
165+
`;
166+
167+
exports[`generateCode > should generate code for %1 inline-subworkflow-with-state-variables.ts > state.py 1`] = `
168+
"from typing import Optional, Union
169+
170+
from vellum.workflows.state import BaseState
171+
172+
173+
class State(BaseState):
174+
parent_counter: Optional[Union[float, int]] = 0
175+
"
176+
`;
177+
156178
exports[`generateCode > should generate code for %1 integration-trigger-with-entrypoint-and-outgoing-edges.ts > triggers/slack_message_trigger.py 1`] = `
157179
"from vellum.workflows.triggers import IntegrationTrigger
158180
Lines changed: 314 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,314 @@
1+
import { v4 as uuidv4 } from "uuid";
2+
3+
const outputVariableId = uuidv4();
4+
const outputNodeId = uuidv4();
5+
const outputInputId = uuidv4();
6+
const outputHandleId = uuidv4();
7+
const entrypointNodeId = uuidv4();
8+
const entrypointHandleId = uuidv4();
9+
10+
const parentStateVariableId = uuidv4();
11+
const innerStateVariableId = uuidv4();
12+
13+
const subworkflowNodeId = uuidv4();
14+
const subworkflowSourceHandleId = uuidv4();
15+
const subworkflowTargetHandleId = uuidv4();
16+
const subworkflowOutputNodeId = uuidv4();
17+
const subworkflowOutputId = uuidv4();
18+
19+
export default {
20+
input_variables: [],
21+
output_variables: [
22+
{
23+
id: outputVariableId,
24+
key: "result",
25+
type: "STRING",
26+
},
27+
],
28+
state_variables: [
29+
{
30+
id: parentStateVariableId,
31+
key: "parent_counter",
32+
type: "NUMBER",
33+
default: {
34+
type: "NUMBER",
35+
value: 0,
36+
},
37+
},
38+
],
39+
workflow_raw_data: {
40+
definition: {
41+
module: ["test", "workflow"],
42+
name: "Workflow",
43+
},
44+
edges: [
45+
{
46+
id: uuidv4(),
47+
source_handle_id: entrypointHandleId,
48+
source_node_id: entrypointNodeId,
49+
target_handle_id: subworkflowTargetHandleId,
50+
target_node_id: subworkflowNodeId,
51+
type: "DEFAULT",
52+
},
53+
{
54+
id: uuidv4(),
55+
source_handle_id: subworkflowSourceHandleId,
56+
source_node_id: subworkflowNodeId,
57+
target_handle_id: outputHandleId,
58+
target_node_id: outputNodeId,
59+
type: "DEFAULT",
60+
},
61+
],
62+
nodes: [
63+
{
64+
base: null,
65+
data: {
66+
label: "Entrypoint Node",
67+
source_handle_id: entrypointHandleId,
68+
},
69+
definition: null,
70+
id: entrypointNodeId,
71+
inputs: [],
72+
type: "ENTRYPOINT",
73+
},
74+
{
75+
base: {
76+
module: [
77+
"vellum",
78+
"workflows",
79+
"nodes",
80+
"displayable",
81+
"inline_subworkflow_node",
82+
"node",
83+
],
84+
name: "InlineSubworkflowNode",
85+
},
86+
data: {
87+
label: "Subworkflow With State",
88+
error_output_id: null,
89+
input_variables: [],
90+
output_variables: [
91+
{
92+
id: subworkflowOutputId,
93+
key: "subworkflow_result",
94+
type: "STRING",
95+
},
96+
],
97+
state_variables: [
98+
{
99+
id: innerStateVariableId,
100+
key: "inner_counter",
101+
type: "NUMBER",
102+
default: {
103+
type: "NUMBER",
104+
value: 0,
105+
},
106+
},
107+
],
108+
source_handle_id: subworkflowSourceHandleId,
109+
target_handle_id: subworkflowTargetHandleId,
110+
variant: "INLINE",
111+
workflow_raw_data: {
112+
definition: {
113+
module: ["test", "workflow", "subworkflow_with_state"],
114+
name: "SubworkflowWithState",
115+
},
116+
edges: [
117+
{
118+
id: uuidv4(),
119+
source_handle_id: "inner-entry-src",
120+
source_node_id: "inner-entry-id",
121+
target_handle_id: "inner-final-handle",
122+
target_node_id: subworkflowOutputNodeId,
123+
type: "DEFAULT",
124+
},
125+
],
126+
nodes: [
127+
{
128+
base: null,
129+
data: {
130+
label: "Entrypoint Node",
131+
source_handle_id: "inner-entry-src",
132+
},
133+
definition: null,
134+
id: "inner-entry-id",
135+
inputs: [],
136+
type: "ENTRYPOINT",
137+
},
138+
{
139+
base: {
140+
module: [
141+
"vellum",
142+
"workflows",
143+
"nodes",
144+
"displayable",
145+
"final_output_node",
146+
"node",
147+
],
148+
name: "FinalOutputNode",
149+
},
150+
data: {
151+
label: "Final Output",
152+
name: "subworkflow_result",
153+
node_input_id: "inner-output-input-id",
154+
output_id: subworkflowOutputId,
155+
output_type: "STRING",
156+
target_handle_id: "inner-final-handle",
157+
},
158+
definition: {
159+
module: [
160+
"test",
161+
"workflow",
162+
"subworkflow_with_state",
163+
"final_output",
164+
],
165+
name: "SubworkflowFinalOutput",
166+
},
167+
id: subworkflowOutputNodeId,
168+
inputs: [
169+
{
170+
id: "inner-output-input-id",
171+
key: "node_input",
172+
value: {
173+
combinator: "OR",
174+
rules: [
175+
{
176+
data: {
177+
type: "STRING",
178+
value: "result from subworkflow",
179+
},
180+
type: "CONSTANT_VALUE",
181+
},
182+
],
183+
},
184+
},
185+
],
186+
outputs: [
187+
{
188+
id: subworkflowOutputId,
189+
name: "value",
190+
type: "STRING",
191+
value: {
192+
type: "CONSTANT_VALUE",
193+
value: {
194+
type: "STRING",
195+
value: "result from subworkflow",
196+
},
197+
},
198+
},
199+
],
200+
ports: [],
201+
trigger: {
202+
id: "inner-final-handle",
203+
merge_behavior: "AWAIT_ANY",
204+
},
205+
type: "TERMINAL",
206+
},
207+
],
208+
output_values: [
209+
{
210+
output_variable_id: subworkflowOutputId,
211+
value: {
212+
node_id: subworkflowOutputNodeId,
213+
node_output_id: subworkflowOutputId,
214+
type: "NODE_OUTPUT",
215+
},
216+
},
217+
],
218+
},
219+
},
220+
definition: {
221+
module: ["test", "workflow", "subworkflow_with_state"],
222+
name: "SubworkflowWithStateNode",
223+
},
224+
id: subworkflowNodeId,
225+
inputs: [],
226+
ports: [
227+
{
228+
id: subworkflowSourceHandleId,
229+
name: "default",
230+
type: "DEFAULT",
231+
},
232+
],
233+
trigger: {
234+
id: subworkflowTargetHandleId,
235+
merge_behavior: "AWAIT_ATTRIBUTES",
236+
},
237+
type: "SUBWORKFLOW",
238+
},
239+
{
240+
base: {
241+
module: [
242+
"vellum",
243+
"workflows",
244+
"nodes",
245+
"displayable",
246+
"final_output_node",
247+
"node",
248+
],
249+
name: "FinalOutputNode",
250+
},
251+
data: {
252+
label: "Workflow Output",
253+
name: "result",
254+
node_input_id: outputInputId,
255+
output_id: outputVariableId,
256+
output_type: "STRING",
257+
target_handle_id: outputHandleId,
258+
},
259+
definition: {
260+
module: ["test", "workflow", "final_output"],
261+
name: "WorkflowFinalOutput",
262+
},
263+
id: outputNodeId,
264+
inputs: [
265+
{
266+
id: outputInputId,
267+
key: "node_input",
268+
value: {
269+
combinator: "OR",
270+
rules: [
271+
{
272+
data: {
273+
node_id: subworkflowNodeId,
274+
output_id: subworkflowOutputId,
275+
},
276+
type: "NODE_OUTPUT",
277+
},
278+
],
279+
},
280+
},
281+
],
282+
outputs: [
283+
{
284+
id: outputVariableId,
285+
name: "value",
286+
type: "STRING",
287+
value: {
288+
node_id: subworkflowNodeId,
289+
node_output_id: subworkflowOutputId,
290+
type: "NODE_OUTPUT",
291+
},
292+
},
293+
],
294+
ports: [],
295+
trigger: {
296+
id: outputHandleId,
297+
merge_behavior: "AWAIT_ANY",
298+
},
299+
type: "TERMINAL",
300+
},
301+
],
302+
output_values: [
303+
{
304+
output_variable_id: outputVariableId,
305+
value: {
306+
node_id: outputNodeId,
307+
node_output_id: outputVariableId,
308+
type: "NODE_OUTPUT",
309+
},
310+
},
311+
],
312+
},
313+
assertions: ["state.py", "nodes/subworkflow_with_state/state.py"],
314+
};

ee/codegen/src/generators/nodes/inline-subworkflow-node.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ export class InlineSubworkflowNode extends BaseNestedWorkflowNode<
128128
workflowVersionExecConfig: {
129129
workflowRawData: inlineSubworkflowNodeData.workflowRawData,
130130
inputVariables: inlineSubworkflowNodeData.inputVariables,
131-
stateVariables: [],
131+
stateVariables: inlineSubworkflowNodeData.stateVariables ?? [],
132132
outputVariables: inlineSubworkflowNodeData.outputVariables,
133133
},
134134
moduleName: nestedWorkflowContext.moduleName,

ee/codegen/src/serializers/vellum.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1239,6 +1239,10 @@ export const InlineSubworkflowNodeDataSerializer: ObjectSchema<
12391239
"output_variables",
12401240
listSchema(VellumVariableSerializer)
12411241
),
1242+
stateVariables: propertySchema(
1243+
"state_variables",
1244+
listSchema(VellumVariableSerializer).optional()
1245+
),
12421246
label: stringSchema(),
12431247
sourceHandleId: propertySchema("source_handle_id", stringSchema()),
12441248
targetHandleId: propertySchema("target_handle_id", stringSchema()),
@@ -1250,6 +1254,7 @@ export declare namespace InlineSubworkflowNodeDataSerializer {
12501254
workflow_raw_data: WorkflowRawDataSerializer.Raw;
12511255
input_variables: VellumVariableSerializer.Raw[];
12521256
output_variables: VellumVariableSerializer.Raw[];
1257+
state_variables?: VellumVariableSerializer.Raw[] | null;
12531258
label: string;
12541259
source_handle_id: string;
12551260
target_handle_id: string;

ee/codegen/src/types/vellum.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -522,6 +522,7 @@ export interface InlineSubworkflowNodeData {
522522
workflowRawData: WorkflowRawData;
523523
inputVariables: VellumVariable[];
524524
outputVariables: VellumVariable[];
525+
stateVariables?: VellumVariable[];
525526
label: string;
526527
sourceHandleId: string;
527528
targetHandleId: string;

0 commit comments

Comments
 (0)