graph-flow

High-Performance Graph Workflow Framework

A type-safe and flexible framework for building complex and interactive workflows in Rust for AI Agents and more

cargo add graph-flow

use async_trait::async_trait;
use graph_flow::{Context, Task, TaskResult, NextAction};

struct HelloTask;

#[async_trait]
impl Task for HelloTask {
    // The id() method now has a default implementation that uses std::any::type_name::<Self>()
    // You can override it if you need a custom identifier

    async fn run(&self, context: Context) -> graph_flow::Result<TaskResult> {
        // Get input from context
        let name: String = context.get_sync("name").unwrap();
        let greeting = format!("Hello, {}", name);
        
        // Store result for next task
        context.set("greeting", greeting.clone()).await;

        // Control flow: Continue to next task 
        // but give control back to workflow manager, to return response to client
        Ok(TaskResult::new(Some(greeting), NextAction::Continue))
    }
}

use graph_flow::{GraphBuilder, InMemorySessionStorage, FlowRunner};
use std::sync::Arc;

// Create task instances
let hello_task = Arc::new(HelloTask);
let excitement_task = Arc::new(ExcitementTask);

// Build the graph
let graph = Arc::new(GraphBuilder::new("greeting_workflow")
    .add_task(hello_task.clone())
    .add_task(excitement_task.clone())
    .add_edge(hello_task.id(), excitement_task.id()) // Connect tasks
    .build());

// Create storage and runner
let session_storage = Arc::new(InMemorySessionStorage::new());
let flow_runner = FlowRunner::new(graph.clone(), session_storage.clone());

// Create a session starting from the first task
let session = Session::new_from_task("session_001".to_string(), hello_task.id());
session.context.set("name", "Batman".to_string()).await;
session_storage.save(session).await?;

// Execute step by step - FlowRunner handles load/execute/save automatically
loop {
    let result = flow_runner.run("session_001").await?;
    println!("Response: {:?}", result.response);
    
    match result.status {
        ExecutionStatus::Completed => break,
        ExecutionStatus::Paused { next_task_id } => continue, // Will auto-continue to next_task_id
        ExecutionStatus::WaitingForInput => continue,
        ExecutionStatus::Error(err) => return Err(err),
    }
}

// Manual session management
loop {
    let mut session = session_storage.get("session_001").await?.unwrap();
    let result = graph.execute_session(&mut session).await?;
    session_storage.save(session).await?;
    
    match result.status {
        ExecutionStatus::Completed => break,
        ExecutionStatus::Paused { next_task_id } => continue, // Will auto-continue to next_task_id
        ExecutionStatus::WaitingForInput => continue,
        ExecutionStatus::Error(err) => return Err(err),
    }
}

// Task returns Continue - gives control back to caller
Ok(TaskResult::new(Some("Done".to_string()), NextAction::Continue))

let flow_runner = FlowRunner::new(graph, session_storage);

loop {
    let result = flow_runner.run(&session_id).await?;
    
    match result.status {
        ExecutionStatus::Completed => break,
        ExecutionStatus::Paused { next_task_id } => continue,  // Will auto-continue to next_task_id
        ExecutionStatus::WaitingForInput => continue,  // Get user input, then continue
        ExecutionStatus::Error(e) => return Err(e),
    }
}

// Task returns ContinueAndExecute - continues automatically
Ok(TaskResult::new(Some("Done".to_string()), NextAction::ContinueAndExecute))

// Runs automatically until End, WaitForInput, or error
let result = flow_runner.run(&session_id).await?;

#[derive(Debug, Clone)]
pub enum ExecutionStatus {
    /// Paused, will continue automatically to the specified next task
    Paused { next_task_id: String },
    /// Waiting for user input to continue
    WaitingForInput,
    /// Workflow completed successfully
    Completed,
    /// Error occurred during execution
    Error(String),
}

struct SentimentAnalysisTask;

#[async_trait]
impl Task for SentimentAnalysisTask {
    async fn run(&self, ctx: Context) -> graph_flow::Result<TaskResult> {
        // Very naive heuristic so the example works without an LLM
        let input: String = ctx.get_sync("user_input").unwrap_or_default();
        let sentiment = if input.to_lowercase().contains("good") {
            "positive"
        } else {
            "negative"
        };
        ctx.set("sentiment", sentiment.to_string()).await;
        Ok(TaskResult::new(None, NextAction::Continue))
    }
}

let graph = GraphBuilder::new("sentiment_flow")
    .add_task(sentiment_task)          // Detect sentiment
    .add_task(positive_task)           // Reply for happy mood
    .add_task(negative_task)           // Reply for unhappy mood
    .add_conditional_edge(
        sentiment_task.id(),
        |ctx| ctx.get_sync::<String>("sentiment")
                 .map(|s| s == "positive")
                 .unwrap_or(false),
        positive_task.id(),            // yes branch → positive
        negative_task.id(),            // else      → negative
    )
    .build();

use rig::{agent::Agent, providers::openrouter};

async fn run(&self, context: Context) -> graph_flow::Result<TaskResult> {
    // Get user input
    let user_input: String = context.get_sync("user_input").unwrap();
    
    // Create LLM agent
    let client = openrouter::Client::new(&api_key);
    let agent = client.agent("openai/gpt-4o-mini")
        .preamble("You are a helpful insurance assistant")
        .build();
    
    // Get chat history for context
    let chat_history = context.get_rig_messages().await;
    
    // Generate response
    let response = agent.chat(&user_input, chat_history).await?;
    
    // Store conversation
    context.add_user_message(user_input).await;
    context.add_assistant_message(response.clone()).await;
    
    Ok(TaskResult::new(Some(response), NextAction::Continue))
}

// Add messages to conversation
context.add_user_message("What's my claim status?".to_string()).await;
context.add_assistant_message("Your claim is being processed".to_string()).await;
context.add_system_message("Claim updated".to_string()).await;

// Retrieve conversation history
let history = context.get_chat_history().await;
let recent = context.get_last_messages(5).await;

// Chat history is automatically serialized with session state
let serialized = serde_json::to_string(&context).unwrap();

use graph_flow::{FanOutTask, GraphBuilder, Task, TaskResult, NextAction, Context};

// Define child tasks that will run in parallel
struct ChildA;
struct ChildB;

#[async_trait]
impl Task for ChildA {
    fn id(&self) -> &str { "child_a" }
    async fn run(&self, ctx: Context) -> graph_flow::Result<TaskResult> {
        let input: String = ctx.get("input").await.unwrap_or_default();
        ctx.set("a_out", format!("{}-A", input)).await;
        Ok(TaskResult::new(Some("A done".to_string()), NextAction::End))
    }
}

// Create FanOut task with children
let fanout = FanOutTask::new("fanout", vec![
    Arc::new(ChildA), 
    Arc::new(ChildB)
]);

// Build graph normally - FanOut is just another task
let graph = GraphBuilder::new("fanout_demo")
    .add_task(prepare_task)
    .add_task(fanout.clone())
    .add_task(consume_task)
    .add_edge(prepare_task.id(), fanout.id())
    .add_edge(fanout.id(), consume_task.id())
    .build();

// Default aggregation keys (without prefix)
// fanout.child_a.response - child's response message
// fanout.child_a.status   - child's status message  
// fanout.child_a.next_action - diagnostic info

// With custom prefix
let fanout = FanOutTask::new("fanout", children)
    .with_prefix("parallel");

// Results stored as:
// parallel.child_a.response
// parallel.child_b.response

struct ConsumeResults;

#[async_trait] 
impl Task for ConsumeResults {
    async fn run(&self, ctx: Context) -> graph_flow::Result<TaskResult> {
        // Read results from both children
        let a_response: Option<String> = ctx.get("fanout.child_a.response").await;
        let b_response: Option<String> = ctx.get("fanout.child_b.response").await;
        
        let summary = format!("A: {:?}, B: {:?}", a_response, b_response);
        Ok(TaskResult::new(Some(summary), NextAction::End))
    }
}

cargo run --bin fanout_basic

// In-memory storage (development)
let session_storage = Arc::new(InMemorySessionStorage::new());

// PostgreSQL storage (production)
let session_storage = Arc::new(
    PostgresSessionStorage::connect(&database_url).await?
);

// Both implement the same SessionStorage trait

// From insurance_type_classifier.rs
let agent = get_llm_agent(
    "You are an insurance classifier. Analyze claim descriptions and classify them as either 'car' or 'apartment' insurance claims."
)?;

let response = agent.chat(&user_input, chat_history).await?;

// Parse JSON from LLM response
let claim_details: ClaimDetails = serde_json::from_str(&json_response)?;
context.set(session_keys::CLAIM_DETAILS, claim_details).await;

// Route based on insurance type determined by LLM
.add_conditional_edge(
    classifier_id,
    |context| {
        context.get_sync::<String>(session_keys::INSURANCE_TYPE)
            .map(|t| t == "car")
            .unwrap_or(false)
    },
    car_details_id,
    apartment_details_id,
)

// In smart_claim_validator.rs
if estimated_cost >= 1000.0 {
    // Wait for manual approval
    return Ok(TaskResult::new(
        Some("Claim requires manual approval. Please review and approve/reject.".to_string()),
        NextAction::WaitForInput
    ));
}

# Set up environment
export OPENROUTER_API_KEY="your-key"
export DATABASE_URL="postgresql://user:pass@localhost/db" # Optional

# Start the service
cargo run --bin insurance-claims-service

# Test with curl
curl -X POST http://localhost:3000/execute \
  -H "Content-Type: application/json" \
  -d '{"content": "I need to file a claim for my car accident"}'

// main.rs - Create components once at startup
struct AppState {
    session_storage: Arc<dyn SessionStorage>,
    flow_runner: FlowRunner,
}

// HTTP handler - execute workflow step
async fn execute_graph(
    State(state): State<AppState>,
    Json(request): Json<ExecuteRequest>,
) -> Result<Json<ExecuteResponse>, StatusCode> {
    // Set user input in session context
    let session = state.session_storage.get(&session_id).await?;
    session.context.set("user_input", request.content).await;
    state.session_storage.save(session).await?;
    
    // Execute one workflow step
    let result = state.flow_runner.run(&session_id).await?;
    
    Ok(Json(ExecuteResponse {
        session_id,
        response: result.response,
        status: format!("{:?}", result.status),
    }))
}

POST /execute
{
  "content": "I had a car accident and need to file a claim"
}

Response:
{
  "session_id": "uuid-here",
  "response": "I'm sorry to hear about your accident. I'm here to help you file your claim...",
  "status": "Continue"
}

POST /execute  
{
  "session_id": "uuid-here",
  "content": "It happened yesterday on Main Street"
}

GET /session/{session_id}

Response:
{
  "id": "uuid-here",
  "current_task_id": "InsuranceTypeClassifierTask",
  "context": {...},
  "status_message": "Determining insurance type based on claim description"
}

pub struct TaskResult {
    pub response: Option<String>,      // Response to user
    pub next_action: NextAction,       // What to do next
}

pub enum NextAction {
    Continue,           // Move to next task, return control to caller
    ContinueAndExecute, // Move to next task and execute it immediately  
    WaitForInput,       // Pause and wait for more user input
    End,               // Complete the workflow
}

// Store typed data
context.set("user_name", "Alice".to_string()).await;
context.set("claim_amount", 1500.0).await;

// Retrieve typed data
let name: String = context.get_sync("user_name").unwrap();
let amount: f64 = context.get("claim_amount").await.unwrap();

// Manage conversations
context.add_user_message("Hello".to_string()).await;
context.add_assistant_message("Hi there!".to_string()).await;

let graph = GraphBuilder::new("workflow_name")
    .add_task(task1)           // Add tasks
    .add_task(task2)
    .add_edge("task1", "task2") // Linear connections
    .add_conditional_edge(      // Conditional routing
        "task2", 
        |ctx| condition_check(ctx),
        "task3_yes",   // yes branch
        "task3_no",    // else branch
    )
    .build();

// High-level: FlowRunner handles session loading/saving automatically
let runner = FlowRunner::new(graph.clone(), session_storage.clone());
let result = runner.run(&session_id).await?;

// Low-level: Manual session management for custom control
let mut session = session_storage.get(&session_id).await?.unwrap();
let result = graph.execute_session(&mut session).await?;
session_storage.save(session).await?;