guardflow 0.2.0

Validators and a retry-until-valid guard for LLM output, in Rust
Documentation
use std::sync::atomic::{AtomicUsize, Ordering};

use async_trait::async_trait;
use graph_flow::{Context, NextAction, Task, TaskResult, error::Result};
use guardflow::Guard;
use guardflow::graphflow::GuardedTask;
use guardflow::validators::MinLength;

struct FlakyTask {
    attempts: AtomicUsize,
    succeeds_on: usize,
}

#[async_trait]
impl Task for FlakyTask {
    fn id(&self) -> &str {
        "flaky"
    }

    async fn run(&self, _context: Context) -> Result<TaskResult> {
        let attempt = self.attempts.fetch_add(1, Ordering::SeqCst) + 1;
        let response = if attempt >= self.succeeds_on {
            "this response is long enough to pass"
        } else {
            "short"
        };
        Ok(TaskResult::new(
            Some(response.to_string()),
            NextAction::Continue,
        ))
    }
}

#[tokio::test]
async fn guarded_task_retries_until_validator_passes() {
    let inner = FlakyTask {
        attempts: AtomicUsize::new(0),
        succeeds_on: 3,
    };
    let guard = Guard::new().with(MinLength(20));
    let guarded = GuardedTask::new(inner, guard).with_max_attempts(5);

    let result = guarded.run(Context::new()).await.unwrap();
    assert_eq!(
        result.response,
        Some("this response is long enough to pass".to_string())
    );
}

#[tokio::test]
async fn guarded_task_gives_up_after_max_attempts() {
    let inner = FlakyTask {
        attempts: AtomicUsize::new(0),
        succeeds_on: 100,
    };
    let guard = Guard::new().with(MinLength(20));
    let guarded = GuardedTask::new(inner, guard).with_max_attempts(2);

    let result = guarded.run(Context::new()).await.unwrap();
    assert_eq!(result.response, Some("short".to_string()));
}