sokuto_server/guardrails/
multitask.rs

1//! # マルチタスク動的マーカー並列スキャンモジュール (1 パス推論)
2//!
3//! インジェクション (INJ)、有害性 (TOX)、個人情報漏洩 (PII) の 3 項目を個別推論せず、
4//! 系列先頭に `[INJ]`, `[TOX]`, `[PII]` の動的マーカーを配置して単一のフォワードパスで同時評価する。
5//! これにより推論パス数を 3 回から 1 回に集約し、66% の計算量を削減する。
6//! また、ヘルムホルツ自由エネルギー $F(x) = -T \ln \sum \exp(z_k / T)$ による OOD 異常検知を行う。
7
8use serde::{Deserialize, Serialize};
9
10use crate::guardrails::fast_fail::{GuardrailAction, GuardrailViolation};
11
12/// 特殊マーカー定義。
13pub const MARKER_INJECTION: &str = "[INJ]";
14pub const MARKER_TOXICITY: &str = "[TOX]";
15pub const MARKER_PII: &str = "[PII]";
16
17/// マルチタスクスキャナー設定。
18#[derive(Debug, Clone, PartialEq)]
19pub struct MultiTaskScannerConfig {
20    /// マルチタスクスキャンを有効にするか。
21    pub enabled: bool,
22    /// インジェクション判定閾値 (デフォルト: 0.5)。
23    pub threshold_injection: f64,
24    /// 有害性判定閾値 (デフォルト: 0.5)。
25    pub threshold_toxicity: f64,
26    /// PII 漏洩判定閾値 (デフォルト: 0.5)。
27    pub threshold_pii: f64,
28    /// ヘルムホルツ自由エネルギー閾値 (この値を超えると OOD と判定、デフォルト: -2.5)。
29    pub threshold_energy: f64,
30    /// 自由エネルギー計算時の温度パラメータ (デフォルト: 1.0)。
31    pub temperature: f64,
32    /// 違反時のアクション。
33    pub action: GuardrailAction,
34}
35
36impl Default for MultiTaskScannerConfig {
37    fn default() -> Self {
38        Self {
39            enabled: true,
40            threshold_injection: 0.5,
41            threshold_toxicity: 0.5,
42            threshold_pii: 0.5,
43            threshold_energy: -2.5,
44            temperature: 1.0,
45            action: GuardrailAction::Block,
46        }
47    }
48}
49
50/// 単一パス並列スキャン評価結果。
51#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52pub struct MultiTaskScanResult {
53    /// プロンプトインジェクション確率 $P(\text{inj} = 1 \mid X)$。
54    pub injection_prob: f64,
55    /// 有害性確率 $P(\text{tox} = 1 \mid X)$。
56    pub toxicity_prob: f64,
57    /// 個人情報漏洩確率 $P(\text{pii} = 1 \mid X)$。
58    pub pii_prob: f64,
59    /// ヘルムホルツ自由エネルギー $F(x)$。
60    pub free_energy: f64,
61    /// 脅威が検知されたか。
62    pub is_threat_detected: bool,
63    /// 分布外 (OOD) または判断不確実であるか。
64    pub is_ood: bool,
65    /// 発生したガードレール違反一覧。
66    pub violations: Vec<GuardrailViolation>,
67}
68
69/// マルチタスク推論用の合成プロンプト系列を構築する。
70///
71/// 先頭に `[INJ] [TOX] [PII]` マーカーを配置した文字列を返却する。
72pub fn build_multitask_prompt(text: &str) -> String {
73    format!(
74        "{} {} {} {}",
75        MARKER_INJECTION, MARKER_TOXICITY, MARKER_PII, text
76    )
77}
78
79/// ロジット列からシグモイド確率を算出する。
80#[inline]
81pub fn sigmoid(logit: f64) -> f64 {
82    1.0 / (1.0 + (-logit).exp())
83}
84
85/// 二値分類マルチタスクにおけるヘルムホルツ自由エネルギーを計算する。
86///
87/// 各タスク $k$ の二値決定状態 $\{0, 1\}$ (未正規化ロジット $[-z_k/2, z_k/2]$)
88/// に対する全自由エネルギーを算出する。
89/// 確信度が高い ($|z_k|$ が大きい) ほど自由エネルギーは低くなり、
90/// 判断が曖昧または不確実な異常入力 ($z_k \approx 0$) では自由エネルギーが上昇する。
91///
92/// $$F(x; T) = -T \sum_{k=1}^K \ln \left( \exp\left(\frac{z_k}{2T}\right) + \exp\left(-\frac{z_k}{2T}\right) \right)$$
93pub fn compute_free_energy(logits: &[f64], temperature: f64) -> f64 {
94    if logits.is_empty() {
95        return 0.0;
96    }
97    let t = if temperature <= 0.0 { 1.0 } else { temperature };
98
99    let mut total_energy = 0.0;
100    for &z in logits {
101        let abs_z = z.abs();
102        // LogSumExp トリックによる安定計算:
103        // ln(exp(a) + exp(-a)) = a + ln(1 + exp(-2a)) (where a = abs_z / (2t))
104        let a = abs_z / (2.0 * t);
105        let log_sum_exp = a + (-2.0 * a).exp().ln_1p();
106        total_energy += -t * log_sum_exp;
107    }
108
109    total_energy
110}
111
112/// ロジットから各タスクの確率と自由エネルギーを算出し、総合判定を下す。
113///
114/// # 引数
115/// - `inj_logit`: `[INJ]` マーカー位置のロジット。
116/// - `tox_logit`: `[TOX]` マーカー位置のロジット。
117/// - `pii_logit`: `[PII]` マーカー位置のロジット。
118/// - `config`: マルチタスク設定。
119pub fn evaluate_multitask_logits(
120    inj_logit: f64,
121    tox_logit: f64,
122    pii_logit: f64,
123    config: &MultiTaskScannerConfig,
124) -> MultiTaskScanResult {
125    let inj_prob = sigmoid(inj_logit);
126    let tox_prob = sigmoid(tox_logit);
127    let pii_prob = sigmoid(pii_logit);
128
129    let logits = [inj_logit, tox_logit, pii_logit];
130    let free_energy = compute_free_energy(&logits, config.temperature);
131
132    let mut violations = Vec::new();
133    let mut threat_detected = false;
134
135    if inj_prob >= config.threshold_injection {
136        threat_detected = true;
137        violations.push(GuardrailViolation {
138            rule: "multitask_prompt_injection".to_string(),
139            confidence: inj_prob,
140            snippet: None,
141            action_taken: config.action,
142        });
143    }
144
145    if tox_prob >= config.threshold_toxicity {
146        threat_detected = true;
147        violations.push(GuardrailViolation {
148            rule: "multitask_toxicity".to_string(),
149            confidence: tox_prob,
150            snippet: None,
151            action_taken: config.action,
152        });
153    }
154
155    if pii_prob >= config.threshold_pii {
156        threat_detected = true;
157        violations.push(GuardrailViolation {
158            rule: "multitask_pii_leakage".to_string(),
159            confidence: pii_prob,
160            snippet: None,
161            action_taken: config.action,
162        });
163    }
164
165    let is_ood = free_energy > config.threshold_energy;
166    if is_ood {
167        violations.push(GuardrailViolation {
168            rule: "multitask_free_energy_ood".to_string(),
169            confidence: sigmoid(free_energy),
170            snippet: Some(format!("free_energy={:.3}", free_energy)),
171            action_taken: config.action,
172        });
173    }
174
175    MultiTaskScanResult {
176        injection_prob: inj_prob,
177        toxicity_prob: tox_prob,
178        pii_prob,
179        free_energy,
180        is_threat_detected: threat_detected,
181        is_ood,
182        violations,
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn test_build_multitask_prompt() {
192        let prompt = build_multitask_prompt("ユーザー入力テスト");
193        assert_eq!(prompt, "[INJ] [TOX] [PII] ユーザー入力テスト");
194    }
195
196    #[test]
197    fn test_evaluate_multitask_safe() {
198        let config = MultiTaskScannerConfig::default();
199        // 低いロジット (すべて安全)
200        let res = evaluate_multitask_logits(-5.0, -4.0, -6.0, &config);
201        assert!(!res.is_threat_detected);
202        assert!(res.injection_prob < 0.1);
203        assert!(res.toxicity_prob < 0.1);
204        assert!(res.pii_prob < 0.1);
205    }
206
207    #[test]
208    fn test_evaluate_multitask_injection_violation() {
209        let config = MultiTaskScannerConfig::default();
210        // インジェクションロジットが高い
211        let res = evaluate_multitask_logits(3.0, -4.0, -5.0, &config);
212        assert!(res.is_threat_detected);
213        assert!(res.injection_prob > 0.9);
214        assert!(
215            res.violations
216                .iter()
217                .any(|v| v.rule == "multitask_prompt_injection")
218        );
219    }
220
221    #[test]
222    fn test_free_energy_ood() {
223        let config = MultiTaskScannerConfig {
224            threshold_energy: -3.0,
225            ..Default::default()
226        };
227        // 高確信度ロジット (安全側または脅威側): 低エネルギー (safe)
228        let res_confident = evaluate_multitask_logits(-5.0, -4.0, -5.0, &config);
229        assert!(!res_confident.is_ood);
230        assert!(res_confident.free_energy < -3.0);
231
232        // 曖昧・判断不能ロジット (0.0 近傍): 高エネルギー (OOD)
233        let res_ambiguous = evaluate_multitask_logits(0.0, 0.0, 0.0, &config);
234        assert!(res_ambiguous.is_ood);
235        assert!(
236            res_ambiguous
237                .violations
238                .iter()
239                .any(|v| v.rule == "multitask_free_energy_ood")
240        );
241    }
242}
sokuto_server/guardrails/multitask.rs

sokuto_server/guardrails/
multitask.rs
