klieo_flows/

typed.rs

//! Typed sequential flow builder.
//!
//! [`SequentialFlow`] round-trips through `serde_json::Value` at every
//! step boundary so it can compose heterogeneous `Arc<dyn Flow>`
//! references. For pipelines where every stage is in-process Rust
//! code, the round-trip is unnecessary overhead and forces every
//! stage type to implement `Serialize` + `Deserialize`.
//!
//! [`TypedSequentialFlow`] threads typed values between stages in
//! memory. Each `.then` appends a stage whose input type is the
//! previous stage's output type. The composed flow is itself a
//! `TypedSequentialFlow<I, O>` where `I` is the first stage's input
//! and `O` is the last stage's output.
//!
//! Use [`SequentialFlow`](crate::SequentialFlow) when the pipeline
//! must accept heterogeneous step types via `Arc<dyn Flow>` (the
//! JSON-marshalled type-erased path).
//!
//! # Example
//!
//! ```no_run
//! use klieo_flows::TypedSequentialFlow;
//! use klieo_flows::FlowError;
//! use klieo_core::agent::AgentContext;
//!
//! struct Ticket { subject: String }
//! struct Classified { ticket: Ticket, category: String }
//! struct Reply(String);
//!
//! # async fn _example(ctx: AgentContext, t: Ticket) -> Result<Reply, FlowError> {
//! let flow = TypedSequentialFlow::new("triage", |_ctx, t: Ticket| async {
//!         Ok::<_, FlowError>(Classified { ticket: t, category: "billing".into() })
//!     })
//!     .then(|_ctx, c: Classified| async move {
//!         Ok::<_, FlowError>(Reply(format!("re: {}", c.ticket.subject)))
//!     });
//!
//! let reply: Reply = flow.run(ctx, t).await?;
//! # Ok(reply)
//! # }
//! ```
//!
//! [`SequentialFlow`]: crate::SequentialFlow

use crate::error::FlowError;
use klieo_core::agent::AgentContext;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;

type StageFn<I, O> = dyn Fn(AgentContext, I) -> Pin<Box<dyn Future<Output = Result<O, FlowError>> + Send>>
    + Send
    + Sync;

/// Pipeline of typed transformations.
///
/// Constructed with [`new`](Self::new); extended with
/// [`then`](Self::then). Drive with [`run`](Self::run).
pub struct TypedSequentialFlow<I, O> {
    name: String,
    stage: Arc<StageFn<I, O>>,
}

impl TypedSequentialFlow<(), ()> {
    /// Start a typed pipeline with a named [`klieo_core::agent::Agent`] as the first stage.
    ///
    /// Equivalent to [`TypedSequentialFlow::new`] with an auto-wrapping
    /// closure, but takes any concrete [`klieo_core::agent::Agent`] directly
    /// — no manual `map_err` needed.
    ///
    /// # Example
    /// ```ignore
    /// use klieo_flows::TypedSequentialFlow;
    ///
    /// // `Classifier` implements `klieo_core::agent::Agent`.
    /// let flow = TypedSequentialFlow::start_with("pipeline", Classifier)
    ///     .then_agent(Retriever)
    ///     .then_agent(Drafter);
    /// ```
    pub fn start_with<A>(
        name: impl Into<String>,
        agent: A,
    ) -> TypedSequentialFlow<A::Input, A::Output>
    where
        A: klieo_core::agent::Agent + Send + Sync + 'static,
        A::Input: Send + 'static,
        A::Output: Send + 'static,
        A::Error: std::error::Error + Send + Sync + 'static,
    {
        let agent = Arc::new(agent);
        TypedSequentialFlow::new(name, move |ctx, input| {
            let agent = Arc::clone(&agent);
            async move {
                agent
                    .run(ctx, input)
                    .await
                    .map_err(|e| FlowError::Agent(e.to_string()))
            }
        })
    }
}

impl<I, O> TypedSequentialFlow<I, O>
where
    I: Send + 'static,
    O: Send + 'static,
{
    /// Start a typed pipeline with the given name and first stage.
    pub fn new<F, Fut>(name: impl Into<String>, stage: F) -> Self
    where
        F: Fn(AgentContext, I) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<O, FlowError>> + Send + 'static,
    {
        let boxed: Arc<StageFn<I, O>> = Arc::new(move |ctx, input| Box::pin(stage(ctx, input)));
        Self {
            name: name.into(),
            stage: boxed,
        }
    }

    /// Append a stage from a concrete [`klieo_core::agent::Agent`].
    ///
    /// The agent's [`Input`](klieo_core::agent::Agent::Input) must match
    /// this flow's current output type `O`. Returns a new
    /// `TypedSequentialFlow<I, A::Output>`.
    pub fn then_agent<A>(self, agent: A) -> TypedSequentialFlow<I, A::Output>
    where
        A: klieo_core::agent::Agent<Input = O> + Send + Sync + 'static,
        A::Output: Send + 'static,
        A::Error: std::error::Error + Send + Sync + 'static,
    {
        let agent = Arc::new(agent);
        self.then(move |ctx, input| {
            let agent = Arc::clone(&agent);
            async move {
                agent
                    .run(ctx, input)
                    .await
                    .map_err(|e| FlowError::Agent(e.to_string()))
            }
        })
    }

    /// Append a stage that consumes the previous output type `O` and
    /// produces a new output type `P`.
    pub fn then<F, Fut, P>(self, stage: F) -> TypedSequentialFlow<I, P>
    where
        F: Fn(AgentContext, O) -> Fut + Send + Sync + 'static,
        Fut: Future<Output = Result<P, FlowError>> + Send + 'static,
        P: Send + 'static,
    {
        let prev = self.stage;
        let next = Arc::new(stage);
        let combined: Arc<StageFn<I, P>> = Arc::new(move |ctx: AgentContext, input: I| {
            let prev = Arc::clone(&prev);
            let next = Arc::clone(&next);
            Box::pin(async move {
                let intermediate = prev(ctx.clone(), input).await?;
                next(ctx, intermediate).await
            })
        });
        TypedSequentialFlow {
            name: self.name,
            stage: combined,
        }
    }

    /// Run the pipeline end-to-end. The composed stage chain is wrapped
    /// in a single `typed_sequential_flow.run` tracing span.
    pub async fn run(&self, ctx: AgentContext, input: I) -> Result<O, FlowError> {
        let span = tracing::info_span!("typed_sequential_flow.run", flow = %self.name);
        let _guard = span.enter();
        crate::flow::bail_if_cancelled(&ctx)?;
        (self.stage)(ctx, input).await
    }

    /// Configured pipeline name.
    pub fn name(&self) -> &str {
        &self.name
    }
}

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

    #[derive(Debug, PartialEq)]
    struct Ticket {
        subject: String,
    }

    #[derive(Debug, PartialEq)]
    struct Classified {
        ticket: Ticket,
        category: String,
    }

    #[derive(Debug, PartialEq)]
    struct Reply(String);

    #[tokio::test]
    async fn single_stage_pipeline_runs() {
        let flow: TypedSequentialFlow<Ticket, Classified> =
            TypedSequentialFlow::new("classify", |_ctx, t: Ticket| async move {
                Ok(Classified {
                    ticket: t,
                    category: "billing".into(),
                })
            });
        let out = flow
            .run(
                ctx(),
                Ticket {
                    subject: "refund".into(),
                },
            )
            .await
            .unwrap();
        assert_eq!(out.category, "billing");
        assert_eq!(out.ticket.subject, "refund");
    }

    #[tokio::test]
    async fn two_stages_compose_typed_values() {
        let flow = TypedSequentialFlow::new("triage", |_ctx, t: Ticket| async move {
            Ok::<_, FlowError>(Classified {
                ticket: t,
                category: "billing".into(),
            })
        })
        .then(|_ctx, c: Classified| async move {
            Ok::<_, FlowError>(Reply(format!("re[{}]: {}", c.category, c.ticket.subject)))
        });
        let reply = flow
            .run(
                ctx(),
                Ticket {
                    subject: "refund".into(),
                },
            )
            .await
            .unwrap();
        assert_eq!(reply, Reply("re[billing]: refund".into()));
    }

    #[tokio::test]
    async fn first_stage_error_short_circuits() {
        let flow = TypedSequentialFlow::new("triage", |_ctx, _t: Ticket| async {
            Err::<Classified, _>(FlowError::Agent("first failed".into()))
        })
        .then(|_ctx, _c: Classified| async { Ok::<_, FlowError>(Reply("never reached".into())) });
        let err = flow
            .run(
                ctx(),
                Ticket {
                    subject: "x".into(),
                },
            )
            .await
            .unwrap_err();
        match err {
            FlowError::Agent(s) => assert_eq!(s, "first failed"),
            other => panic!("expected Agent, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn second_stage_error_propagates() {
        let flow = TypedSequentialFlow::new("triage", |_ctx, t: Ticket| async move {
            Ok::<_, FlowError>(Classified {
                ticket: t,
                category: "billing".into(),
            })
        })
        .then(|_ctx, _c: Classified| async {
            Err::<Reply, _>(FlowError::Agent("second failed".into()))
        });
        let err = flow
            .run(
                ctx(),
                Ticket {
                    subject: "x".into(),
                },
            )
            .await
            .unwrap_err();
        match err {
            FlowError::Agent(s) => assert_eq!(s, "second failed"),
            other => panic!("expected Agent, got {other:?}"),
        }
    }

    #[tokio::test]
    async fn three_stages_thread_in_order() {
        let flow = TypedSequentialFlow::new("counter", |_ctx, n: u32| async move {
            Ok::<_, FlowError>(n + 1)
        })
        .then(|_ctx, n: u32| async move { Ok::<_, FlowError>(n * 2) })
        .then(|_ctx, n: u32| async move { Ok::<_, FlowError>(n.to_string()) });
        // ((0 + 1) * 2).to_string() == "2"
        let out = flow.run(ctx(), 0_u32).await.unwrap();
        assert_eq!(out, "2");
    }

    #[tokio::test]
    async fn name_returns_configured_name() {
        let flow: TypedSequentialFlow<u32, u32> =
            TypedSequentialFlow::new("mine", |_ctx, n: u32| async move { Ok(n) });
        assert_eq!(flow.name(), "mine");
    }

    #[tokio::test]
    async fn cancelled_ctx_returns_cancelled_before_stage() {
        use crate::test_helpers::cancelled_ctx;
        let flow = TypedSequentialFlow::new("c", |_ctx, _n: u32| async move {
            panic!("stage must not run under a cancelled context");
            #[allow(unreachable_code)]
            Ok::<u32, FlowError>(0)
        });
        let err = flow.run(cancelled_ctx(), 0_u32).await.unwrap_err();
        assert!(matches!(err, FlowError::Cancelled), "got {err:?}");
    }

    struct Double;

    #[async_trait::async_trait]
    impl klieo_core::agent::Agent for Double {
        type Input = i32;
        type Output = i32;
        type Error = klieo_core::Error;

        fn name(&self) -> &str {
            "double"
        }
        fn system_prompt(&self) -> &str {
            ""
        }
        fn tools(&self) -> &[klieo_core::llm::ToolDef] {
            &[]
        }
        async fn run(
            &self,
            _ctx: klieo_core::agent::AgentContext,
            n: i32,
        ) -> Result<i32, klieo_core::Error> {
            Ok(n * 2)
        }
    }

    struct AddTen;

    #[async_trait::async_trait]
    impl klieo_core::agent::Agent for AddTen {
        type Input = i32;
        type Output = i32;
        type Error = klieo_core::Error;

        fn name(&self) -> &str {
            "add-ten"
        }
        fn system_prompt(&self) -> &str {
            ""
        }
        fn tools(&self) -> &[klieo_core::llm::ToolDef] {
            &[]
        }
        async fn run(
            &self,
            _ctx: klieo_core::agent::AgentContext,
            n: i32,
        ) -> Result<i32, klieo_core::Error> {
            Ok(n + 10)
        }
    }

    #[tokio::test]
    async fn start_with_and_then_agent_compose_correctly() {
        let flow = TypedSequentialFlow::start_with("test", Double).then_agent(AddTen);
        let result = flow.run(ctx(), 3).await.unwrap();
        assert_eq!(result, 16, "3 * 2 = 6, + 10 = 16");
    }

    #[tokio::test]
    async fn start_with_sets_name() {
        let flow: TypedSequentialFlow<i32, i32> =
            TypedSequentialFlow::start_with("my-pipeline", Double);
        assert_eq!(flow.name(), "my-pipeline");
    }
}