stepflow_flow/workflow/
builders.rs

1// Copyright 2025 DataStax Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
4// in compliance with the License. You may obtain a copy of the License at
5//
6//     http://www.apache.org/licenses/LICENSE-2.0
7//
8// Unless required by applicable law or agreed to in writing, software distributed under the License
9// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
10// or implied. See the License for the specific language governing permissions and limitations under
11// the License.
12
13use std::collections::HashMap;
14
15use indexmap::IndexMap;
16
17use super::{
18    Component, ErrorAction, ExampleInput, Flow, FlowSchema, JsonPath, Step, TestConfig,
19    VariableSchema,
20};
21use crate::{ValueExpr, schema::SchemaRef};
22
23/// Builder for creating Flow instances with reduced boilerplate.
24#[derive(Default)]
25pub struct FlowBuilder {
26    name: Option<String>,
27    description: Option<String>,
28    version: Option<String>,
29    input_schema: Option<SchemaRef>,
30    output_schema: Option<SchemaRef>,
31    steps: Vec<Step>,
32    output: Option<ValueExpr>,
33    variables: Option<VariableSchema>,
34    test: Option<TestConfig>,
35    examples: Option<Vec<ExampleInput>>,
36    metadata: HashMap<String, serde_json::Value>,
37}
38
39impl FlowBuilder {
40    /// Create a new FlowBuilder with default values.
41    pub fn new() -> Self {
42        Default::default()
43    }
44
45    /// Set the flow name.
46    pub fn name<S: Into<String>>(mut self, name: S) -> Self {
47        self.name = Some(name.into());
48        self
49    }
50
51    /// Set the flow description.
52    pub fn description<S: Into<String>>(mut self, desc: S) -> Self {
53        self.description = Some(desc.into());
54        self
55    }
56
57    /// Set the flow version.
58    pub fn version<S: Into<String>>(mut self, version: S) -> Self {
59        self.version = Some(version.into());
60        self
61    }
62
63    /// Set the input schema.
64    pub fn input_schema(mut self, schema: SchemaRef) -> Self {
65        self.input_schema = Some(schema);
66        self
67    }
68
69    /// Set the output schema.
70    pub fn output_schema(mut self, schema: SchemaRef) -> Self {
71        self.output_schema = Some(schema);
72        self
73    }
74
75    /// Add a single step to the flow.
76    pub fn step(mut self, step: Step) -> Self {
77        self.steps.push(step);
78        self
79    }
80
81    /// Add multiple steps to the flow.
82    pub fn steps<I: IntoIterator<Item = Step>>(mut self, steps: I) -> Self {
83        self.steps.extend(steps);
84        self
85    }
86
87    /// Set the flow output.
88    pub fn output(mut self, output: ValueExpr) -> Self {
89        self.output = Some(output);
90        self
91    }
92
93    /// Set the variables schema for the flow.
94    pub fn variables(mut self, variables: VariableSchema) -> Self {
95        self.variables = Some(variables);
96        self
97    }
98
99    /// Set the test configuration.
100    pub fn test_config(mut self, test: TestConfig) -> Self {
101        self.test = Some(test);
102        self
103    }
104
105    /// Set the examples.
106    pub fn examples(mut self, examples: Vec<ExampleInput>) -> Self {
107        self.examples = Some(examples);
108        self
109    }
110
111    /// Add metadata.
112    pub fn metadata<S: Into<String>>(mut self, key: S, value: serde_json::Value) -> Self {
113        self.metadata.insert(key.into(), value);
114        self
115    }
116
117    /// Create a builder for a test flow with a default name.
118    pub fn test_flow() -> Self {
119        Self::new().name("test_workflow")
120    }
121
122    /// Create a builder for a mock flow with common test defaults.
123    pub fn mock_flow() -> Self {
124        Self::new()
125            .name("mock_flow")
126            .description("A test flow for mocking")
127    }
128
129    /// Build the final Flow instance.
130    pub fn build(self) -> Flow {
131        Flow {
132            name: self.name,
133            description: self.description,
134            version: self.version,
135            schemas: FlowSchema {
136                defs: HashMap::new(),
137                input: self.input_schema,
138                output: self.output_schema,
139                variables: self.variables.map(SchemaRef::from),
140                steps: IndexMap::new(),
141            },
142            steps: self.steps,
143            output: self.output.unwrap_or_default(),
144            test: self.test,
145            examples: self.examples,
146            metadata: self.metadata,
147        }
148    }
149}
150
151/// Builder for creating Step instances with reduced boilerplate.
152pub struct StepBuilder {
153    id: Option<String>,
154    component: Option<Component>,
155    input: Option<ValueExpr>,
156    on_error: Option<ErrorAction>,
157    must_execute: Option<bool>,
158    metadata: HashMap<String, serde_json::Value>,
159}
160
161impl StepBuilder {
162    /// Create a new StepBuilder with the given step ID.
163    pub fn new<S: Into<String>>(id: S) -> Self {
164        Self {
165            id: Some(id.into()),
166            component: None,
167            input: None,
168            on_error: None,
169            must_execute: None,
170            metadata: HashMap::new(),
171        }
172    }
173
174    /// Set the component for this step.
175    pub fn component<S: Into<String>>(mut self, component: S) -> Self {
176        self.component = Some(Component::from_string(component.into()));
177        self
178    }
179
180    /// Set the input template for this step.
181    pub fn input(mut self, input: ValueExpr) -> Self {
182        self.input = Some(input);
183        self
184    }
185
186    /// Set the input from a JSON value, parsing it as a ValueExpr.
187    pub fn input_json(mut self, input: serde_json::Value) -> Self {
188        self.input = Some(serde_json::from_value(input).unwrap());
189        self
190    }
191
192    /// Set the input as a literal JSON value.
193    pub fn input_literal(mut self, input: serde_json::Value) -> Self {
194        self.input = Some(ValueExpr::literal(input));
195        self
196    }
197
198    /// Set the error action.
199    pub fn on_error(mut self, action: ErrorAction) -> Self {
200        self.on_error = Some(action);
201        self
202    }
203
204    /// Set whether this step must execute even if its output is not used.
205    pub fn must_execute(mut self, must_execute: bool) -> Self {
206        self.must_execute = Some(must_execute);
207        self
208    }
209
210    /// Add metadata.
211    pub fn metadata<S: Into<String>>(mut self, key: S, value: serde_json::Value) -> Self {
212        self.metadata.insert(key.into(), value);
213        self
214    }
215
216    /// Create a mock step with a default mock component.
217    pub fn mock_step<S: Into<String>>(id: S) -> Self {
218        Self::new(id).component("/mock/test")
219    }
220
221    /// Create a builtin step with the specified builtin component.
222    pub fn builtin_step<S: Into<String>>(id: S, component: S) -> Self {
223        Self::new(id).component(format!("/builtin/{}", component.into()))
224    }
225
226    /// Create a step that references another step's output.
227    pub fn step_ref<S: Into<String>>(id: S, ref_step: S) -> Self {
228        Self::new(id)
229            .component("/mock/test")
230            .input(ValueExpr::step_output(ref_step))
231    }
232
233    /// Create a step that references workflow input.
234    pub fn workflow_input<S: Into<String>>(id: S) -> Self {
235        Self::new(id)
236            .component("/mock/test")
237            .input(ValueExpr::workflow_input(JsonPath::default()))
238    }
239
240    /// Build the final Step instance.
241    pub fn build(self) -> Step {
242        Step {
243            id: self.id.expect("Step ID is required"),
244            component: self
245                .component
246                .unwrap_or_else(|| Component::from_string("/mock/test")),
247            input: self.input.unwrap_or_else(ValueExpr::null),
248            on_error: self.on_error,
249            must_execute: self.must_execute,
250            metadata: self.metadata,
251        }
252    }
253}
254
255impl Flow {
256    /// Create a new FlowBuilder.
257    pub fn builder() -> FlowBuilder {
258        FlowBuilder::new()
259    }
260}
261
262impl Step {
263    /// Create a new StepBuilder with the given ID.
264    pub fn builder<S: Into<String>>(id: S) -> StepBuilder {
265        StepBuilder::new(id)
266    }
267}
268
269/// Create a FlowBuilder for a mock flow with common test defaults.
270pub fn mock_flow() -> FlowBuilder {
271    FlowBuilder::mock_flow()
272}
273
274/// Create a FlowBuilder for a test flow with a default name.
275pub fn test_flow() -> FlowBuilder {
276    FlowBuilder::test_flow()
277}
278
279/// Create a StepBuilder for a mock step.
280pub fn mock_step<S: Into<String>>(id: S) -> StepBuilder {
281    StepBuilder::mock_step(id)
282}
283
284/// Create a StepBuilder for a builtin step.
285pub fn builtin_step<S: Into<String>>(id: S, component: S) -> StepBuilder {
286    StepBuilder::builtin_step(id, component)
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use serde_json::json;
293
294    #[test]
295    fn test_flow_builder_basic() {
296        let flow = FlowBuilder::new()
297            .name("test")
298            .description("A test flow")
299            .build();
300
301        assert_eq!(flow.name, Some("test".to_string()));
302        assert_eq!(flow.description, Some("A test flow".to_string()));
303        assert!(flow.steps.is_empty());
304    }
305
306    #[test]
307    fn test_flow_builder_with_steps() {
308        let step1 = StepBuilder::mock_step("step1")
309            .input_literal(json!({"value": 42}))
310            .build();
311
312        let step2 = StepBuilder::step_ref("step2", "step1").build();
313
314        let flow = FlowBuilder::test_flow()
315            .description("Test with steps")
316            .step(step1)
317            .step(step2)
318            .output(ValueExpr::step_output("step2"))
319            .build();
320
321        assert_eq!(flow.name, Some("test_workflow".to_string()));
322        assert_eq!(flow.steps.len(), 2);
323        assert_eq!(flow.steps[0].id, "step1");
324        assert_eq!(flow.steps[1].id, "step2");
325    }
326
327    #[test]
328    fn test_step_builder_basic() {
329        let step = StepBuilder::new("test_step")
330            .component("/test/component")
331            .input_literal(json!({"key": "value"}))
332            .build();
333
334        assert_eq!(step.id, "test_step");
335        assert_eq!(step.component.path(), "/test/component");
336        assert_eq!(step.on_error, None);
337        assert_eq!(step.on_error_or_default(), ErrorAction::Fail);
338    }
339
340    #[test]
341    fn test_step_builder_convenience_methods() {
342        let mock_step = StepBuilder::mock_step("mock1").build();
343        assert_eq!(mock_step.component.path(), "/mock/test");
344
345        let builtin_step = StepBuilder::builtin_step("builtin1", "openai").build();
346        assert_eq!(builtin_step.component.path(), "/builtin/openai");
347
348        let workflow_input_step = StepBuilder::workflow_input("input1").build();
349        assert_eq!(workflow_input_step.component.path(), "/mock/test");
350
351        let step_ref = StepBuilder::step_ref("ref1", "step1").build();
352        assert_eq!(step_ref.component.path(), "/mock/test");
353    }
354
355    #[test]
356    fn test_convenience_functions() {
357        let flow = mock_flow().step(mock_step("step1").build()).build();
358
359        assert_eq!(flow.name, Some("mock_flow".to_string()));
360        assert_eq!(flow.steps.len(), 1);
361    }
362}
stepflow_flow/workflow/builders.rs

stepflow_flow/workflow/
builders.rs
