klieo-flows 3.3.0

//! End-to-end smoke test: SequentialFlow > ParallelFlow > LoopFlow > AgentFlow.

use async_trait::async_trait;
use klieo_core::agent::{Agent, AgentContext};
use klieo_core::error::Error as CoreError;
use klieo_core::llm::ToolDef;
use klieo_core::test_utils::fake_context;
use klieo_flows::{
    AgentFlow, Flow, FlowError, LoopFlow, LoopVerdict, ParallelFlow, SequentialFlow,
};
use serde_json::Value;
use std::sync::Arc;

struct TaggingAgent {
    tag: String,
}

#[async_trait]
impl Agent for TaggingAgent {
    type Input = Value;
    type Output = Value;
    type Error = CoreError;
    fn name(&self) -> &str {
        &self.tag
    }
    fn system_prompt(&self) -> &str {
        ""
    }
    fn tools(&self) -> &[ToolDef] {
        &[]
    }
    async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, CoreError> {
        let mut v = input;
        if let Value::Object(ref mut m) = v {
            m.insert("touched_by".into(), Value::String(self.tag.clone()));
        }
        Ok(v)
    }
}

struct EchoFlow {
    name: String,
}

#[async_trait]
impl Flow for EchoFlow {
    fn name(&self) -> &str {
        &self.name
    }
    async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        Ok(input)
    }
}

struct AddOneFlow;

#[async_trait]
impl Flow for AddOneFlow {
    fn name(&self) -> &str {
        "add_one"
    }
    async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
        // After ParallelFlow runs, input is `{ "left": <orig>, "right": <orig> }`.
        // After Loop's first iteration, the previous body output threads back
        // in: `{ "n": <prev+1> }`. Pull `n` from either shape so the body can
        // run repeatedly with stable JSON shape between iterations.
        let n = input
            .get("left")
            .and_then(|v| v.get("n"))
            .and_then(|v| v.as_u64())
            .or_else(|| input.get("n").and_then(|v| v.as_u64()))
            .unwrap_or(0);
        Ok(serde_json::json!({ "n": n + 1 }))
    }
}

fn ctx() -> AgentContext {
    fake_context("integration")
}

#[tokio::test]
async fn end_to_end_sequential_parallel_loop_agent() {
    // Sequential pipeline:
    //   1. Parallel { left: echo, right: echo } -> {"left": input, "right": input}
    //   2. Loop until `n >= 1` over AddOneFlow -> {"n": 1} (one iter is enough)
    //   3. AgentFlow over TaggingAgent -> adds "touched_by"

    let parallel = Arc::new(
        ParallelFlow::new("split")
            .branch("left", Arc::new(EchoFlow { name: "l".into() }))
            .branch("right", Arc::new(EchoFlow { name: "r".into() })),
    );

    let loop_step = Arc::new(
        LoopFlow::new("inc_until_one", Arc::new(AddOneFlow))
            .until(|v| {
                let n = v.get("n").and_then(|v| v.as_u64()).unwrap_or(0);
                if n >= 1 {
                    LoopVerdict::Done
                } else {
                    LoopVerdict::Continue
                }
            })
            .with_max_iters(4),
    );

    let agent_step = Arc::new(AgentFlow::new(TaggingAgent {
        tag: "tagger".into(),
    }));

    let pipe = SequentialFlow::new("pipe")
        .step(parallel)
        .step(loop_step)
        .step(agent_step);

    let out = pipe.run(ctx(), serde_json::json!({"n": 0})).await.unwrap();

    // After Loop: {"n": 1}; after AgentFlow over TaggingAgent: {"n": 1, "touched_by": "tagger"}.
    assert_eq!(out, serde_json::json!({"n": 1, "touched_by": "tagger"}));
}