use tinyagents::graph::{ClosureStateReducer, adapter_subgraph_node, shared_subgraph_node};
use tinyagents::{Command, CompiledGraph, GraphBuilder, NodeContext, NodeResult, Result};
#[derive(Clone, Debug, Default)]
struct Pipe {
log: Vec<String>,
total: i32,
}
#[derive(Clone, Debug, Default)]
struct Patch {
log: Vec<String>,
add: i32,
}
fn grandchild() -> CompiledGraph<i32, i32> {
GraphBuilder::<i32, i32>::overwrite()
.add_node("g_add", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 1))
})
.set_entry("g_add")
.set_finish("g_add")
.compile()
.expect("grandchild compiles")
}
fn enrich_child() -> CompiledGraph<i32, i32> {
GraphBuilder::<i32, i32>::overwrite()
.add_node("c_add", |s: i32, _c: NodeContext| async move {
Ok(NodeResult::Update(s + 5))
})
.add_node("grandchild", shared_subgraph_node(grandchild()))
.set_entry("c_add")
.add_edge("c_add", "grandchild")
.set_finish("grandchild")
.compile()
.expect("enrich child compiles")
}
#[tokio::main]
async fn main() -> Result<()> {
let graph = GraphBuilder::<Pipe, Patch>::new()
.set_reducer(ClosureStateReducer::new(|mut s: Pipe, p: Patch| {
s.log.extend(p.log);
s.total += p.add;
Ok(s)
}))
.with_parallel(true)
.add_node("dispatch", |_s: Pipe, _c: NodeContext| async move {
Ok(NodeResult::Command(
Command::goto(["branch_sub", "branch_plain"]).with_update(Patch {
log: vec!["dispatch".to_string()],
add: 0,
}),
))
})
.mark_command_routing("dispatch")
.add_node(
"branch_sub",
adapter_subgraph_node(
enrich_child(),
|p: &Pipe| p.total,
|_p: &Pipe, child_value: i32| Patch {
log: vec![format!("branch_sub={child_value}")],
add: child_value,
},
),
)
.add_node("branch_plain", |_s: Pipe, _c: NodeContext| async move {
Ok(NodeResult::Update(Patch {
log: vec!["branch_plain".to_string()],
add: 3,
}))
})
.add_node("join", |_s: Pipe, _c: NodeContext| async move {
Ok(NodeResult::Update(Patch {
log: vec!["join".to_string()],
add: 0,
}))
})
.set_entry("dispatch")
.add_edge("branch_sub", "join")
.add_edge("branch_plain", "join")
.set_finish("join")
.compile()?;
let run = graph.run(Pipe::default()).await?;
let path: Vec<&str> = run.visited.iter().map(|n| n.as_str()).collect();
println!("path : {}", path.join(" -> "));
println!("log : {}", run.state.log.join(" | "));
println!("total : {}", run.state.total); println!("steps : {}", run.steps);
Ok(())
}