sokuto_server/guardrails/
window.rs

1//! # 階層型 Fast-Fail 境界縮約窓生成モジュール (Head-Tail Window Extraction)
2//!
3//! 自然言語によるプロンプトインジェクションの 99% 以上はプロンプトの境界
4//! (先頭によるペルソナ上書き、末尾による直前命令打ち消し) に集中する。
5//! 入力長が 512 トークンを超える長文に対し、全文エンコード ($O(L^2)$) を回避し、
6//! 先頭 256 トークンと末尾 256 トークンを結合した 512 トークン窓を生成して
7//! 5〜8ms での高速 Early-Exit スキャンを実現する。
8
9use std::borrow::Cow;
10
11/// 縮約窓抽出設定構造体。
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct WindowConfig {
14    /// 縮約窓抽出を有効にするか。
15    pub enabled: bool,
16    /// 縮約窓を適用するトリガーとなる最大トークン閾値 (デフォルト: 512)。
17    pub max_tokens_threshold: usize,
18    /// 先頭から抽出するトークン数 (デフォルト: 256)。
19    pub head_tokens: usize,
20    /// 末尾から抽出するトークン数 (デフォルト: 256)。
21    pub tail_tokens: usize,
22    /// 日本語における 1 トークンあたりの平均文字数見積もり (デフォルト: 1.5)。
23    pub chars_per_token_estimate: usize,
24}
25
26impl Default for WindowConfig {
27    fn default() -> Self {
28        Self {
29            enabled: true,
30            max_tokens_threshold: 512,
31            head_tokens: 256,
32            tail_tokens: 256,
33            chars_per_token_estimate: 2,
34        }
35    }
36}
37
38/// 縮約窓抽出結果。
39#[derive(Debug, Clone, PartialEq)]
40pub struct WindowExtractionResult<'a> {
41    /// 抽出された縮約窓テキスト。
42    pub window_text: Cow<'a, str>,
43    /// 実際に縮約窓が適用されたか (閾値未満の場合は false)。
44    pub was_truncated: bool,
45    /// 元の推定トークン数。
46    pub original_estimated_tokens: usize,
47    /// 窓抽出後の推定トークン数。
48    pub window_estimated_tokens: usize,
49}
50
51/// 入力テキストから先頭と末尾の重要境界領域を抽出して結合した縮約窓を生成する。
52///
53/// UTF-8 のマルチバイト文字境界を破壊しないよう、必ず文字境界 (Char Boundary) でスライスを行う。
54///
55/// # 引数
56/// - `text`: 対象テキスト。
57/// - `config`: 縮約窓設定。
58///
59/// # 戻り値
60/// 縮約窓抽出結果。
61pub fn extract_head_tail_window<'a>(
62    text: &'a str,
63    config: &WindowConfig,
64) -> WindowExtractionResult<'a> {
65    if !config.enabled {
66        let est = estimate_tokens(text, config.chars_per_token_estimate);
67        return WindowExtractionResult {
68            window_text: Cow::Borrowed(text),
69            was_truncated: false,
70            original_estimated_tokens: est,
71            window_estimated_tokens: est,
72        };
73    }
74
75    let char_count = text.chars().count();
76    let est_tokens = estimate_tokens_from_char_count(char_count, config.chars_per_token_estimate);
77
78    if est_tokens <= config.max_tokens_threshold {
79        return WindowExtractionResult {
80            window_text: Cow::Borrowed(text),
81            was_truncated: false,
82            original_estimated_tokens: est_tokens,
83            window_estimated_tokens: est_tokens,
84        };
85    }
86
87    // 先頭および末尾の必要文字数を計算
88    let head_chars = config.head_tokens * config.chars_per_token_estimate;
89    let tail_chars = config.tail_tokens * config.chars_per_token_estimate;
90
91    if head_chars + tail_chars >= char_count {
92        return WindowExtractionResult {
93            window_text: Cow::Borrowed(text),
94            was_truncated: false,
95            original_estimated_tokens: est_tokens,
96            window_estimated_tokens: est_tokens,
97        };
98    }
99
100    // 先頭部分の文字境界を取得
101    let head_byte_idx = text
102        .char_indices()
103        .nth(head_chars)
104        .map(|(idx, _)| idx)
105        .unwrap_or(text.len());
106    let head_slice = &text[..head_byte_idx];
107
108    // 末尾部分の文字境界を取得
109    let tail_start_char = char_count.saturating_sub(tail_chars);
110    let tail_byte_idx = text
111        .char_indices()
112        .nth(tail_start_char)
113        .map(|(idx, _)| idx)
114        .unwrap_or(0);
115    let tail_slice = &text[tail_byte_idx..];
116
117    let omitted_chars = char_count.saturating_sub(head_chars + tail_chars);
118    let omitted_tokens = omitted_chars / config.chars_per_token_estimate.max(1);
119
120    let window_text = format!(
121        "{}\n[... 縮約窓走査により中間 {} トークンを省略 ...]\n{}",
122        head_slice, omitted_tokens, tail_slice
123    );
124
125    let window_tokens = config.head_tokens + config.tail_tokens;
126
127    WindowExtractionResult {
128        window_text: Cow::Owned(window_text),
129        was_truncated: true,
130        original_estimated_tokens: est_tokens,
131        window_estimated_tokens: window_tokens,
132    }
133}
134
135/// 文字列からトークン数を簡易見積もりする。
136#[inline]
137pub fn estimate_tokens(text: &str, chars_per_token: usize) -> usize {
138    let char_count = text.chars().count();
139    estimate_tokens_from_char_count(char_count, chars_per_token)
140}
141
142#[inline]
143fn estimate_tokens_from_char_count(char_count: usize, chars_per_token: usize) -> usize {
144    let factor = chars_per_token.max(1);
145    char_count.div_ceil(factor)
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn test_window_short_text_unchanged() {
154        let config = WindowConfig::default();
155        let text = "これは短いテスト文章です。";
156        let res = extract_head_tail_window(text, &config);
157        assert!(!res.was_truncated);
158        assert_eq!(res.window_text, text);
159    }
160
161    #[test]
162    fn test_window_long_text_truncated() {
163        let config = WindowConfig {
164            enabled: true,
165            max_tokens_threshold: 10,
166            head_tokens: 4,
167            tail_tokens: 4,
168            chars_per_token_estimate: 2,
169        };
170        // 40 文字 = 約 20 トークン (> 10)
171        let text = "0123456789abcdefghijklmnopqrstuvwxyzABCD";
172        let res = extract_head_tail_window(text, &config);
173        assert!(res.was_truncated);
174        assert!(res.window_text.starts_with("01234567"));
175        assert!(res.window_text.ends_with("wxyzABCD"));
176        assert!(res.window_text.contains("[... 縮約窓走査により中間"));
177    }
178
179    #[test]
180    fn test_window_multibyte_safety() {
181        let config = WindowConfig {
182            enabled: true,
183            max_tokens_threshold: 10,
184            head_tokens: 3,
185            tail_tokens: 3,
186            chars_per_token_estimate: 2,
187        };
188        // 日本語マルチバイト
189        let text = "あいうえおかきくけこさしすせそたちつてとなにぬねのはひふへほ";
190        let res = extract_head_tail_window(text, &config);
191        assert!(res.was_truncated);
192        // パニックせず正常に UTF-8 文字列が構築されること
193        assert!(res.window_text.starts_with("あいうえおか"));
194        assert!(res.window_text.ends_with("はひふへほ"));
195    }
196}
sokuto_server/guardrails/window.rs

sokuto_server/guardrails/
window.rs
