klieo-flows 0.8.0

//! 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<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 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();
        (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");
    }
}