llm_optimizer_processor/aggregation/
sum.rs

1use super::trait_::{Aggregator, ToF64};
2
3use serde::{Deserialize, Serialize};
4
5/// Accumulator for sum aggregation
6#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
7pub struct SumAccumulator {
8    sum: f64,
9    count: u64,
10}
11
12/// Sum aggregator - computes the sum of all values
13///
14/// This aggregator computes the running sum of all values processed.
15/// Useful for tracking total latency, total cost, total tokens, etc.
16///
17/// # Examples
18///
19/// ```
20/// use processor::aggregation::{Aggregator, SumAggregator};
21///
22/// let mut agg = SumAggregator::new();
23/// agg.update(1.0).unwrap();
24/// agg.update(2.0).unwrap();
25/// agg.update(3.0).unwrap();
26///
27/// assert_eq!(agg.finalize().unwrap(), 6.0);
28/// ```
29#[derive(Debug, Clone)]
30pub struct SumAggregator<T> {
31    sum: f64,
32    count: u64,
33    _phantom: std::marker::PhantomData<T>,
34}
35
36impl<T: ToF64 + Clone> SumAggregator<T> {
37    /// Create a new sum aggregator
38    pub fn new() -> Self {
39        Self {
40            sum: 0.0,
41            count: 0,
42            _phantom: std::marker::PhantomData,
43        }
44    }
45}
46
47impl<T: ToF64 + Clone> Default for SumAggregator<T> {
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53impl<T: ToF64 + Clone + Send + Sync + std::fmt::Debug> Aggregator for SumAggregator<T> {
54    type Input = T;
55    type Output = f64;
56    type Accumulator = SumAccumulator;
57
58    fn new() -> Self {
59        Self::new()
60    }
61
62    fn update(&mut self, value: <Self as Aggregator>::Input) -> anyhow::Result<()> {
63        self.sum += value.to_f64();
64        self.count += 1;
65        Ok(())
66    }
67
68    fn finalize(&self) -> anyhow::Result<<Self as Aggregator>::Output> {
69        Ok(self.sum)
70    }
71
72    fn accumulator(&self) -> <Self as Aggregator>::Accumulator {
73        SumAccumulator {
74            sum: self.sum,
75            count: self.count,
76        }
77    }
78
79    fn merge(&mut self, other: <Self as Aggregator>::Accumulator) -> anyhow::Result<()> {
80        self.sum += other.sum;
81        self.count += other.count;
82        Ok(())
83    }
84
85    fn reset(&mut self) {
86        self.sum = 0.0;
87        self.count = 0;
88    }
89
90    fn count(&self) -> u64 {
91        self.count
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn test_sum_basic() {
101        let mut agg = SumAggregator::<f64>::new();
102        assert_eq!(agg.count(), 0);
103        assert!(agg.is_empty());
104
105        agg.update(1.0).unwrap();
106        agg.update(2.0).unwrap();
107        agg.update(3.0).unwrap();
108
109        assert_eq!(agg.count(), 3);
110        assert_eq!(agg.finalize().unwrap(), 6.0);
111    }
112
113    #[test]
114    fn test_sum_batch() {
115        let mut agg = SumAggregator::<f64>::new();
116        agg.update_batch(&[1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
117
118        assert_eq!(agg.count(), 5);
119        assert_eq!(agg.finalize().unwrap(), 15.0);
120    }
121
122    #[test]
123    fn test_sum_merge() {
124        let mut agg1 = SumAggregator::<i32>::new();
125        agg1.update(1).unwrap();
126        agg1.update(2).unwrap();
127
128        let mut agg2 = SumAggregator::<i32>::new();
129        agg2.update(3).unwrap();
130        agg2.update(4).unwrap();
131        agg2.update(5).unwrap();
132
133        let acc2 = agg2.accumulator();
134        agg1.merge(acc2).unwrap();
135
136        assert_eq!(agg1.finalize().unwrap(), 15.0);
137        assert_eq!(agg1.count(), 5);
138    }
139
140    #[test]
141    fn test_sum_reset() {
142        let mut agg = SumAggregator::<f64>::new();
143        agg.update(1.0).unwrap();
144        agg.update(2.0).unwrap();
145
146        assert_eq!(agg.finalize().unwrap(), 3.0);
147
148        agg.reset();
149        assert_eq!(agg.count(), 0);
150        assert_eq!(agg.finalize().unwrap(), 0.0);
151    }
152
153    #[test]
154    fn test_sum_serialization() {
155        let mut agg = SumAggregator::<f64>::new();
156        agg.update(1.5).unwrap();
157        agg.update(2.5).unwrap();
158
159        let acc = agg.accumulator();
160        let serialized = serde_json::to_string(&acc).unwrap();
161        let deserialized: SumAccumulator = serde_json::from_str(&serialized).unwrap();
162
163        assert_eq!(acc, deserialized);
164        assert_eq!(deserialized.sum, 4.0);
165        assert_eq!(deserialized.count, 2);
166    }
167
168    #[test]
169    fn test_sum_with_different_types() {
170        // Test with i64
171        let mut agg = SumAggregator::<i64>::new();
172        agg.update(10).unwrap();
173        agg.update(20).unwrap();
174        agg.update(30).unwrap();
175        assert_eq!(agg.finalize().unwrap(), 60.0);
176
177        // Test with f32
178        let mut agg = SumAggregator::<f32>::new();
179        agg.update(1.5).unwrap();
180        agg.update(2.5).unwrap();
181        assert_eq!(agg.finalize().unwrap(), 4.0);
182    }
183
184    #[test]
185    fn test_sum_negative_values() {
186        let mut agg = SumAggregator::<f64>::new();
187        agg.update(10.0).unwrap();
188        agg.update(-5.0).unwrap();
189        agg.update(3.0).unwrap();
190
191        assert_eq!(agg.finalize().unwrap(), 8.0);
192    }
193
194    #[test]
195    fn test_sum_empty() {
196        let agg = SumAggregator::<f64>::new();
197        assert_eq!(agg.finalize().unwrap(), 0.0);
198    }
199}
llm_optimizer_processor/aggregation/sum.rs

llm_optimizer_processor/aggregation/
sum.rs
