klieo-flows 3.4.0

//! Ordered pipeline of flows; output of step N feeds input of step N+1.

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;

/// Ordered pipeline. Each step's output is threaded as the next step's
/// input. Empty pipelines return their input unchanged (identity).
/// First error short-circuits.
pub struct SequentialFlow {
    name: String,
    steps: Vec<Arc<dyn Flow>>,
}

impl SequentialFlow {
    /// New empty pipeline with the given name.
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            steps: Vec::new(),
        }
    }

    /// Append a step. Builder-style: returns self.
    pub fn step(mut self, flow: Arc<dyn Flow>) -> Self {
        self.steps.push(flow);
        self
    }

    /// Convenience builder that wraps a concrete [`klieo_core::agent::Agent`]
    /// in [`crate::flow::AgentFlow`] + upcasts to `Arc<dyn Flow>` in one
    /// call.
    ///
    /// Equivalent to:
    /// ```ignore
    /// .step(std::sync::Arc::new(AgentFlow::new(agent)) as std::sync::Arc<dyn Flow>)
    /// ```
    pub fn step_agent<A>(self, agent: A) -> Self
    where
        A: klieo_core::agent::Agent + 'static,
    {
        self.step(std::sync::Arc::new(crate::flow::AgentFlow::new(agent))
            as std::sync::Arc<dyn crate::flow::Flow>)
    }
}

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

    async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        let span = tracing::info_span!("sequential_flow.run", flow = %self.name);
        let _guard = span.enter();
        let mut current = input;
        for (i, step) in self.steps.iter().enumerate() {
            crate::flow::bail_if_cancelled(&ctx)?;
            tracing::debug!(step_index = i, step_name = step.name(), "sequential step");
            current = step.run(ctx.clone(), current).await?;
        }
        Ok(current)
    }
}

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

    struct Boom;
    #[async_trait]
    impl Flow for Boom {
        fn name(&self) -> &str {
            "boom"
        }
        async fn run(&self, _c: AgentContext, _i: Value) -> Result<Value, FlowError> {
            panic!("step must not run under a cancelled context")
        }
    }

    #[tokio::test]
    async fn cancelled_ctx_returns_cancelled_before_any_step() {
        let f = SequentialFlow::new("s").step(Arc::new(Boom));
        let err = f
            .run(cancelled_ctx(), serde_json::json!({}))
            .await
            .unwrap_err();
        assert!(matches!(err, FlowError::Cancelled), "got {err:?}");
    }

    struct IncrementFlow {
        name: String,
    }

    #[async_trait]
    impl Flow for IncrementFlow {
        fn name(&self) -> &str {
            &self.name
        }
        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 }))
        }
    }

    struct FailFlow;

    #[async_trait]
    impl Flow for FailFlow {
        fn name(&self) -> &str {
            "fail"
        }
        async fn run(&self, _ctx: AgentContext, _input: Value) -> Result<Value, FlowError> {
            Err(FlowError::Agent("boom".into()))
        }
    }

    #[tokio::test]
    async fn empty_pipeline_returns_input_unchanged() {
        let f = SequentialFlow::new("empty");
        let out = f.run(ctx(), serde_json::json!({"n": 7})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 7}));
    }

    #[tokio::test]
    async fn three_steps_increment_in_order() {
        let f = SequentialFlow::new("incs")
            .step(Arc::new(IncrementFlow { name: "a".into() }))
            .step(Arc::new(IncrementFlow { name: "b".into() }))
            .step(Arc::new(IncrementFlow { name: "c".into() }));
        let out = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();
        assert_eq!(out, serde_json::json!({"n": 3}));
    }

    #[tokio::test]
    async fn step_2_error_short_circuits() {
        let f = SequentialFlow::new("with_fail")
            .step(Arc::new(IncrementFlow { name: "a".into() }))
            .step(Arc::new(FailFlow))
            .step(Arc::new(IncrementFlow { name: "c".into() }));
        let err = f.run(ctx(), serde_json::json!({"n": 0})).await.unwrap_err();
        match err {
            FlowError::Agent(s) => assert_eq!(s, "boom"),
            other => panic!("expected Agent, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn name_returns_configured_name() {
        let f = SequentialFlow::new("my_pipe");
        assert_eq!(f.name(), "my_pipe");
    }
}

#[cfg(test)]
mod step_agent_tests {
    use super::*;
    use crate::test_helpers::ctx;
    use async_trait::async_trait;
    use klieo_core::agent::{Agent, AgentContext};
    use klieo_core::error::Error as CoreError;
    use klieo_core::llm::ToolDef;
    use serde::{Deserialize, Serialize};

    #[derive(Deserialize, Serialize, PartialEq, Debug)]
    struct EchoIn {
        msg: String,
    }
    #[derive(Deserialize, Serialize, PartialEq, Debug)]
    struct EchoOut {
        echoed: String,
    }

    struct EchoAgent;

    #[async_trait]
    impl Agent for EchoAgent {
        type Input = EchoIn;
        type Output = EchoOut;
        type Error = CoreError;
        fn name(&self) -> &str {
            "echo"
        }
        fn system_prompt(&self) -> &str {
            ""
        }
        fn tools(&self) -> &[ToolDef] {
            &[]
        }
        async fn run(&self, _ctx: AgentContext, input: EchoIn) -> Result<EchoOut, CoreError> {
            Ok(EchoOut { echoed: input.msg })
        }
    }

    #[tokio::test]
    async fn step_agent_chains_typed_agent_without_upcast_boilerplate() {
        let pipeline = SequentialFlow::new("step-agent-smoke").step_agent(EchoAgent);
        let out = pipeline
            .run(ctx(), serde_json::json!({"msg": "hi"}))
            .await
            .expect("run");
        assert_eq!(out, serde_json::json!({"echoed": "hi"}));
    }
}