Skip to main content

sokuto_server/guardrails/
arithmetic.rs

1//! # 算術・数え上げ事前集計ヘルパーモジュール (Arithmetic Annotation)
2//!
3//! 小規模エンコーダーモデルが原理的に苦手とする厳密な四則演算 (金額合算など) や
4//! 配列要素のカウントを前処理で決定論的に代行し、State 末尾に集計メタデータを透過付加する。
5
6use serde_json::Value;
7
8/// 算術事前集計設定構造体。
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ArithmeticConfig {
11    /// 構造化 State に対する事前集計を有効にするか。
12    pub enable_aggregation: bool,
13    /// 集計対象とする代表的な数値フィールド名のリスト。
14    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/// 集計サマリー情報。
34#[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/// キー名が集計対象ターゲットに厳格に一致するか判定する。
42#[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
52/// JSON 値から統計サマリーを再帰的または表層走査で集計する。
53fn 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            // 数値配列の場合の合計
61            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            // オブジェクト配列の場合の指定キー合計
76            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                // 配列フィールドのカウント
97                if let Some(arr) = v.as_array() {
98                    summary
99                        .numeric_sums
100                        .push((format!("{k}_count"), arr.len() as f64));
101
102                    // 配列内のオブジェクトの数値キー合計
103                    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                // 特定数値フィールド
121                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
138/// State 文字列および JSON Value を検査し、事前集計サマリーブロックを生成する。
139///
140/// # 引数
141/// - `state_val`: リクエストの State (`serde_json::Value`)。
142/// - `config`: 算術設定。
143///
144/// # 戻り値
145/// 集計可能な項目が存在した場合は `[Derived Stats: ...]` の注記文字列を返却する。
146pub fn compute_arithmetic_summary(state_val: &Value, config: &ArithmeticConfig) -> Option<String> {
147    if !config.enable_aggregation {
148        return None;
149    }
150
151    // 文字列の場合は JSON パースを試行
152    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}