ai_agents_tools/builtin/
math.rs

1use async_trait::async_trait;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::generate_schema;
7use ai_agents_core::{Tool, ToolResult, ToolSafetyMetadata};
8
9pub struct MathTool;
10
11impl MathTool {
12    pub fn new() -> Self {
13        Self
14    }
15}
16
17impl Default for MathTool {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23#[derive(Debug, Deserialize, JsonSchema)]
24struct MathInput {
25    /// Operation: mean, median, mode, stdev, variance, sum, min, max, abs, round, floor, ceil, clamp, percentage, sqrt, pow, log, range, count
26    operation: String,
27    /// Array of numbers (for statistical operations)
28    #[serde(default)]
29    values: Option<Vec<f64>>,
30    /// Single number input
31    #[serde(default)]
32    value: Option<f64>,
33    /// Decimal places (for round)
34    #[serde(default)]
35    decimals: Option<i32>,
36    /// Minimum value (for clamp/range)
37    #[serde(default)]
38    min: Option<f64>,
39    /// Maximum value (for clamp/range)
40    #[serde(default)]
41    max: Option<f64>,
42    /// Base for pow/log
43    #[serde(default)]
44    base: Option<f64>,
45    /// Exponent for pow
46    #[serde(default)]
47    exponent: Option<f64>,
48    /// Total for percentage calculation
49    #[serde(default)]
50    total: Option<f64>,
51    /// Step for range
52    #[serde(default)]
53    step: Option<f64>,
54}
55
56#[derive(Debug, Serialize, Deserialize)]
57struct StatOutput {
58    result: f64,
59    count: usize,
60}
61
62#[derive(Debug, Serialize, Deserialize)]
63struct StdevOutput {
64    stdev: f64,
65    variance: f64,
66    mean: f64,
67    count: usize,
68}
69
70#[derive(Debug, Serialize, Deserialize)]
71struct ModeOutput {
72    mode: Vec<f64>,
73    frequency: usize,
74}
75
76#[derive(Debug, Serialize, Deserialize)]
77struct SingleOutput {
78    result: f64,
79}
80
81#[derive(Debug, Serialize, Deserialize)]
82struct ClampOutput {
83    result: f64,
84    clamped: bool,
85}
86
87#[derive(Debug, Serialize, Deserialize)]
88struct RangeOutput {
89    range: Vec<f64>,
90    count: usize,
91}
92
93#[derive(Debug, Serialize, Deserialize)]
94struct MinMaxOutput {
95    min: f64,
96    max: f64,
97    range: f64,
98}
99
100#[async_trait]
101impl Tool for MathTool {
102    fn id(&self) -> &str {
103        "math"
104    }
105
106    fn name(&self) -> &str {
107        "Advanced Math"
108    }
109
110    fn description(&self) -> &str {
111        "Advanced math operations: mean (average), median, mode, stdev (standard deviation), variance, sum, min, max, minmax (both), abs, round, floor, ceil, clamp, percentage, sqrt, pow, log, log10, range, count."
112    }
113
114    fn input_schema(&self) -> Value {
115        generate_schema::<MathInput>()
116    }
117
118    fn safety_metadata(&self) -> ToolSafetyMetadata {
119        ToolSafetyMetadata::compute()
120    }
121
122    async fn execute(&self, args: Value, _ctx: ai_agents_core::ToolExecutionContext) -> ToolResult {
123        let input: MathInput = match serde_json::from_value(args) {
124            Ok(input) => input,
125            Err(e) => return ToolResult::error(format!("Invalid input: {}", e)),
126        };
127
128        match input.operation.to_lowercase().as_str() {
129            "mean" | "average" | "avg" => self.handle_mean(&input),
130            "median" => self.handle_median(&input),
131            "mode" => self.handle_mode(&input),
132            "stdev" | "std" => self.handle_stdev(&input),
133            "variance" | "var" => self.handle_variance(&input),
134            "sum" => self.handle_sum(&input),
135            "min" => self.handle_min(&input),
136            "max" => self.handle_max(&input),
137            "minmax" => self.handle_minmax(&input),
138            "abs" => self.handle_abs(&input),
139            "round" => self.handle_round(&input),
140            "floor" => self.handle_floor(&input),
141            "ceil" => self.handle_ceil(&input),
142            "clamp" => self.handle_clamp(&input),
143            "percentage" | "percent" => self.handle_percentage(&input),
144            "sqrt" => self.handle_sqrt(&input),
145            "pow" | "power" => self.handle_pow(&input),
146            "log" => self.handle_log(&input),
147            "log10" => self.handle_log10(&input),
148            "range" => self.handle_range(&input),
149            "count" => self.handle_count(&input),
150            _ => ToolResult::error(format!(
151                "Unknown operation: {}. Valid: mean, median, mode, stdev, variance, sum, min, max, minmax, abs, round, floor, ceil, clamp, percentage, sqrt, pow, log, log10, range, count",
152                input.operation
153            )),
154        }
155    }
156}
157
158impl MathTool {
159    fn get_values(&self, input: &MathInput) -> Result<Vec<f64>, ToolResult> {
160        input
161            .values
162            .clone()
163            .ok_or_else(|| ToolResult::error("'values' array is required"))
164    }
165
166    fn handle_mean(&self, input: &MathInput) -> ToolResult {
167        let values = match self.get_values(input) {
168            Ok(v) => v,
169            Err(e) => return e,
170        };
171        if values.is_empty() {
172            return ToolResult::error("values array cannot be empty");
173        }
174        let mean = values.iter().sum::<f64>() / values.len() as f64;
175        let output = StatOutput {
176            result: mean,
177            count: values.len(),
178        };
179        self.to_result(&output)
180    }
181
182    fn handle_median(&self, input: &MathInput) -> ToolResult {
183        let values = match self.get_values(input) {
184            Ok(v) => v,
185            Err(e) => return e,
186        };
187        if values.is_empty() {
188            return ToolResult::error("values array cannot be empty");
189        }
190
191        let mut sorted = values.clone();
192        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
193
194        let mid = sorted.len() / 2;
195        let median = if sorted.len() % 2 == 0 {
196            (sorted[mid - 1] + sorted[mid]) / 2.0
197        } else {
198            sorted[mid]
199        };
200
201        let output = StatOutput {
202            result: median,
203            count: values.len(),
204        };
205        self.to_result(&output)
206    }
207
208    fn handle_mode(&self, input: &MathInput) -> ToolResult {
209        let values = match self.get_values(input) {
210            Ok(v) => v,
211            Err(e) => return e,
212        };
213        if values.is_empty() {
214            return ToolResult::error("values array cannot be empty");
215        }
216
217        use std::collections::HashMap;
218        let mut counts: HashMap<String, usize> = HashMap::new();
219
220        for v in &values {
221            let key = format!("{:.10}", v);
222            *counts.entry(key).or_insert(0) += 1;
223        }
224
225        let max_count = *counts.values().max().unwrap_or(&0);
226        let modes: Vec<f64> = counts
227            .iter()
228            .filter(|&(_, &c)| c == max_count)
229            .filter_map(|(k, _)| k.parse().ok())
230            .collect();
231
232        let output = ModeOutput {
233            mode: modes,
234            frequency: max_count,
235        };
236        self.to_result(&output)
237    }
238
239    fn handle_stdev(&self, input: &MathInput) -> ToolResult {
240        let values = match self.get_values(input) {
241            Ok(v) => v,
242            Err(e) => return e,
243        };
244        if values.len() < 2 {
245            return ToolResult::error("stdev requires at least 2 values");
246        }
247
248        let mean = values.iter().sum::<f64>() / values.len() as f64;
249        let variance =
250            values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (values.len() - 1) as f64;
251        let stdev = variance.sqrt();
252
253        let output = StdevOutput {
254            stdev,
255            variance,
256            mean,
257            count: values.len(),
258        };
259        self.to_result(&output)
260    }
261
262    fn handle_variance(&self, input: &MathInput) -> ToolResult {
263        let values = match self.get_values(input) {
264            Ok(v) => v,
265            Err(e) => return e,
266        };
267        if values.len() < 2 {
268            return ToolResult::error("variance requires at least 2 values");
269        }
270
271        let mean = values.iter().sum::<f64>() / values.len() as f64;
272        let variance =
273            values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (values.len() - 1) as f64;
274
275        let output = StatOutput {
276            result: variance,
277            count: values.len(),
278        };
279        self.to_result(&output)
280    }
281
282    fn handle_sum(&self, input: &MathInput) -> ToolResult {
283        let values = match self.get_values(input) {
284            Ok(v) => v,
285            Err(e) => return e,
286        };
287
288        let sum: f64 = values.iter().sum();
289        let output = StatOutput {
290            result: sum,
291            count: values.len(),
292        };
293        self.to_result(&output)
294    }
295
296    fn handle_min(&self, input: &MathInput) -> ToolResult {
297        let values = match self.get_values(input) {
298            Ok(v) => v,
299            Err(e) => return e,
300        };
301        if values.is_empty() {
302            return ToolResult::error("values array cannot be empty");
303        }
304
305        let min = values.iter().cloned().fold(f64::INFINITY, f64::min);
306        let output = SingleOutput { result: min };
307        self.to_result(&output)
308    }
309
310    fn handle_max(&self, input: &MathInput) -> ToolResult {
311        let values = match self.get_values(input) {
312            Ok(v) => v,
313            Err(e) => return e,
314        };
315        if values.is_empty() {
316            return ToolResult::error("values array cannot be empty");
317        }
318
319        let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
320        let output = SingleOutput { result: max };
321        self.to_result(&output)
322    }
323
324    fn handle_minmax(&self, input: &MathInput) -> ToolResult {
325        let values = match self.get_values(input) {
326            Ok(v) => v,
327            Err(e) => return e,
328        };
329        if values.is_empty() {
330            return ToolResult::error("values array cannot be empty");
331        }
332
333        let min = values.iter().cloned().fold(f64::INFINITY, f64::min);
334        let max = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
335        let output = MinMaxOutput {
336            min,
337            max,
338            range: max - min,
339        };
340        self.to_result(&output)
341    }
342
343    fn handle_abs(&self, input: &MathInput) -> ToolResult {
344        let value = match input.value {
345            Some(v) => v,
346            None => return ToolResult::error("'value' is required for abs operation"),
347        };
348        let output = SingleOutput {
349            result: value.abs(),
350        };
351        self.to_result(&output)
352    }
353
354    fn handle_round(&self, input: &MathInput) -> ToolResult {
355        let value = match input.value {
356            Some(v) => v,
357            None => return ToolResult::error("'value' is required for round operation"),
358        };
359        let decimals = input.decimals.unwrap_or(0);
360        let multiplier = 10_f64.powi(decimals);
361        let rounded = (value * multiplier).round() / multiplier;
362        let output = SingleOutput { result: rounded };
363        self.to_result(&output)
364    }
365
366    fn handle_floor(&self, input: &MathInput) -> ToolResult {
367        let value = match input.value {
368            Some(v) => v,
369            None => return ToolResult::error("'value' is required for floor operation"),
370        };
371        let output = SingleOutput {
372            result: value.floor(),
373        };
374        self.to_result(&output)
375    }
376
377    fn handle_ceil(&self, input: &MathInput) -> ToolResult {
378        let value = match input.value {
379            Some(v) => v,
380            None => return ToolResult::error("'value' is required for ceil operation"),
381        };
382        let output = SingleOutput {
383            result: value.ceil(),
384        };
385        self.to_result(&output)
386    }
387
388    fn handle_clamp(&self, input: &MathInput) -> ToolResult {
389        let value = match input.value {
390            Some(v) => v,
391            None => return ToolResult::error("'value' is required for clamp operation"),
392        };
393        let min = match input.min {
394            Some(m) => m,
395            None => return ToolResult::error("'min' is required for clamp operation"),
396        };
397        let max = match input.max {
398            Some(m) => m,
399            None => return ToolResult::error("'max' is required for clamp operation"),
400        };
401
402        let clamped_value = value.max(min).min(max);
403        let output = ClampOutput {
404            result: clamped_value,
405            clamped: value != clamped_value,
406        };
407        self.to_result(&output)
408    }
409
410    fn handle_percentage(&self, input: &MathInput) -> ToolResult {
411        let value = match input.value {
412            Some(v) => v,
413            None => return ToolResult::error("'value' is required for percentage operation"),
414        };
415        let total = match input.total {
416            Some(t) => t,
417            None => return ToolResult::error("'total' is required for percentage operation"),
418        };
419        if total == 0.0 {
420            return ToolResult::error("total cannot be zero");
421        }
422
423        let percentage = (value / total) * 100.0;
424        let output = SingleOutput { result: percentage };
425        self.to_result(&output)
426    }
427
428    fn handle_sqrt(&self, input: &MathInput) -> ToolResult {
429        let value = match input.value {
430            Some(v) => v,
431            None => return ToolResult::error("'value' is required for sqrt operation"),
432        };
433        if value < 0.0 {
434            return ToolResult::error("cannot calculate sqrt of negative number");
435        }
436        let output = SingleOutput {
437            result: value.sqrt(),
438        };
439        self.to_result(&output)
440    }
441
442    fn handle_pow(&self, input: &MathInput) -> ToolResult {
443        let base = input.value.or(input.base);
444        let base = match base {
445            Some(b) => b,
446            None => return ToolResult::error("'value' or 'base' is required for pow operation"),
447        };
448        let exponent = match input.exponent {
449            Some(e) => e,
450            None => return ToolResult::error("'exponent' is required for pow operation"),
451        };
452        let output = SingleOutput {
453            result: base.powf(exponent),
454        };
455        self.to_result(&output)
456    }
457
458    fn handle_log(&self, input: &MathInput) -> ToolResult {
459        let value = match input.value {
460            Some(v) => v,
461            None => return ToolResult::error("'value' is required for log operation"),
462        };
463        if value <= 0.0 {
464            return ToolResult::error("cannot calculate log of non-positive number");
465        }
466        let result = match input.base {
467            Some(b) if b > 0.0 && b != 1.0 => value.log(b),
468            Some(_) => return ToolResult::error("log base must be positive and not equal to 1"),
469            None => value.ln(),
470        };
471        let output = SingleOutput { result };
472        self.to_result(&output)
473    }
474
475    fn handle_log10(&self, input: &MathInput) -> ToolResult {
476        let value = match input.value {
477            Some(v) => v,
478            None => return ToolResult::error("'value' is required for log10 operation"),
479        };
480        if value <= 0.0 {
481            return ToolResult::error("cannot calculate log of non-positive number");
482        }
483        let output = SingleOutput {
484            result: value.log10(),
485        };
486        self.to_result(&output)
487    }
488
489    fn handle_range(&self, input: &MathInput) -> ToolResult {
490        let min = input.min.unwrap_or(0.0);
491        let max = match input.max {
492            Some(m) => m,
493            None => return ToolResult::error("'max' is required for range operation"),
494        };
495        let step = input.step.unwrap_or(1.0);
496
497        if step == 0.0 {
498            return ToolResult::error("step cannot be zero");
499        }
500        if (max > min && step < 0.0) || (max < min && step > 0.0) {
501            return ToolResult::error("step direction doesn't match min/max range");
502        }
503
504        let mut values = Vec::new();
505        let mut current = min;
506
507        if step > 0.0 {
508            while current < max {
509                values.push(current);
510                current += step;
511            }
512        } else {
513            while current > max {
514                values.push(current);
515                current += step;
516            }
517        }
518
519        let output = RangeOutput {
520            count: values.len(),
521            range: values,
522        };
523        self.to_result(&output)
524    }
525
526    fn handle_count(&self, input: &MathInput) -> ToolResult {
527        let values = match self.get_values(input) {
528            Ok(v) => v,
529            Err(e) => return e,
530        };
531        let output = StatOutput {
532            result: values.len() as f64,
533            count: values.len(),
534        };
535        self.to_result(&output)
536    }
537
538    fn to_result<T: Serialize>(&self, output: &T) -> ToolResult {
539        match serde_json::to_string(output) {
540            Ok(json) => ToolResult::ok(json),
541            Err(e) => ToolResult::error(format!("Serialization error: {}", e)),
542        }
543    }
544}
545
546#[cfg(test)]
547mod tests {
548    use super::*;
549
550    #[tokio::test]
551    async fn test_mean() {
552        let tool = MathTool::new();
553        let result = tool
554            .execute(
555                serde_json::json!({
556                    "operation": "mean",
557                    "values": [1, 2, 3, 4, 5]
558                }),
559                ai_agents_core::ToolExecutionContext::test("test"),
560            )
561            .await;
562        assert!(result.success);
563        let output: StatOutput = serde_json::from_str(&result.output).unwrap();
564        assert!((output.result - 3.0).abs() < f64::EPSILON);
565    }
566
567    #[tokio::test]
568    async fn test_median_odd() {
569        let tool = MathTool::new();
570        let result = tool
571            .execute(
572                serde_json::json!({
573                    "operation": "median",
574                    "values": [1, 3, 2, 5, 4]
575                }),
576                ai_agents_core::ToolExecutionContext::test("test"),
577            )
578            .await;
579        assert!(result.success);
580        let output: StatOutput = serde_json::from_str(&result.output).unwrap();
581        assert!((output.result - 3.0).abs() < f64::EPSILON);
582    }
583
584    #[tokio::test]
585    async fn test_median_even() {
586        let tool = MathTool::new();
587        let result = tool
588            .execute(
589                serde_json::json!({
590                    "operation": "median",
591                    "values": [1, 2, 3, 4]
592                }),
593                ai_agents_core::ToolExecutionContext::test("test"),
594            )
595            .await;
596        assert!(result.success);
597        let output: StatOutput = serde_json::from_str(&result.output).unwrap();
598        assert!((output.result - 2.5).abs() < f64::EPSILON);
599    }
600
601    #[tokio::test]
602    async fn test_stdev() {
603        let tool = MathTool::new();
604        let result = tool
605            .execute(
606                serde_json::json!({
607                    "operation": "stdev",
608                    "values": [2, 4, 4, 4, 5, 5, 7, 9]
609                }),
610                ai_agents_core::ToolExecutionContext::test("test"),
611            )
612            .await;
613        assert!(result.success);
614        let output: StdevOutput = serde_json::from_str(&result.output).unwrap();
615        assert!((output.stdev - 2.138).abs() < 0.01);
616    }
617
618    #[tokio::test]
619    async fn test_sum() {
620        let tool = MathTool::new();
621        let result = tool
622            .execute(
623                serde_json::json!({
624                    "operation": "sum",
625                    "values": [1, 2, 3, 4, 5]
626                }),
627                ai_agents_core::ToolExecutionContext::test("test"),
628            )
629            .await;
630        assert!(result.success);
631        let output: StatOutput = serde_json::from_str(&result.output).unwrap();
632        assert!((output.result - 15.0).abs() < f64::EPSILON);
633    }
634
635    #[tokio::test]
636    async fn test_minmax() {
637        let tool = MathTool::new();
638        let result = tool
639            .execute(
640                serde_json::json!({
641                    "operation": "minmax",
642                    "values": [3, 1, 4, 1, 5, 9, 2, 6]
643                }),
644                ai_agents_core::ToolExecutionContext::test("test"),
645            )
646            .await;
647        assert!(result.success);
648        let output: MinMaxOutput = serde_json::from_str(&result.output).unwrap();
649        assert!((output.min - 1.0).abs() < f64::EPSILON);
650        assert!((output.max - 9.0).abs() < f64::EPSILON);
651        assert!((output.range - 8.0).abs() < f64::EPSILON);
652    }
653
654    #[tokio::test]
655    async fn test_round() {
656        let tool = MathTool::new();
657        let result = tool
658            .execute(
659                serde_json::json!({
660                    "operation": "round",
661                    "value": 3.14159,
662                    "decimals": 2
663                }),
664                ai_agents_core::ToolExecutionContext::test("test"),
665            )
666            .await;
667        assert!(result.success);
668        let output: SingleOutput = serde_json::from_str(&result.output).unwrap();
669        assert!((output.result - 3.14).abs() < f64::EPSILON);
670    }
671
672    #[tokio::test]
673    async fn test_clamp() {
674        let tool = MathTool::new();
675        let result = tool
676            .execute(
677                serde_json::json!({
678                    "operation": "clamp",
679                    "value": 15,
680                    "min": 0,
681                    "max": 10
682                }),
683                ai_agents_core::ToolExecutionContext::test("test"),
684            )
685            .await;
686        assert!(result.success);
687        let output: ClampOutput = serde_json::from_str(&result.output).unwrap();
688        assert!((output.result - 10.0).abs() < f64::EPSILON);
689        assert!(output.clamped);
690    }
691
692    #[tokio::test]
693    async fn test_percentage() {
694        let tool = MathTool::new();
695        let result = tool
696            .execute(
697                serde_json::json!({
698                    "operation": "percentage",
699                    "value": 25,
700                    "total": 100
701                }),
702                ai_agents_core::ToolExecutionContext::test("test"),
703            )
704            .await;
705        assert!(result.success);
706        let output: SingleOutput = serde_json::from_str(&result.output).unwrap();
707        assert!((output.result - 25.0).abs() < f64::EPSILON);
708    }
709
710    #[tokio::test]
711    async fn test_sqrt() {
712        let tool = MathTool::new();
713        let result = tool
714            .execute(
715                serde_json::json!({
716                    "operation": "sqrt",
717                    "value": 16
718                }),
719                ai_agents_core::ToolExecutionContext::test("test"),
720            )
721            .await;
722        assert!(result.success);
723        let output: SingleOutput = serde_json::from_str(&result.output).unwrap();
724        assert!((output.result - 4.0).abs() < f64::EPSILON);
725    }
726
727    #[tokio::test]
728    async fn test_pow() {
729        let tool = MathTool::new();
730        let result = tool
731            .execute(
732                serde_json::json!({
733                    "operation": "pow",
734                    "value": 2,
735                    "exponent": 10
736                }),
737                ai_agents_core::ToolExecutionContext::test("test"),
738            )
739            .await;
740        assert!(result.success);
741        let output: SingleOutput = serde_json::from_str(&result.output).unwrap();
742        assert!((output.result - 1024.0).abs() < f64::EPSILON);
743    }
744
745    #[tokio::test]
746    async fn test_range() {
747        let tool = MathTool::new();
748        let result = tool
749            .execute(
750                serde_json::json!({
751                    "operation": "range",
752                    "min": 0,
753                    "max": 5,
754                    "step": 1
755                }),
756                ai_agents_core::ToolExecutionContext::test("test"),
757            )
758            .await;
759        assert!(result.success);
760        let output: RangeOutput = serde_json::from_str(&result.output).unwrap();
761        assert_eq!(output.range, vec![0.0, 1.0, 2.0, 3.0, 4.0]);
762    }
763
764    #[tokio::test]
765    async fn test_invalid_operation() {
766        let tool = MathTool::new();
767        let result = tool
768            .execute(
769                serde_json::json!({
770                    "operation": "invalid"
771                }),
772                ai_agents_core::ToolExecutionContext::test("test"),
773            )
774            .await;
775        assert!(!result.success);
776    }
777
778    #[tokio::test]
779    async fn test_empty_values() {
780        let tool = MathTool::new();
781        let result = tool
782            .execute(
783                serde_json::json!({
784                    "operation": "mean",
785                    "values": []
786                }),
787                ai_agents_core::ToolExecutionContext::test("test"),
788            )
789            .await;
790        assert!(!result.success);
791    }
792}
ai_agents_tools/builtin/math.rs

ai_agents_tools/builtin/
math.rs
