Skip to main content

ri_agent_graph/
node.rs

1use crate::command::NodeOutput;
2use crate::config::GraphConfig;
3use crate::error::Result;
4use crate::state::AgentState;
5use async_trait::async_trait;
6use std::future::Future;
7use std::pin::Pin;
8
9/// A node in the agent graph.
10/// Nodes perform actions and can modify the shared state.
11#[async_trait]
12pub trait Node: Send + Sync {
13    /// Execute this node
14    async fn execute(&self, state: &AgentState, config: &GraphConfig) -> Result<NodeOutput>;
15
16    /// Optional: Get a name for this node (for debugging)
17    fn name(&self) -> Option<&str> {
18        None
19    }
20}
21
22/// Helper to create a node from an async function
23pub struct FnNode<F>
24where
25    F: Fn(&AgentState, &GraphConfig) -> Pin<Box<dyn Future<Output = Result<NodeOutput>> + Send>>
26        + Send
27        + Sync,
28{
29    func: F,
30    name: Option<String>,
31}
32
33impl<F> FnNode<F>
34where
35    F: Fn(&AgentState, &GraphConfig) -> Pin<Box<dyn Future<Output = Result<NodeOutput>> + Send>>
36        + Send
37        + Sync,
38{
39    pub fn new(func: F) -> Self {
40        Self { func, name: None }
41    }
42
43    pub fn with_name(mut self, name: impl Into<String>) -> Self {
44        self.name = Some(name.into());
45        self
46    }
47}
48
49#[async_trait]
50impl<F> Node for FnNode<F>
51where
52    F: Fn(&AgentState, &GraphConfig) -> Pin<Box<dyn Future<Output = Result<NodeOutput>> + Send>>
53        + Send
54        + Sync,
55{
56    async fn execute(&self, state: &AgentState, config: &GraphConfig) -> Result<NodeOutput> {
57        (self.func)(state, config).await
58    }
59
60    fn name(&self) -> Option<&str> {
61        self.name.as_deref()
62    }
63}
64
65/// Helper macro to create a node from an async closure.
66///
67/// # Forms
68///
69/// ```ignore
70/// // Basic form (backward compatible) - body returns Result<()>
71/// node!(|state| async move {
72///     state.set("key", "value").await?;
73///     Ok(())
74/// })
75///
76/// // Named form - body returns Result<()>
77/// node!("my_node", |state| async move {
78///     state.set("key", "value").await?;
79///     Ok(())
80/// })
81///
82/// // With config - body returns Result<NodeOutput> or Result<()>
83/// node!(|state, config| async move {
84///     state.set("key", "value").await?;
85///     Ok(NodeOutput::Done)
86/// })
87///
88/// // Named with config
89/// node!("my_node", |state, config| async move {
90///     state.set("key", "value").await?;
91///     Ok(())
92/// })
93/// ```
94#[macro_export]
95macro_rules! node {
96    // Form 1: |state| - backward compatible, body returns Result<impl Into<NodeOutput>>
97    (|$state:ident| async move $body:block) => {
98        Box::new($crate::node::FnNode::new(
99            |__state: &$crate::state::AgentState, __config: &$crate::config::GraphConfig| {
100                let $state = __state.clone();
101                let _ = __config;
102                Box::pin(async move {
103                    let __result = (|| async move { $body })().await;
104                    __result.map(::std::convert::Into::into)
105                })
106            },
107        ))
108    };
109    // Form 2: named |state|
110    ($name:expr, |$state:ident| async move $body:block) => {
111        Box::new(
112            $crate::node::FnNode::new(
113                |__state: &$crate::state::AgentState, __config: &$crate::config::GraphConfig| {
114                    let $state = __state.clone();
115                    let _ = __config;
116                    Box::pin(async move {
117                        let __result = (|| async move { $body })().await;
118                        __result.map(::std::convert::Into::into)
119                    })
120                },
121            )
122            .with_name($name),
123        )
124    };
125    // Form 3: |state, config| - has access to config
126    (|$state:ident, $config:ident| async move $body:block) => {
127        Box::new($crate::node::FnNode::new(
128            |__state: &$crate::state::AgentState, __config: &$crate::config::GraphConfig| {
129                let $state = __state.clone();
130                let $config = __config.clone();
131                Box::pin(async move {
132                    let __result = (|| async move { $body })().await;
133                    __result.map(::std::convert::Into::into)
134                })
135            },
136        ))
137    };
138    // Form 4: named |state, config|
139    ($name:expr, |$state:ident, $config:ident| async move $body:block) => {
140        Box::new(
141            $crate::node::FnNode::new(
142                |__state: &$crate::state::AgentState, __config: &$crate::config::GraphConfig| {
143                    let $state = __state.clone();
144                    let $config = __config.clone();
145                    Box::pin(async move {
146                        let __result = (|| async move { $body })().await;
147                        __result.map(::std::convert::Into::into)
148                    })
149                },
150            )
151            .with_name($name),
152        )
153    };
154}