Skip to main content

workflow/
workflow.rs

1//! Workflow runtime example.
2//!
3//! This example uses a mock task runner so it can be run without configuring an
4//! LLM provider:
5//!
6//! ```powershell
7//! cargo run -p enki-next --example workflow
8//! ```
9
10use async_trait::async_trait;
11use enki_next::tooling::types::WorkflowToolContext;
12use enki_next::workflow::{
13    TaskDefinition, TaskTarget, WorkflowDefinition, WorkflowEdgeDefinition, WorkflowEdgeTransition,
14    WorkflowFailurePolicy, WorkflowNodeDefinition, WorkflowNodeKind, WorkflowRequest,
15    WorkflowRuntime, WorkflowTaskResult, WorkflowTaskRunner,
16};
17use serde_json::json;
18use std::path::Path;
19use std::sync::Arc;
20use std::time::{SystemTime, UNIX_EPOCH};
21
22struct DemoTaskRunner;
23
24#[async_trait(?Send)]
25impl WorkflowTaskRunner for DemoTaskRunner {
26    async fn run_task(
27        &self,
28        target: &TaskTarget,
29        metadata: &WorkflowToolContext,
30        workspace_dir: &Path,
31        prompt: &str,
32    ) -> Result<WorkflowTaskResult, String> {
33        let agent_id = match target {
34            TaskTarget::AgentId(agent_id) => agent_id.clone(),
35            TaskTarget::Capabilities(capabilities) => {
36                format!("agent-with-{}", capabilities.join("-"))
37            }
38        };
39
40        Ok(WorkflowTaskResult {
41            content: format!(
42                "[{agent_id}] handled node '{}' in {}\n\n{}",
43                metadata.node_id,
44                workspace_dir.display(),
45                prompt
46            ),
47            value: json!({
48                "agent_id": agent_id,
49                "node_id": metadata.node_id,
50                "content": prompt,
51                "workspace_dir": workspace_dir.display().to_string(),
52            }),
53            agent_id,
54            steps: Vec::new(),
55        })
56    }
57}
58
59#[tokio::main]
60async fn main() -> Result<(), String> {
61    let reusable_task = TaskDefinition {
62        id: "draft_release_note".to_string(),
63        target: TaskTarget::Capabilities(vec!["writing".to_string()]),
64        prompt: "Draft a concise release note for {{input.topic}}.".to_string(),
65        input_bindings: Default::default(),
66        input_transform: None,
67        output_transform: None,
68        output_key: Some("draft".to_string()),
69        retry_policy: None,
70        failure_policy: None,
71    };
72
73    let workflow = WorkflowDefinition {
74        id: "release-note-review".to_string(),
75        name: "Release Note Review".to_string(),
76        retry_policy: None,
77        failure_policy: Some(WorkflowFailurePolicy::ContinueBestEffort),
78        nodes: vec![
79            WorkflowNodeDefinition {
80                id: "draft".to_string(),
81                kind: WorkflowNodeKind::Task {
82                    task_id: Some("draft_release_note".to_string()),
83                    task: None,
84                },
85                output_key: None,
86                retry_policy: None,
87                failure_policy: None,
88            },
89            WorkflowNodeDefinition {
90                id: "review".to_string(),
91                kind: WorkflowNodeKind::Task {
92                    task_id: None,
93                    task: Some(TaskDefinition {
94                        id: "review_inline".to_string(),
95                        target: TaskTarget::AgentId("reviewer".to_string()),
96                        prompt:
97                            "Review the draft and suggest improvements:\n{{context.draft.content}}"
98                                .to_string(),
99                        input_bindings: Default::default(),
100                        input_transform: None,
101                        output_transform: None,
102                        output_key: Some("review".to_string()),
103                        retry_policy: None,
104                        failure_policy: None,
105                    }),
106                },
107                output_key: None,
108                retry_policy: None,
109                failure_policy: None,
110            },
111            WorkflowNodeDefinition {
112                id: "fact_check".to_string(),
113                kind: WorkflowNodeKind::Task {
114                    task_id: None,
115                    task: Some(TaskDefinition {
116                        id: "fact_check_inline".to_string(),
117                        target: TaskTarget::Capabilities(vec!["analysis".to_string()]),
118                        prompt:
119                            "Fact-check this draft for {{input.topic}}:\n{{context.draft.content}}"
120                                .to_string(),
121                        input_bindings: Default::default(),
122                        input_transform: None,
123                        output_transform: None,
124                        output_key: Some("fact_check".to_string()),
125                        retry_policy: None,
126                        failure_policy: None,
127                    }),
128                },
129                output_key: None,
130                retry_policy: None,
131                failure_policy: None,
132            },
133            WorkflowNodeDefinition {
134                id: "review_text".to_string(),
135                kind: WorkflowNodeKind::Transform {
136                    transform_id: "extract_content".to_string(),
137                    input_key: Some("review".to_string()),
138                },
139                output_key: Some("review_text".to_string()),
140                retry_policy: None,
141                failure_policy: None,
142            },
143            WorkflowNodeDefinition {
144                id: "review_ready".to_string(),
145                kind: WorkflowNodeKind::Decision {
146                    condition: "context.review_text != null".to_string(),
147                },
148                output_key: Some("review_ready".to_string()),
149                retry_policy: None,
150                failure_policy: None,
151            },
152            WorkflowNodeDefinition {
153                id: "merge".to_string(),
154                kind: WorkflowNodeKind::Join,
155                output_key: Some("merged".to_string()),
156                retry_policy: None,
157                failure_policy: None,
158            },
159        ],
160        edges: vec![
161            WorkflowEdgeDefinition {
162                from: "draft".to_string(),
163                to: "review".to_string(),
164                transition: WorkflowEdgeTransition::OnSuccess,
165            },
166            WorkflowEdgeDefinition {
167                from: "draft".to_string(),
168                to: "fact_check".to_string(),
169                transition: WorkflowEdgeTransition::OnSuccess,
170            },
171            WorkflowEdgeDefinition {
172                from: "review".to_string(),
173                to: "review_text".to_string(),
174                transition: WorkflowEdgeTransition::OnSuccess,
175            },
176            WorkflowEdgeDefinition {
177                from: "review_text".to_string(),
178                to: "review_ready".to_string(),
179                transition: WorkflowEdgeTransition::OnSuccess,
180            },
181            WorkflowEdgeDefinition {
182                from: "review_ready".to_string(),
183                to: "merge".to_string(),
184                transition: WorkflowEdgeTransition::Condition(
185                    "context.review_ready.matched == true".to_string(),
186                ),
187            },
188            WorkflowEdgeDefinition {
189                from: "fact_check".to_string(),
190                to: "merge".to_string(),
191                transition: WorkflowEdgeTransition::OnSuccess,
192            },
193        ],
194    };
195
196    let workspace_home = std::env::temp_dir().join(format!(
197        "enki-workflow-example-{}",
198        SystemTime::now()
199            .duration_since(UNIX_EPOCH)
200            .map_err(|err| err.to_string())?
201            .as_nanos()
202    ));
203
204    let runtime = WorkflowRuntime::builder()
205        .with_workspace_home(&workspace_home)
206        .with_task_runner(Arc::new(DemoTaskRunner))
207        .add_task(reusable_task)
208        .add_workflow(workflow)
209        .build()
210        .await?;
211
212    let response = runtime
213        .start(WorkflowRequest::new(
214            "release-note-review",
215            json!({ "topic": "runtime-managed workflows" }),
216        ))
217        .await?;
218
219    println!("Workflow: {}", response.workflow_id);
220    println!("Run: {}", response.run_id);
221    println!("Status: {:?}", response.status);
222    println!("Workspace: {}", workspace_home.display());
223    println!(
224        "Context:\n{}",
225        serde_json::to_string_pretty(&response.context.to_value()).map_err(|err| err.to_string())?
226    );
227
228    Ok(())
229}