graph_flow/lib.rs
1//! # graph-flow
2//!
3//! A high-performance, type-safe framework for building multi-agent workflow systems in Rust.
4//!
5//! ## Features
6//!
7//! - **Type-Safe Workflows**: Compile-time guarantees for workflow correctness
8//! - **Flexible Execution**: Step-by-step, batch, or mixed execution modes
9//! - **Built-in Persistence**: PostgreSQL and in-memory storage backends with
10//! optimistic locking
11//! - **LLM Integration**: Optional integration with Rig for AI agent capabilities
12//! - **Human-in-the-Loop**: Natural workflow interruption and resumption
13//! - **Async/Await Native**: Built from the ground up for async Rust
14//!
15//! ## Quick Start
16//!
17//! ```rust
18//! use graph_flow::{Context, Task, TaskResult, NextAction, GraphBuilder, FlowRunner, InMemorySessionStorage, Session, SessionStorage};
19//! use async_trait::async_trait;
20//! use std::sync::Arc;
21//!
22//! // Define a task
23//! struct HelloTask;
24//!
25//! #[async_trait]
26//! impl Task for HelloTask {
27//! fn id(&self) -> &str {
28//! "hello_task"
29//! }
30//!
31//! async fn run(&self, context: Context) -> graph_flow::Result<TaskResult> {
32//! let name: String = context.get("name").unwrap_or("World".to_string());
33//! let greeting = format!("Hello, {}!", name);
34//!
35//! context.set("greeting", greeting.clone())?;
36//! Ok(TaskResult::new(Some(greeting), NextAction::Continue))
37//! }
38//! }
39//!
40//! # #[tokio::main]
41//! # async fn main() -> graph_flow::Result<()> {
42//! // Build the workflow
43//! let hello_task = Arc::new(HelloTask);
44//! let graph = Arc::new(
45//! GraphBuilder::new("greeting_workflow")
46//! .add_task(hello_task.clone())
47//! .build()?
48//! );
49//!
50//! // Set up session storage and runner
51//! let session_storage = Arc::new(InMemorySessionStorage::new());
52//! let flow_runner = FlowRunner::new(graph.clone(), session_storage.clone());
53//!
54//! // Create and execute session
55//! let session = Session::new_from_task("user_123".to_string(), hello_task.id());
56//! session.context.set("name", "Alice".to_string())?;
57//! session_storage.save(session).await?;
58//!
59//! let result = flow_runner.run("user_123").await?;
60//! println!("Response: {:?}", result.response);
61//! # Ok(())
62//! # }
63//! ```
64//!
65//! ## Core Concepts
66//!
67//! ### Tasks
68//!
69//! Tasks are the building blocks of your workflow. They implement the [`Task`] trait:
70//!
71//! ```rust
72//! use graph_flow::{Task, TaskResult, NextAction, Context};
73//! use async_trait::async_trait;
74//!
75//! struct MyTask;
76//!
77//! #[async_trait]
78//! impl Task for MyTask {
79//! fn id(&self) -> &str {
80//! "my_task"
81//! }
82//!
83//! async fn run(&self, context: Context) -> graph_flow::Result<TaskResult> {
84//! // Your task logic here
85//! Ok(TaskResult::new(Some("Done!".to_string()), NextAction::End))
86//! }
87//! }
88//! ```
89//!
90//! ### Context
91//!
92//! The [`Context`] provides thread-safe state management across your workflow.
93//! Its methods are synchronous, so they work both in async tasks and in
94//! edge-condition closures:
95//!
96//! ```rust
97//! # use graph_flow::Context;
98//! # fn main() -> graph_flow::Result<()> {
99//! let context = Context::new();
100//!
101//! // Store and retrieve data
102//! context.set("key", "value")?;
103//! let value: Option<String> = context.get("key");
104//!
105//! // Chat history management
106//! context.add_user_message("Hello!".to_string());
107//! context.add_assistant_message("Hi there!".to_string());
108//! # Ok(())
109//! # }
110//! ```
111//!
112//! ### Graph Building
113//!
114//! Use [`GraphBuilder`] to create workflows. [`GraphBuilder::build`] validates
115//! the graph (edge endpoints and the start task must exist) and returns an
116//! immutable [`Graph`]:
117//!
118//! ```rust
119//! # use graph_flow::{GraphBuilder, Task, TaskResult, NextAction, Context};
120//! # use async_trait::async_trait;
121//! # use std::sync::Arc;
122//! # struct Task1; struct Task2; struct Task3;
123//! # #[async_trait] impl Task for Task1 { fn id(&self) -> &str { "task1" } async fn run(&self, _: Context) -> graph_flow::Result<TaskResult> { Ok(TaskResult::new(None, NextAction::End)) } }
124//! # #[async_trait] impl Task for Task2 { fn id(&self) -> &str { "task2" } async fn run(&self, _: Context) -> graph_flow::Result<TaskResult> { Ok(TaskResult::new(None, NextAction::End)) } }
125//! # #[async_trait] impl Task for Task3 { fn id(&self) -> &str { "task3" } async fn run(&self, _: Context) -> graph_flow::Result<TaskResult> { Ok(TaskResult::new(None, NextAction::End)) } }
126//! # fn main() -> graph_flow::Result<()> {
127//! # let task1 = Arc::new(Task1); let task2 = Arc::new(Task2); let task3 = Arc::new(Task3);
128//! let graph = GraphBuilder::new("my_workflow")
129//! .add_task(task1.clone())
130//! .add_task(task2.clone())
131//! .add_task(task3.clone())
132//! .add_edge(task1.id(), task2.id())
133//! .add_conditional_edge(
134//! task2.id(),
135//! |ctx| ctx.get::<bool>("condition").unwrap_or(false),
136//! task3.id(), // if true
137//! task1.id(), // if false
138//! )
139//! .build()?;
140//! # Ok(())
141//! # }
142//! ```
143//!
144//! ### Execution
145//!
146//! Use [`FlowRunner`] for convenient session-based execution:
147//!
148//! ```rust,no_run
149//! # use graph_flow::{FlowRunner, InMemorySessionStorage, Session, Graph, SessionStorage};
150//! # use std::sync::Arc;
151//! # #[tokio::main]
152//! # async fn main() -> graph_flow::Result<()> {
153//! # let graph = Arc::new(Graph::new("test"));
154//! let storage = Arc::new(InMemorySessionStorage::new());
155//! let runner = FlowRunner::new(graph, storage.clone());
156//!
157//! // Create session
158//! let session = Session::new_from_task("session_id".to_string(), "start_task");
159//! storage.save(session).await?;
160//!
161//! // Execute workflow
162//! let result = runner.run("session_id").await?;
163//! # Ok(())
164//! # }
165//! ```
166//!
167//! ## Features
168//!
169//! - **Default**: Core workflow functionality
170//! - **`rig`**: Enables LLM integration via the Rig crate
171//!
172//! ## Storage Backends
173//!
174//! - [`InMemorySessionStorage`]: For development and testing
175//! - [`PostgresSessionStorage`]: For production use with PostgreSQL
176
177pub mod context;
178pub mod error;
179pub mod graph;
180pub mod runner;
181pub mod storage;
182pub mod storage_postgres;
183pub mod task;
184pub mod fanout;
185
186// Re-export commonly used types
187pub use context::{ChatHistory, Context, MessageRole, SerializableMessage};
188pub use error::{GraphError, Result};
189pub use graph::{ExecutionResult, ExecutionStatus, Graph, GraphBuilder};
190pub use runner::FlowRunner;
191pub use storage::{
192 GraphStorage, InMemoryGraphStorage, InMemorySessionStorage, Session, SessionStorage,
193};
194pub use storage_postgres::PostgresSessionStorage;
195pub use task::{NextAction, Task, TaskResult};
196pub use fanout::FanOutTask;
197
198#[cfg(test)]
199mod tests {
200 use super::*;
201 use async_trait::async_trait;
202 use std::sync::Arc;
203
204 struct TestTask {
205 id: String,
206 }
207
208 #[async_trait]
209 impl Task for TestTask {
210 fn id(&self) -> &str {
211 &self.id
212 }
213
214 async fn run(&self, context: Context) -> Result<TaskResult> {
215 let input: String = context.get("input").unwrap_or_default();
216 context.set("output", format!("Processed: {}", input))?;
217
218 Ok(TaskResult::new(
219 Some("Task completed".to_string()),
220 NextAction::End,
221 ))
222 }
223 }
224
225 #[tokio::test]
226 async fn test_simple_graph_execution() {
227 let task = Arc::new(TestTask {
228 id: "test_task".to_string(),
229 });
230
231 let graph = GraphBuilder::new("test_graph")
232 .add_task(task)
233 .build()
234 .unwrap();
235
236 let mut session = Session::new_from_task("s1".to_string(), "test_task");
237 session.context.set("input", "Hello, World!").unwrap();
238
239 let result = graph.execute_session(&mut session).await.unwrap();
240
241 assert!(result.response.is_some());
242 assert!(matches!(result.status, ExecutionStatus::Completed));
243
244 let output: String = session.context.get("output").unwrap();
245 assert_eq!(output, "Processed: Hello, World!");
246 }
247
248 #[tokio::test]
249 async fn test_build_rejects_dangling_edge() {
250 let task = Arc::new(TestTask {
251 id: "only_task".to_string(),
252 });
253
254 let result = GraphBuilder::new("bad_graph")
255 .add_task(task)
256 .add_edge("only_task", "missing_task")
257 .build();
258
259 assert!(matches!(result, Err(GraphError::InvalidEdge(_))));
260 }
261
262 #[tokio::test]
263 async fn test_build_rejects_unknown_start_task() {
264 let task = Arc::new(TestTask {
265 id: "only_task".to_string(),
266 });
267
268 let result = GraphBuilder::new("bad_graph")
269 .add_task(task)
270 .set_start_task("missing_task")
271 .build();
272
273 assert!(matches!(result, Err(GraphError::TaskNotFound(_))));
274 }
275
276 #[tokio::test]
277 async fn test_storage() {
278 let graph_storage = InMemoryGraphStorage::new();
279 let session_storage = InMemorySessionStorage::new();
280
281 let graph = Arc::new(Graph::new("test"));
282 graph_storage
283 .save("test".to_string(), graph.clone())
284 .await
285 .unwrap();
286
287 let retrieved = graph_storage.get("test").await.unwrap();
288 assert!(retrieved.is_some());
289
290 let session =
291 Session::new_from_task("session1".to_string(), "task1").with_graph_id("test");
292
293 session_storage.save(session).await.unwrap();
294 let retrieved_session = session_storage.get("session1").await.unwrap();
295 assert!(retrieved_session.is_some());
296 assert_eq!(retrieved_session.unwrap().graph_id, "test");
297 }
298}