Skip to content

Commit f5034f3

Browse files
devin-ai-integration[bot]dvargas92495
devin-ai-integration[bot]
and
dvargas92495
authored
Change EXECUTION_COUNTER nodeId to node_id and add test for existing node (#3336)
* Add test for execution counter referencing existing node in mocks Co-Authored-By: [email protected] <[email protected]> * Change EXECUTION_COUNTER nodeId to node_id (snake_case) Co-Authored-By: [email protected] <[email protected]> * Fix additional nodeId to node_id references in tests and serializer Co-Authored-By: [email protected] <[email protected]> * Finish up support --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: [email protected] <[email protected]> Co-authored-by: David Vargas <[email protected]>
1 parent 8aa9a97 commit f5034f3

File tree

8 files changed

+293
-22
lines changed

8 files changed

+293
-22
lines changed

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,33 @@ class InnerMapNode(MapNode):
203203
"
204204
`;
205205

206+
exports[`generateCode > should generate code for %1 node-mock.ts > sandbox.py 1`] = `
207+
"from vellum.workflows import MockNodeExecution
208+
from vellum.workflows.inputs import DatasetRow
209+
from vellum.workflows.sandbox import WorkflowSandboxRunner
210+
211+
from .nodes.start_node import StartNode
212+
from .workflow import Workflow
213+
214+
dataset = [
215+
DatasetRow(
216+
label="First Scenario",
217+
mocks=[
218+
MockNodeExecution(
219+
when_condition=StartNode.Execution.count.greater_than_or_equal_to("Hello, World!"),
220+
then_outputs=StartNode.Outputs(result="Hello, World!"),
221+
),
222+
],
223+
),
224+
]
225+
226+
runner = WorkflowSandboxRunner(workflow=Workflow(), dataset=dataset)
227+
228+
if __name__ == "__main__":
229+
runner.run()
230+
"
231+
`;
232+
206233
exports[`generateCode > should generate code for %1 scheduled-trigger.ts > triggers/scheduled.py 1`] = `
207234
"from vellum_ee.workflows.display.editor import NodeDisplayComment
208235
from vellum.workflows.triggers import ScheduleTrigger

‎ee/codegen/src/__test__/__snapshots__/workflow-sandbox.test.ts.snap‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,35 @@ if __name__ == "__main__":
235235
"
236236
`;
237237

238+
exports[`Workflow Sandbox > write > should handle mocks with when_condition referencing existing node ID via execution counter 1`] = `
239+
"from vellum.workflows import MockNodeExecution
240+
from vellum.workflows.inputs import DatasetRow
241+
from vellum.workflows.sandbox import WorkflowSandboxRunner
242+
243+
from .inputs import Inputs
244+
from .nodes.my_custom_node import MyCustomNode
245+
from .workflow import TestWorkflow
246+
247+
dataset = [
248+
DatasetRow(
249+
label="Scenario with existing node in when_condition",
250+
inputs=Inputs(test_input="test-value"),
251+
mocks=[
252+
MockNodeExecution(
253+
when_condition=MyCustomNode.Execution.count.equals(1),
254+
then_outputs=MyCustomNode.Outputs(result="mocked_result"),
255+
),
256+
],
257+
),
258+
]
259+
260+
runner = WorkflowSandboxRunner(workflow=TestWorkflow(), dataset=dataset)
261+
262+
if __name__ == "__main__":
263+
runner.run()
264+
"
265+
`;
266+
238267
exports[`Workflow Sandbox > write > should handle mocks with when_condition referencing non-existent node ID via execution counter 1`] = `
239268
"from vellum.workflows import MockNodeExecution
240269
from vellum.workflows.inputs import DatasetRow

‎ee/codegen/src/__test__/generate-code-fixtures/node-mock.ts‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
export default {
2+
workflow_raw_data: {
3+
nodes: [
4+
{
5+
id: "entrypoint-node",
6+
type: "ENTRYPOINT",
7+
data: {
8+
label: "Entrypoint",
9+
source_handle_id: "entrypoint-source",
10+
},
11+
inputs: [],
12+
},
13+
{
14+
id: "start-node",
15+
type: "GENERIC",
16+
label: "Start Node",
17+
base: {
18+
name: "BaseNode",
19+
module: ["vellum", "workflows", "nodes", "bases", "base"],
20+
},
21+
definition: {
22+
name: "StartNode",
23+
module: ["testing", "nodes", "start_node"],
24+
},
25+
trigger: {
26+
id: "start-target",
27+
merge_behavior: "AWAIT_ATTRIBUTES",
28+
},
29+
outputs: [
30+
{
31+
id: "node-output-id",
32+
name: "result",
33+
type: "STRING",
34+
value: null,
35+
},
36+
],
37+
ports: [],
38+
attributes: [],
39+
},
40+
],
41+
edges: [
42+
{
43+
id: "edge-1",
44+
source_node_id: "entrypoint-node",
45+
source_handle_id: "entrypoint-source",
46+
target_node_id: "start-node",
47+
target_handle_id: "start-target",
48+
type: "DEFAULT",
49+
},
50+
],
51+
output_values: [
52+
{
53+
output_variable_id: "workflow-output-variable-id",
54+
value: {
55+
type: "NODE_OUTPUT",
56+
node_id: "start-node",
57+
node_output_id: "node-output-id",
58+
},
59+
},
60+
],
61+
},
62+
input_variables: [],
63+
output_variables: [
64+
{
65+
id: "workflow-output-variable-id",
66+
key: "result",
67+
type: "STRING",
68+
},
69+
],
70+
dataset: [
71+
{
72+
id: "first-scenario",
73+
label: "First Scenario",
74+
inputs: [],
75+
mocks: [
76+
{
77+
node_id: "start-node",
78+
when_condition: {
79+
type: "BINARY_EXPRESSION",
80+
operator: ">=",
81+
lhs: {
82+
type: "EXECUTION_COUNTER",
83+
node_id: "start-node",
84+
},
85+
rhs: {
86+
type: "CONSTANT_VALUE",
87+
value: {
88+
type: "STRING",
89+
value: "Hello, World!",
90+
},
91+
},
92+
},
93+
then_outputs: {
94+
result: "Hello, World!",
95+
},
96+
},
97+
],
98+
},
99+
],
100+
assertions: ["sandbox.py"],
101+
};

‎ee/codegen/src/__test__/generate-code.test.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ describe("generateCode", () => {
3535

3636
it.each(fixtures)(`should generate code for %1`, async (fixture) => {
3737
const fixtureModule = await import(path.join(fixturesDir, fixture));
38-
const { assertions, ...workflowVersionExecConfigData } =
38+
const { assertions, dataset, ...workflowVersionExecConfigData } =
3939
fixtureModule.default;
4040

4141
const fixtureDir = path.join(tempDir, fixture.replace(/\.ts$/, ""));
@@ -62,6 +62,7 @@ describe("generateCode", () => {
6262
workflowVersionExecConfigData,
6363
moduleName,
6464
vellumApiKey,
65+
sandboxInputs: dataset,
6566
});
6667
await project.generateCode();
6768

‎ee/codegen/src/__test__/workflow-sandbox.test.ts‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -658,5 +658,82 @@ describe("Workflow Sandbox", () => {
658658
// THEN it should generate MockNodeExecution with None for the unresolved reference
659659
expect(result).toMatchSnapshot();
660660
});
661+
662+
it("should handle mocks with when_condition referencing existing node ID via execution counter", async () => {
663+
/**
664+
* Tests that mocks with when_condition containing EXECUTION_COUNTER reference
665+
* to an existing node ID generate the correct node execution counter reference.
666+
*/
667+
668+
const writer = new Writer();
669+
const uniqueWorkflowContext = workflowContextFactory();
670+
const inputVariable: VellumVariable = {
671+
id: "1",
672+
key: "test_input",
673+
type: "STRING",
674+
};
675+
676+
uniqueWorkflowContext.addInputVariableContext(
677+
inputVariableContextFactory({
678+
inputVariableData: inputVariable,
679+
workflowContext: uniqueWorkflowContext,
680+
})
681+
);
682+
683+
const genericNodeData = genericNodeFactory();
684+
await nodeContextFactory({
685+
workflowContext: uniqueWorkflowContext,
686+
nodeData: genericNodeData,
687+
});
688+
689+
// GIVEN a mock with when_condition referencing an existing node ID via execution counter
690+
const sandboxInputs: WorkflowSandboxDatasetRow[] = [
691+
{
692+
label: "Scenario with existing node in when_condition",
693+
inputs: [
694+
{
695+
name: inputVariable.key,
696+
type: "STRING",
697+
value: "test-value",
698+
},
699+
],
700+
mocks: [
701+
{
702+
node_id: genericNodeData.id,
703+
when_condition: {
704+
type: "BINARY_EXPRESSION",
705+
operator: "=",
706+
lhs: {
707+
type: "EXECUTION_COUNTER",
708+
nodeId: genericNodeData.id,
709+
},
710+
rhs: {
711+
type: "CONSTANT_VALUE",
712+
value: {
713+
type: "NUMBER",
714+
value: 1,
715+
},
716+
},
717+
},
718+
then_outputs: {
719+
result: "mocked_result",
720+
},
721+
},
722+
],
723+
},
724+
];
725+
726+
// WHEN we generate the sandbox file
727+
const sandbox = codegen.workflowSandboxFile({
728+
workflowContext: uniqueWorkflowContext,
729+
sandboxInputs,
730+
});
731+
732+
sandbox.write(writer);
733+
const result = await writer.toStringFormatted();
734+
735+
// THEN it should generate MockNodeExecution with the node's Execution.count reference
736+
expect(result).toMatchSnapshot();
737+
});
661738
});
662739
});

‎ee/codegen/src/project.ts‎

Lines changed: 17 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,10 @@ import { TemplatingNode } from "src/generators/nodes/templating-node";
6767
import { IntegrationTrigger } from "src/generators/triggers/integration-trigger";
6868
import { ScheduledTrigger } from "src/generators/triggers/scheduled-trigger";
6969
import { WorkflowSandboxFile } from "src/generators/workflow-sandbox-file";
70-
import { WorkflowVersionExecConfigSerializer } from "src/serializers/vellum";
70+
import {
71+
WorkflowSandboxDatasetRowsSerializer,
72+
WorkflowVersionExecConfigSerializer,
73+
} from "src/serializers/vellum";
7174
import {
7275
CodeResourceDefinition,
7376
FinalOutputNode as FinalOutputNodeType,
@@ -100,7 +103,7 @@ export declare namespace WorkflowProjectGenerator {
100103
workflowVersionExecConfigData: unknown;
101104
vellumApiKey?: string;
102105
vellumApiEnvironment?: VellumEnvironmentUrls;
103-
sandboxInputs?: WorkflowSandboxDatasetRow[];
106+
sandboxInputs?: unknown;
104107
options?: WorkflowProjectGeneratorOptions;
105108
}
106109

@@ -173,7 +176,18 @@ ${errors.slice(0, 3).map((err) => {
173176
pythonCodeMergeableNodeFiles: new Set<string>(),
174177
triggers: this.workflowVersionExecConfig.triggers,
175178
});
176-
this.sandboxInputs = rest.sandboxInputs;
179+
180+
const sandboxInputsResult = WorkflowSandboxDatasetRowsSerializer.parse(
181+
rest.sandboxInputs,
182+
{
183+
allowUnrecognizedUnionMembers: true,
184+
allowUnrecognizedEnumValues: true,
185+
unrecognizedObjectKeys: "strip",
186+
}
187+
);
188+
this.sandboxInputs = sandboxInputsResult.ok
189+
? sandboxInputsResult.value
190+
: undefined;
177191
}
178192
}
179193

‎ee/codegen/src/serializers/vellum.ts‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ import {
2525
VellumValue as VellumValueSerializer,
2626
VellumVariable as VellumVariableSerializer,
2727
VellumVariableType as VellumVariableTypeSerializer,
28+
WorkflowInput as WorkflowInputSerializer,
2829
} from "vellum-ai/serialization";
2930
import { ConditionCombinator as ConditionCombinatorSerializer } from "vellum-ai/serialization/types/ConditionCombinator";
3031

@@ -119,6 +120,7 @@ import {
119120
WorkflowNode,
120121
WorkflowOutputValue,
121122
WorkflowRawData,
123+
WorkflowSandboxDatasetRowMock,
122124
WorkflowSandboxRoutingConfig,
123125
WorkflowStatePointer,
124126
WorkflowStateVariableWorkflowReference,
@@ -2589,3 +2591,39 @@ export declare namespace WorkflowDisplayDataSerializer {
25892591
viewport: WorkflowDisplayDataViewportSerializer.Raw;
25902592
}
25912593
}
2594+
2595+
export const WorkflowSandboxInputsSerializer = listSchema(
2596+
WorkflowInputSerializer
2597+
);
2598+
2599+
export declare namespace WorkflowSandboxDatasetRowMockSerializer {
2600+
interface Raw {
2601+
node_id: string;
2602+
when_condition?: WorkflowValueDescriptorSerializer.Raw | null;
2603+
then_outputs?: Record<string, unknown> | null;
2604+
}
2605+
}
2606+
2607+
export const WorkflowSandboxDatasetRowMockSerializer: ObjectSchema<
2608+
WorkflowSandboxDatasetRowMockSerializer.Raw,
2609+
WorkflowSandboxDatasetRowMock
2610+
> = objectSchema({
2611+
node_id: stringSchema(),
2612+
when_condition: WorkflowValueDescriptorSerializer.optional(),
2613+
then_outputs: recordSchema(stringSchema(), unknownSchema()).optional(),
2614+
});
2615+
2616+
const WorkflowSandboxDatasetRowSerializer = undiscriminatedUnionSchema([
2617+
WorkflowSandboxInputsSerializer,
2618+
objectSchema({
2619+
id: stringSchema().optional(),
2620+
label: stringSchema(),
2621+
inputs: WorkflowSandboxInputsSerializer.optional(),
2622+
workflow_trigger_id: stringSchema().optional(),
2623+
mocks: listSchema(WorkflowSandboxDatasetRowMockSerializer).optional(),
2624+
}),
2625+
]);
2626+
2627+
export const WorkflowSandboxDatasetRowsSerializer = listSchema(
2628+
WorkflowSandboxDatasetRowSerializer
2629+
);

0 commit comments

Comments
 (0)