ri_agent_graph/
join.rs

1//! JoinNode for deterministic fan-in merging.
2//!
3//! When parallel branches converge, a [`JoinNode`] provides explicit
4//! merge logic. It reads specified keys from state (set by parallel branches),
5//! applies a merge function, and writes the result to an output key.
6
7use crate::command::NodeOutput;
8use crate::config::GraphConfig;
9use crate::error::AgentGraphError;
10use crate::node::Node;
11use crate::state::AgentState;
12use serde_json::Value;
13
14/// Merge function signature for JoinNode.
15pub type MergeFn = Box<dyn Fn(Vec<(String, Value)>) -> crate::Result<Value> + Send + Sync>;
16
17/// A node that merges results from parallel branches.
18///
19/// After fan-out, parallel branches write their results to known state keys.
20/// The JoinNode reads those keys, applies a merge function, and writes
21/// the merged result to an output key.
22pub struct JoinNode {
23    name: Option<String>,
24    /// State keys to read from parallel branches.
25    input_keys: Vec<String>,
26    /// State key to write the merged result to.
27    output_key: String,
28    /// Merge function: receives `Vec<(key, value)>` and produces the merged value.
29    merge_fn: MergeFn,
30}
31
32impl JoinNode {
33    /// Create a new JoinNode.
34    ///
35    /// - `input_keys`: state keys to collect from parallel branches.
36    /// - `output_key`: state key to write the merged result to.
37    /// - `merge_fn`: function that merges the collected values.
38    pub fn new(
39        input_keys: Vec<String>,
40        output_key: impl Into<String>,
41        merge_fn: impl Fn(Vec<(String, Value)>) -> crate::Result<Value> + Send + Sync + 'static,
42    ) -> Self {
43        Self {
44            name: None,
45            input_keys,
46            output_key: output_key.into(),
47            merge_fn: Box::new(merge_fn),
48        }
49    }
50
51    /// Set a name for this node.
52    pub fn with_name(mut self, name: impl Into<String>) -> Self {
53        self.name = Some(name.into());
54        self
55    }
56
57    /// Convenience: create a JoinNode that collects values into an array.
58    pub fn collect_array(input_keys: Vec<String>, output_key: impl Into<String>) -> Self {
59        Self::new(input_keys, output_key, |values| {
60            let arr: Vec<Value> = values.into_iter().map(|(_, v)| v).collect();
61            Ok(Value::Array(arr))
62        })
63    }
64
65    /// Convenience: create a JoinNode that merges objects (shallow).
66    pub fn merge_objects(input_keys: Vec<String>, output_key: impl Into<String>) -> Self {
67        Self::new(input_keys, output_key, |values| {
68            let mut result = serde_json::Map::new();
69            for (key, value) in values {
70                if let Value::Object(map) = value {
71                    for (k, v) in map {
72                        result.insert(k, v);
73                    }
74                } else {
75                    result.insert(key, value);
76                }
77            }
78            Ok(Value::Object(result))
79        })
80    }
81}
82
83#[async_trait::async_trait]
84impl Node for JoinNode {
85    async fn execute(
86        &self,
87        state: &AgentState,
88        _config: &GraphConfig,
89    ) -> crate::Result<NodeOutput> {
90        // Collect values from input keys
91        let mut inputs = Vec::new();
92        for key in &self.input_keys {
93            let value: Value = state.get_opt::<Value>(key).await?.unwrap_or(Value::Null);
94            inputs.push((key.clone(), value));
95        }
96
97        // Validate that we have at least some non-null inputs
98        let has_data = inputs.iter().any(|(_, v)| !v.is_null());
99        if !has_data {
100            return Err(AgentGraphError::ExecutionError(format!(
101                "JoinNode: no data found for input keys {:?}",
102                self.input_keys
103            )));
104        }
105
106        // Apply merge function
107        let merged = (self.merge_fn)(inputs)?;
108
109        // Write to output key
110        state.set_raw(&self.output_key, merged).await?;
111
112        Ok(NodeOutput::Done)
113    }
114
115    fn name(&self) -> Option<&str> {
116        self.name.as_deref()
117    }
118}
119
120impl std::fmt::Debug for JoinNode {
121    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
122        f.debug_struct("JoinNode")
123            .field("name", &self.name)
124            .field("input_keys", &self.input_keys)
125            .field("output_key", &self.output_key)
126            .finish()
127    }
128}
ri_agent_graph/join.rs

ri_agent_graph/
join.rs
