sokuto_runtime/engine/
gating.rs

1//! # 推論エンジン向け確信度ゲーティング統合モジュール
2//!
3//! モデル成果物の較正情報 (`CalibrationConfig.gating_thresholds`) または
4//! リクエスト指定の `GatingConfig` を用いて、推論直後の判定結果 (`Answer`) に対して
5//! 透過的に 3 系統ルーティング (`AutoExecute` / `ConfirmOrEscalate` / `Fallback`) を確定・付与する。
6
7use indexmap::IndexMap;
8use sokuto_core::contract::CalibrationConfig;
9pub use sokuto_core::gating::DecisionRoute;
10use sokuto_core::gating::{GatingConfig, SystemRoutingSummary, evaluate_response_routing};
11use sokuto_core::schema::{Answer, Question};
12
13/// リクエスト固有設定とモデル較正設定から実効的なゲーティング設定をカスケード解決する。
14///
15/// # 解決優先順位
16/// 1. `gating_config`: リクエストで明示された設定。
17/// 2. `calib_config.gating_thresholds`: モデル成果物同梱の閾値設定。
18///
19/// # 引数
20/// - `gating_config`: リクエストで指定された設定 (任意)。
21/// - `calib_config`: モデルのキャリブレーション設定。
22///
23/// # 戻り値
24/// 解決された `GatingConfig`。
25pub fn resolve_gating_config(
26    gating_config: Option<&GatingConfig>,
27    calib_config: &CalibrationConfig,
28) -> GatingConfig {
29    match gating_config {
30        Some(cfg) => cfg.clone(),
31        None => calib_config.gating_config(),
32    }
33}
34
35use crate::engine::escalation::build_rich_escalation_prompt;
36use sokuto_core::gating::evaluate_answer_gating_with_energy;
37use sokuto_core::math::normalized_free_energy;
38use sokuto_core::schema::QuestionType;
39
40/// Choice 型の質問とロジット配列から、正規化ヘルムホルツ自由エネルギーを算出する。
41///
42/// モデル較正設定の `energy_temperature` (または候補数バケット温度 $T = \text{calib\_config.get\_temperature(QuestionType::Choice, K)}$) を用い、
43/// $E_{\text{norm}}(x) = -T \ln \sum_{i=1}^K \exp(z_i / T) + T \ln K$ を算出する。
44/// Choice 型以外 (Score 型、Noul 型) またはロジットが空の場合は `None` を返す。
45///
46/// # 引数
47/// - `question`: 質問定義。
48/// - `logits`: 決定ロジットスライス。
49/// - `calib_config`: キャリブレーション設定。
50///
51/// # 戻り値
52/// 正規化ヘルムホルツ自由エネルギー値 (任意)。
53pub fn calculate_choice_energy(
54    question: &Question,
55    logits: &[f64],
56    calib_config: &CalibrationConfig,
57) -> Option<f64> {
58    calculate_choice_energy_with_config(question, logits, calib_config, None)
59}
60
61/// Choice 型の質問とロジット配列、およびリクエスト固有のゲーティング設定から、正規化ヘルムホルツ自由エネルギーを算出する。
62///
63/// # 温度解決の優先順位
64/// 1. `gating_config.energy_temperature`: リクエストで明示されたエネルギー温度。
65/// 2. `calib_config.gating_thresholds.energy_temperature`: モデル成果物で事前較正されたエネルギー専用温度。
66/// 3. `calib_config.get_temperature(QuestionType::Choice, K)`: 候補数バケットに基づく Softmax 較正温度。
67///
68/// # 引数
69/// - `question`: 質問定義。
70/// - `logits`: 決定ロジットスライス。
71/// - `calib_config`: キャリブレーション設定。
72/// - `gating_config`: リクエスト固有のゲーティング設定 (任意)。
73///
74/// # 戻り値
75/// 正規化ヘルムホルツ自由エネルギー値 (任意)。
76pub fn calculate_choice_energy_with_config(
77    question: &Question,
78    logits: &[f64],
79    calib_config: &CalibrationConfig,
80    gating_config: Option<&GatingConfig>,
81) -> Option<f64> {
82    if question.question_type != QuestionType::Choice || logits.is_empty() {
83        return None;
84    }
85    let k = logits.len();
86    let temp = gating_config
87        .and_then(|cfg| cfg.energy_temperature)
88        .or(calib_config.gating_thresholds.energy_temperature)
89        .unwrap_or_else(|| calib_config.get_temperature(QuestionType::Choice, k));
90    normalized_free_energy(logits, temp).ok()
91}
92
93/// 単一の判定結果 (`Answer`) に対し、エネルギー値と解決されたゲーティング設定を適用する。
94///
95/// 設定の `enabled` が `true` の場合、`evaluate_answer_gating_with_energy` を適用して
96/// `answer.gating` に監査メタデータを付与する。無効 (`enabled == false`) の場合は何もしない。
97/// Choice 型かつ `is_ood == true` の場合、または確信度不足によるエスカレーション時は、
98/// 文脈や診断情報を含む高度な CoT エスカレーションプロンプトを構築して格納する。
99///
100/// # 引数
101/// - `answer`: 判定結果 (可変参照)。
102/// - `energy`: 正規化ヘルムホルツ自由エネルギー値 (任意)。
103/// - `state`: 文脈テキスト (プロンプト構築用、任意)。
104/// - `question_id`: 質問識別子 (プロンプト構築用、任意)。
105/// - `question`: 質問定義 (任意)。
106/// - `gating_config`: 明示的なゲーティング設定 (任意)。
107/// - `calib_config`: モデルのキャリブレーション設定。
108pub fn apply_gating_to_answer_with_energy(
109    answer: &mut Answer,
110    energy: Option<f64>,
111    state: Option<&str>,
112    question_id: Option<&str>,
113    question: Option<&Question>,
114    gating_config: Option<&GatingConfig>,
115    calib_config: &CalibrationConfig,
116) {
117    let config = resolve_gating_config(gating_config, calib_config);
118    if config.enabled {
119        let mut meta =
120            evaluate_answer_gating_with_energy(answer, question_id, question, &config, energy);
121        if meta.route != DecisionRoute::AutoExecute {
122            let rich_prompt = meta.escalation.as_ref().map(|esc| {
123                build_rich_escalation_prompt(
124                    state,
125                    question_id,
126                    question,
127                    &meta,
128                    &esc.top_candidates,
129                )
130            });
131            if let (Some(esc), Some(prompt)) = (&mut meta.escalation, rich_prompt) {
132                esc.prompt_template = Some(prompt);
133            }
134        }
135        answer.gating = Some(meta);
136    }
137}
138
139/// 単一の判定結果 (`Answer`) とロジットスライスからエネルギーを自動算出し、ゲーティング設定を適用する。
140///
141/// # 引数
142/// - `answer`: 判定結果 (可変参照)。
143/// - `logits`: 決定ロジットスライス (任意)。
144/// - `state`: 文脈テキスト (任意)。
145/// - `question_id`: 質問識別子 (任意)。
146/// - `question`: 質問定義 (任意)。
147/// - `gating_config`: 明示的なゲーティング設定 (任意)。
148/// - `calib_config`: モデルのキャリブレーション設定。
149pub fn apply_gating_to_answer_with_logits(
150    answer: &mut Answer,
151    logits: Option<&[f64]>,
152    state: Option<&str>,
153    question_id: Option<&str>,
154    question: Option<&Question>,
155    gating_config: Option<&GatingConfig>,
156    calib_config: &CalibrationConfig,
157) {
158    let energy = match (question, logits) {
159        (Some(q), Some(lg)) => {
160            calculate_choice_energy_with_config(q, lg, calib_config, gating_config)
161        }
162        _ => None,
163    };
164    apply_gating_to_answer_with_energy(
165        answer,
166        energy,
167        state,
168        question_id,
169        question,
170        gating_config,
171        calib_config,
172    );
173}
174
175/// 単一の判定結果 (`Answer`) に対し、解決されたゲーティング設定を適用する。
176///
177/// エネルギー値は `None` として処理されるため、OOD チェックを伴わない標準の確信度ゲーティングが行われる。
178/// ロジットに基づく OOD 安全弁を有効化する場合は [`apply_gating_to_answer_with_energy`] または
179/// [`apply_gating_to_answer_with_logits`] を使用する。
180///
181/// # 引数
182/// - `answer`: 判定結果 (可変参照)。
183/// - `state`: 文脈テキスト (プロンプト構築用、任意)。
184/// - `question_id`: 質問識別子 (プロンプト構築用、任意)。
185/// - `question`: 質問定義 (任意)。
186/// - `gating_config`: 明示的なゲーティング設定 (任意)。
187/// - `calib_config`: モデルのキャリブレーション設定。
188pub fn apply_gating_to_answer(
189    answer: &mut Answer,
190    state: Option<&str>,
191    question_id: Option<&str>,
192    question: Option<&Question>,
193    gating_config: Option<&GatingConfig>,
194    calib_config: &CalibrationConfig,
195) {
196    apply_gating_to_answer_with_energy(
197        answer,
198        None,
199        state,
200        question_id,
201        question,
202        gating_config,
203        calib_config,
204    );
205}
206
207/// 複数質問の判定結果マップに対し、ロジット配列を用いたエネルギー付きゲーティング判定を一括適用する。
208///
209/// 各回答にゲーティングメタデータ (および OOD 診断) を付与し、有効な場合は `SystemRoutingSummary` を返却する。
210///
211/// # 引数
212/// - `answers`: 判定結果マップ (可変参照)。
213/// - `logits_map`: 各質問のロジット配列マップ。
214/// - `state`: 共通の文脈テキスト (任意)。
215/// - `questions`: 質問定義マップ。
216/// - `gating_config`: 明示的なゲーティング設定 (任意)。
217/// - `calib_config`: モデルのキャリブレーション設定。
218///
219/// # 戻り値
220/// ゲーティングが有効な場合は `Some(SystemRoutingSummary)`、無効な場合は `None`。
221pub fn apply_gating_to_answers_with_logits(
222    answers: &mut IndexMap<String, Answer>,
223    logits_map: &IndexMap<String, Vec<f64>>,
224    state: Option<&str>,
225    questions: &IndexMap<String, Question>,
226    gating_config: Option<&GatingConfig>,
227    calib_config: &CalibrationConfig,
228) -> Option<SystemRoutingSummary> {
229    let config = resolve_gating_config(gating_config, calib_config);
230    if !config.enabled {
231        return None;
232    }
233
234    for (qid, answer) in answers.iter_mut() {
235        let q_def = questions.get(qid);
236        let q_logits = logits_map.get(qid).map(|v| v.as_slice());
237        apply_gating_to_answer_with_logits(
238            answer,
239            q_logits,
240            state,
241            Some(qid),
242            q_def,
243            gating_config,
244            calib_config,
245        );
246    }
247
248    Some(evaluate_response_routing(answers))
249}
250
251/// 複数質問の判定結果マップに対し、ゲーティング判定を一括適用し集約サマリーを算出する。
252///
253/// 各回答にゲーティングメタデータを付与し、有効な場合は `SystemRoutingSummary` を返却する。
254///
255/// # 引数
256/// - `answers`: 判定結果マップ (可変参照)。
257/// - `state`: 共通の文脈テキスト (任意)。
258/// - `questions`: 質問定義マップ。
259/// - `gating_config`: 明示的なゲーティング設定 (任意)。
260/// - `calib_config`: モデルのキャリブレーション設定。
261///
262/// # 戻り値
263/// ゲーティングが有効な場合は `Some(SystemRoutingSummary)`、無効な場合は `None`。
264pub fn apply_gating_to_answers(
265    answers: &mut IndexMap<String, Answer>,
266    state: Option<&str>,
267    questions: &IndexMap<String, Question>,
268    gating_config: Option<&GatingConfig>,
269    calib_config: &CalibrationConfig,
270) -> Option<SystemRoutingSummary> {
271    let config = resolve_gating_config(gating_config, calib_config);
272    if !config.enabled {
273        return None;
274    }
275
276    for (qid, answer) in answers.iter_mut() {
277        let q_def = questions.get(qid);
278        apply_gating_to_answer(answer, state, Some(qid), q_def, gating_config, calib_config);
279    }
280
281    Some(evaluate_response_routing(answers))
282}
283
284#[cfg(test)]
285mod tests {
286    use super::*;
287    use sokuto_core::contract::calibration::GatingThresholds;
288
289    #[test]
290    fn test_resolve_gating_config_cascade() {
291        let calib = CalibrationConfig {
292            gating_thresholds: GatingThresholds {
293                high_threshold: 0.72,
294                low_threshold: 0.38,
295                top_margin_threshold: 0.18,
296                ..Default::default()
297            },
298            ..Default::default()
299        };
300
301        // 1. 指定なしの場合は calib_config から導出
302        let resolved_default = resolve_gating_config(None, &calib);
303        assert!(resolved_default.enabled);
304        assert!((resolved_default.high_threshold - 0.72).abs() < 1e-9);
305
306        // 2. 明示指定された場合はそれを優先
307        let explicit = GatingConfig {
308            enabled: false,
309            high_threshold: 0.90,
310            low_threshold: 0.40,
311            top_margin_threshold: 0.20,
312            ood_enabled: true,
313            energy_threshold: -1.0,
314            ..GatingConfig::default()
315        };
316        let resolved_explicit = resolve_gating_config(Some(&explicit), &calib);
317        assert!(!resolved_explicit.enabled);
318        assert!((resolved_explicit.high_threshold - 0.90).abs() < 1e-9);
319    }
320
321    #[test]
322    fn test_apply_gating_to_answer() {
323        let calib = CalibrationConfig::default();
324        let mut probs = IndexMap::new();
325        probs.insert("opt_a".to_string(), 0.95);
326        probs.insert("opt_b".to_string(), 0.05);
327
328        let mut ans = Answer::choice("opt_a", probs, 0.90);
329        assert!(ans.gating.is_none());
330
331        apply_gating_to_answer(&mut ans, Some("文脈"), Some("q1"), None, None, &calib);
332        assert!(ans.gating.is_some());
333        assert_eq!(
334            ans.gating.as_ref().unwrap().route,
335            DecisionRoute::AutoExecute
336        );
337        // AutoExecute 時はプロンプト生成がスキップされ、escalation は None のまま
338        assert!(ans.gating.as_ref().unwrap().escalation.is_none());
339    }
340
341    #[test]
342    fn test_apply_gating_to_answer_disabled() {
343        let calib = CalibrationConfig::default();
344        let disabled_cfg = GatingConfig {
345            enabled: false,
346            ..GatingConfig::default()
347        };
348
349        let mut probs = IndexMap::new();
350        probs.insert("opt_a".to_string(), 0.95);
351        probs.insert("opt_b".to_string(), 0.05);
352
353        let mut ans = Answer::choice("opt_a", probs, 0.90);
354        apply_gating_to_answer(
355            &mut ans,
356            None,
357            Some("q1"),
358            None,
359            Some(&disabled_cfg),
360            &calib,
361        );
362        assert!(ans.gating.is_none());
363    }
364
365    #[test]
366    fn test_apply_gating_to_answer_escalation_rich_prompt() {
367        let calib = CalibrationConfig::default();
368        let mut criteria = IndexMap::new();
369        criteria.insert("opt_a".to_string(), "オプションAの詳細説明".to_string());
370        criteria.insert("opt_b".to_string(), "オプションBの詳細説明".to_string());
371        let q = Question::new_choice("指示文", criteria);
372
373        let mut probs = IndexMap::new();
374        probs.insert("opt_a".to_string(), 0.51);
375        probs.insert("opt_b".to_string(), 0.49);
376
377        let mut ans = Answer::choice("opt_a", probs, 0.51);
378        apply_gating_to_answer(
379            &mut ans,
380            Some("ユーザー問い合わせ文脈"),
381            Some("q_test"),
382            Some(&q),
383            None,
384            &calib,
385        );
386
387        assert!(ans.gating.is_some());
388        let meta = ans.gating.as_ref().unwrap();
389        assert_eq!(meta.route, DecisionRoute::ConfirmOrEscalate);
390        assert!(meta.escalation.is_some());
391
392        let prompt = meta
393            .escalation
394            .as_ref()
395            .unwrap()
396            .prompt_template
397            .as_ref()
398            .unwrap();
399
400        assert!(prompt.contains("<context>"));
401        assert!(prompt.contains("ユーザー問い合わせ文脈"));
402        assert!(prompt.contains("</context>"));
403        assert!(prompt.contains("オプションAの詳細説明"));
404        assert!(prompt.contains("System 1 の判定では候補"));
405        assert!(prompt.contains("思考連鎖 (Chain-of-Thought)"));
406        assert!(prompt.contains("```json"));
407    }
408
409    #[test]
410    fn test_calculate_choice_energy() {
411        let calib = CalibrationConfig::default();
412        let mut criteria = IndexMap::new();
413        criteria.insert("a".to_string(), "説明A".to_string());
414        criteria.insert("b".to_string(), "説明B".to_string());
415        let q = Question::new_choice("指示", criteria);
416
417        // ID 的ロジット (顕著な差)
418        let id_logits = [5.0, -2.0];
419        let energy_id = calculate_choice_energy(&q, &id_logits, &calib).unwrap();
420        // OOD 的ロジット (低確信度・平坦)
421        let ood_logits = [0.1, 0.1];
422        let energy_ood = calculate_choice_energy(&q, &ood_logits, &calib).unwrap();
423
424        // 顕著な差がある ID は低エネルギー、平坦な OOD は高エネルギー
425        assert!(energy_id < energy_ood);
426    }
427
428    #[test]
429    fn test_apply_gating_to_answer_with_energy_ood() {
430        let calib = CalibrationConfig::default();
431        let mut criteria = IndexMap::new();
432        criteria.insert("opt_a".to_string(), "説明A".to_string());
433        criteria.insert("opt_b".to_string(), "説明B".to_string());
434        let q = Question::new_choice("指示文", criteria);
435
436        let mut probs = IndexMap::new();
437        probs.insert("opt_a".to_string(), 0.95);
438        probs.insert("opt_b".to_string(), 0.05);
439
440        // 確信度は高くても、エネルギーが閾値 (-1.0) を超えて OOD (平坦/低信頼性) と検知された場合
441        let mut ans = Answer::choice("opt_a", probs, 0.95);
442        apply_gating_to_answer_with_energy(
443            &mut ans,
444            Some(-0.2), // > -1.0
445            Some("未定義カテゴリの文脈"),
446            Some("q_ood"),
447            Some(&q),
448            None,
449            &calib,
450        );
451
452        assert!(ans.gating.is_some());
453        let meta = ans.gating.as_ref().unwrap();
454        assert_eq!(meta.route, DecisionRoute::Fallback);
455        assert!(meta.is_ood);
456        assert_eq!(meta.energy, Some(-0.2));
457
458        // エスカレーションプロンプトに OOD 警告が含まれることを検証
459        let esc = meta.escalation.as_ref().unwrap();
460        let prompt = esc.prompt_template.as_ref().unwrap();
461        assert!(prompt.contains("None of the above / 該当なし"));
462        assert!(prompt.contains("未定義カテゴリ"));
463    }
464
465    #[test]
466    fn test_apply_gating_to_answer_with_logits() {
467        let calib = CalibrationConfig::default();
468        let mut criteria = IndexMap::new();
469        criteria.insert("opt_a".to_string(), "説明A".to_string());
470        criteria.insert("opt_b".to_string(), "説明B".to_string());
471        let q = Question::new_choice("指示文", criteria);
472
473        let mut probs = IndexMap::new();
474        probs.insert("opt_a".to_string(), 0.95);
475        probs.insert("opt_b".to_string(), 0.05);
476
477        let mut ans = Answer::choice("opt_a", probs, 0.95);
478        // ID 的な明確なロジット -> 低エネルギー -> AutoExecute 通過
479        let id_logits = [5.0, -2.0];
480        apply_gating_to_answer_with_logits(
481            &mut ans,
482            Some(&id_logits),
483            Some("通常文脈"),
484            Some("q_id"),
485            Some(&q),
486            None,
487            &calib,
488        );
489
490        let meta = ans.gating.as_ref().unwrap();
491        assert_eq!(meta.route, DecisionRoute::AutoExecute);
492        assert!(!meta.is_ood);
493        assert!(meta.energy.is_some());
494    }
495}
sokuto_runtime/engine/gating.rs

sokuto_runtime/engine/
gating.rs
