klieo-flows 0.8.0

//! Iterative flow. Repeats `body` until predicate returns `Done` or `max_iters` is hit.

use crate::error::FlowError;
use crate::flow::Flow;
use async_trait::async_trait;
use klieo_core::agent::AgentContext;
use serde_json::Value;
use std::sync::Arc;

/// Predicate verdict for [`LoopFlow`].
#[derive(Debug, Clone)]
pub enum LoopVerdict {
    /// Run another iteration.
    Continue,
    /// Terminate; return the most recent output.
    Done,
    /// Terminate with a [`FlowError::Loop`] containing this reason.
    Failed(String),
}

const DEFAULT_MAX_ITERS: u32 = 16;

/// Loop a single body flow until the predicate returns `Done`.
pub struct LoopFlow {
    name: String,
    body: Arc<dyn Flow>,
    predicate: Arc<dyn Fn(&Value) -> LoopVerdict + Send + Sync>,
    max_iters: u32,
}

impl LoopFlow {
    /// New loop. Default predicate returns `Done` immediately (single iter)
    /// so the caller is forced to call [`LoopFlow::until`].
    pub fn new(name: impl Into<String>, body: Arc<dyn Flow>) -> Self {
        Self {
            name: name.into(),
            body,
            predicate: Arc::new(|_| LoopVerdict::Done),
            max_iters: DEFAULT_MAX_ITERS,
        }
    }

    /// Replace the termination predicate.
    pub fn until<F>(mut self, predicate: F) -> Self
    where
        F: Fn(&Value) -> LoopVerdict + Send + Sync + 'static,
    {
        self.predicate = Arc::new(predicate);
        self
    }

    /// Override the default cap (16). The cap bounds the number of body
    /// invocations — the body runs at most `max_iters` times.
    pub fn with_max_iters(mut self, n: u32) -> Self {
        self.max_iters = n;
        self
    }
}

#[async_trait]
impl Flow for LoopFlow {
    fn name(&self) -> &str {
        &self.name
    }

    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        let span = tracing::info_span!(
            "loop_flow.run",
            flow = %self.name,
            max_iters = self.max_iters
        );
        let _guard = span.enter();

        let mut current = input;
        for iter in 0..self.max_iters {
            current = self.body.run(ctx.clone(), current).await?;
            match (self.predicate)(&current) {
                LoopVerdict::Continue => continue,
                LoopVerdict::Done => return Ok(current),
                LoopVerdict::Failed(reason) => {
                    return Err(FlowError::Loop {
                        reason,
                        iter: iter + 1,
                    });
                }
            }
        }
        Err(FlowError::Loop {
            reason: "max_iters exceeded".into(),
            iter: self.max_iters,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::test_helpers::ctx;

    struct IncrementFlow;

    #[async_trait]
    impl Flow for IncrementFlow {
        fn name(&self) -> &str {
            "inc"
        }
        async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
            let n = input.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
            Ok(serde_json::json!({ "n": n + 1 }))
        }
    }

    #[tokio::test]
    async fn predicate_done_immediately_runs_once() {
        let f = LoopFlow::new("once", Arc::new(IncrementFlow));
        let out = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 1}));
    }

    #[tokio::test]
    async fn continues_then_done_after_three() {
        let f = LoopFlow::new("3x", Arc::new(IncrementFlow)).until(|v| {
            let n = v.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
            if n >= 4 {
                LoopVerdict::Done
            } else {
                LoopVerdict::Continue
            }
        });
        let out = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 4}));
    }

    #[tokio::test]
    async fn predicate_failed_returns_loop_error() {
        let f = LoopFlow::new("fail", Arc::new(IncrementFlow))
            .until(|_| LoopVerdict::Failed("nope".into()));
        let err = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Loop { reason, iter } => {
                assert_eq!(reason, "nope");
                assert_eq!(iter, 1);
            }
            other => panic!("expected Loop, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn max_iters_exceeded() {
        let f = LoopFlow::new("inf", Arc::new(IncrementFlow))
            .until(|_| LoopVerdict::Continue)
            .with_max_iters(2);
        let err = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Loop { reason, iter } => {
                assert_eq!(reason, "max_iters exceeded");
                assert_eq!(iter, 2);
            }
            other => panic!("expected Loop, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn output_threads_to_next_iteration() {
        let f = LoopFlow::new("thread", Arc::new(IncrementFlow)).until(|v| {
            let n = v.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
            if n >= 5 {
                LoopVerdict::Done
            } else {
                LoopVerdict::Continue
            }
        });
        let out = f.run(ctx(), serde_json::json!({"n": 2})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 5}));
    }
}