sokuto_server/guardrails/
arithmetic.rs1use serde_json::Value;
7
8#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ArithmeticConfig {
11 pub enable_aggregation: bool,
13 pub target_numeric_keys: Vec<String>,
15}
16
17impl Default for ArithmeticConfig {
18 fn default() -> Self {
19 Self {
20 enable_aggregation: true,
21 target_numeric_keys: vec![
22 "amount".to_string(),
23 "total".to_string(),
24 "total_amount".to_string(),
25 "price".to_string(),
26 "quantity".to_string(),
27 "cost".to_string(),
28 ],
29 }
30 }
31}
32
33#[derive(Debug, Default, Clone)]
35struct DerivedSummary {
36 item_count: Option<usize>,
37 key_count: Option<usize>,
38 numeric_sums: Vec<(String, f64)>,
39}
40
41#[inline]
43fn matches_target_key(key: &str, target_keys: &[String]) -> bool {
44 let lower = key.to_lowercase();
45 target_keys.iter().any(|target| {
46 lower == *target
47 || lower.ends_with(&format!("_{target}"))
48 || lower.starts_with(&format!("{target}_"))
49 })
50}
51
52fn inspect_and_aggregate(val: &Value, target_keys: &[String]) -> Option<DerivedSummary> {
54 let mut summary = DerivedSummary::default();
55
56 match val {
57 Value::Array(arr) => {
58 summary.item_count = Some(arr.len());
59
60 let mut num_sum = 0.0;
62 let mut num_count = 0;
63 for item in arr {
64 if let Some(n) = item.as_f64() {
65 num_sum += n;
66 num_count += 1;
67 }
68 }
69 if num_count > 0 && num_count == arr.len() {
70 summary
71 .numeric_sums
72 .push(("elements_sum".to_string(), num_sum));
73 }
74
75 for key in target_keys {
77 let mut sum = 0.0;
78 let mut match_count = 0;
79 for item in arr {
80 if let Some(obj) = item.as_object()
81 && let Some(val_num) = obj.get(key).and_then(|v| v.as_f64())
82 {
83 sum += val_num;
84 match_count += 1;
85 }
86 }
87 if match_count > 0 {
88 summary.numeric_sums.push((format!("total_{key}"), sum));
89 }
90 }
91 }
92 Value::Object(obj) => {
93 summary.key_count = Some(obj.len());
94
95 for (k, v) in obj {
96 if let Some(arr) = v.as_array() {
98 summary
99 .numeric_sums
100 .push((format!("{k}_count"), arr.len() as f64));
101
102 for key in target_keys {
104 let mut sum = 0.0;
105 let mut match_count = 0;
106 for item in arr {
107 if let Some(item_obj) = item.as_object()
108 && let Some(val_num) = item_obj.get(key).and_then(|v| v.as_f64())
109 {
110 sum += val_num;
111 match_count += 1;
112 }
113 }
114 if match_count > 0 {
115 summary.numeric_sums.push((format!("total_{key}"), sum));
116 }
117 }
118 }
119
120 if matches_target_key(k, target_keys)
122 && let Some(num) = v.as_f64()
123 {
124 summary.numeric_sums.push((k.clone(), num));
125 }
126 }
127 }
128 _ => return None,
129 }
130
131 if summary.item_count.is_some() || summary.key_count.is_some() {
132 Some(summary)
133 } else {
134 None
135 }
136}
137
138pub fn compute_arithmetic_summary(state_val: &Value, config: &ArithmeticConfig) -> Option<String> {
147 if !config.enable_aggregation {
148 return None;
149 }
150
151 let parsed_val: Option<Value> = match state_val {
153 Value::String(s) => serde_json::from_str(s).ok(),
154 other => Some(other.clone()),
155 };
156
157 let target_val = parsed_val?;
158 let summary = inspect_and_aggregate(&target_val, &config.target_numeric_keys)?;
159
160 let mut parts = Vec::new();
161
162 if let Some(items) = summary.item_count {
163 parts.push(format!("item_count={items}"));
164 }
165 if let Some(keys) = summary.key_count {
166 parts.push(format!("key_count={keys}"));
167 }
168 for (k, v) in summary.numeric_sums {
169 if v.fract() == 0.0 {
170 parts.push(format!("{k}={}", v as i64));
171 } else {
172 parts.push(format!("{k}={v:.2}"));
173 }
174 }
175
176 if parts.is_empty() {
177 None
178 } else {
179 Some(format!("\n[Derived Stats: {}]", parts.join(", ")))
180 }
181}
182
183#[cfg(test)]
184mod tests {
185 use super::*;
186 use serde_json::json;
187
188 #[test]
189 fn test_compute_arithmetic_summary_array_objects() {
190 let config = ArithmeticConfig::default();
191 let val = json!([
192 {"name": "item A", "amount": 1000},
193 {"name": "item B", "amount": 2500},
194 {"name": "item C", "amount": 1500}
195 ]);
196
197 let summary = compute_arithmetic_summary(&val, &config);
198 assert!(summary.is_some());
199 let s = summary.unwrap();
200 assert!(s.contains("item_count=3"));
201 assert!(s.contains("total_amount=5000"));
202 }
203
204 #[test]
205 fn test_compute_arithmetic_summary_numbers_array() {
206 let config = ArithmeticConfig::default();
207 let val = json!([10, 20, 30, 40]);
208
209 let summary = compute_arithmetic_summary(&val, &config);
210 assert!(summary.is_some());
211 let s = summary.unwrap();
212 assert!(s.contains("item_count=4"));
213 assert!(s.contains("elements_sum=100"));
214 }
215}