use anyhow::Result;
use async_trait::async_trait;
use tokio::time::{sleep, Duration};
use sein::{AgentEngine, Node, Tool};
#[derive(Debug)]
pub struct CounterContext {
current: u32,
max: u32,
}
pub struct ThinkingNode;
#[async_trait]
impl Node<CounterContext> for ThinkingNode {
async fn enter(&self, ctx: &mut CounterContext, _tools: &sein::ToolRegistry<CounterContext>) -> Result<()> {
sleep(Duration::from_millis(500)).await;
println!("Thinking... current count: {}/{}", ctx.current, ctx.max);
Ok(())
}
}
pub struct GreetNode;
#[async_trait]
impl Node<CounterContext> for GreetNode {
async fn enter(&self, ctx: &mut CounterContext, tools: &sein::ToolRegistry<CounterContext>) -> Result<()> {
let greet_tool = tools
.get("greet")
.ok_or_else(|| anyhow::anyhow!("Tool 'greet' not found in registry"))?;
greet_tool.execute(ctx).await?;
Ok(())
}
}
pub struct ExitNode;
#[async_trait]
impl Node<CounterContext> for ExitNode {
async fn enter(&self, _ctx: &mut CounterContext, _tools: &sein::ToolRegistry<CounterContext>) -> Result<()> {
println!("Done! Exiting...");
Ok(())
}
}
pub struct GreetTool;
#[async_trait]
impl Tool<CounterContext> for GreetTool {
async fn execute(&self, ctx: &mut CounterContext) -> Result<()> {
sleep(Duration::from_secs(1)).await;
ctx.current += 1;
println!("Hello! This is greeting number {}", ctx.current);
Ok(())
}
fn name(&self) -> &str {
"greet"
}
}
#[tokio::main]
async fn main() -> Result<()> {
let mut engine = AgentEngine::new(CounterContext { current: 0, max: 3 });
engine.add_tools(vec![GreetTool]);
let thinking_node = engine.add_node(ThinkingNode);
let greet_node = engine.add_node(GreetNode);
let exit_node = engine.add_node(ExitNode);
engine.add_transition(thinking_node, greet_node, |ctx| ctx.current < ctx.max)?;
engine.add_transition(thinking_node, exit_node, |ctx| ctx.current >= ctx.max)?;
engine.add_transition(greet_node, thinking_node, |_| true)?;
engine.set_start_node(thinking_node)?;
engine.run().await?;
Ok(())
}