Skip to main content

ri_agent_graph/
reducer.rs

1use crate::error::{AgentGraphError, Result};
2use serde_json::Value;
3
4/// A reducer combines old and new values during state updates.
5/// Critical for correctness during parallel execution.
6pub trait Reducer: Send + Sync {
7    /// Combine the current value with the new value
8    fn reduce(&self, current: &Value, update: &Value) -> Result<Value>;
9}
10
11/// Default reducer: the new value replaces the old value.
12pub struct LastWriteWins;
13
14impl Reducer for LastWriteWins {
15    fn reduce(&self, _current: &Value, update: &Value) -> Result<Value> {
16        Ok(update.clone())
17    }
18}
19
20/// Append reducer: appends items from the update array to the current array.
21pub struct AppendReducer;
22
23impl Reducer for AppendReducer {
24    fn reduce(&self, current: &Value, update: &Value) -> Result<Value> {
25        match (current, update) {
26            (Value::Null, val) => {
27                // First write: if update is array, use it directly; otherwise wrap in array
28                if let Value::Array(_) = val {
29                    Ok(val.clone())
30                } else {
31                    Ok(Value::Array(vec![val.clone()]))
32                }
33            }
34            (Value::Array(curr), Value::Array(upd)) => {
35                let mut result = curr.clone();
36                result.extend(upd.iter().cloned());
37                Ok(Value::Array(result))
38            }
39            (Value::Array(curr), val) => {
40                let mut result = curr.clone();
41                result.push(val.clone());
42                Ok(Value::Array(result))
43            }
44            (_, Value::Array(upd)) => {
45                let mut result = vec![current.clone()];
46                result.extend(upd.iter().cloned());
47                Ok(Value::Array(result))
48            }
49            _ => Ok(Value::Array(vec![current.clone(), update.clone()])),
50        }
51    }
52}
53
54/// Add reducer: adds numeric values together.
55pub struct AddReducer;
56
57impl Reducer for AddReducer {
58    fn reduce(&self, current: &Value, update: &Value) -> Result<Value> {
59        let a_f = match current {
60            Value::Number(n) => n.as_f64().unwrap_or(0.0),
61            Value::Null => 0.0,
62            _ => {
63                return Err(AgentGraphError::StateError(
64                    "AddReducer: current value must be a number".to_string(),
65                ))
66            }
67        };
68        let b_f = match update {
69            Value::Number(n) => n.as_f64().ok_or_else(|| {
70                AgentGraphError::StateError("AddReducer: cannot convert update to f64".to_string())
71            })?,
72            _ => {
73                return Err(AgentGraphError::StateError(
74                    "AddReducer: update value must be a number".to_string(),
75                ))
76            }
77        };
78        Ok(serde_json::json!(a_f + b_f))
79    }
80}
81
82/// Merge reducer: deep-merges JSON objects.
83pub struct MergeReducer;
84
85impl Reducer for MergeReducer {
86    fn reduce(&self, current: &Value, update: &Value) -> Result<Value> {
87        match (current, update) {
88            (Value::Object(curr), Value::Object(upd)) => {
89                let mut result = curr.clone();
90                for (k, v) in upd {
91                    if let Some(existing) = result.get(k) {
92                        if existing.is_object() && v.is_object() {
93                            result.insert(k.clone(), self.reduce(existing, v)?);
94                        } else {
95                            result.insert(k.clone(), v.clone());
96                        }
97                    } else {
98                        result.insert(k.clone(), v.clone());
99                    }
100                }
101                Ok(Value::Object(result))
102            }
103            _ => Ok(update.clone()),
104        }
105    }
106}
107
108/// Closure-based reducer for custom logic.
109pub struct FnReducer<F>
110where
111    F: Fn(&Value, &Value) -> Result<Value> + Send + Sync,
112{
113    func: F,
114}
115
116impl<F> FnReducer<F>
117where
118    F: Fn(&Value, &Value) -> Result<Value> + Send + Sync,
119{
120    pub fn new(func: F) -> Self {
121        Self { func }
122    }
123}
124
125impl<F> Reducer for FnReducer<F>
126where
127    F: Fn(&Value, &Value) -> Result<Value> + Send + Sync,
128{
129    fn reduce(&self, current: &Value, update: &Value) -> Result<Value> {
130        (self.func)(current, update)
131    }
132}