sokuto_server/guardrails/
streaming.rs

1//! # 投機的ストリーミング検証モジュール (SSE Realtime Interceptor)
2//!
3//! System 2 (上位 LLM) からの SSE (Server-Sent Events) ストリーミングレスポンスに対し、
4//! 句読点 (「。」「\n」「！」「？」) 単位でリアルタイムに文境界をバッファリングし、
5//! 投機的 NLI ハルシネーション検証を並行して実行する。
6//! 不整合を検知した瞬間、ダウンストリームへの送信を即時切断し、安全なエラー・置換メッセージを
7//! 注入して安全にコネクションをクローズする。
8
9use std::sync::Arc;
10
11use bytes::Bytes;
12use futures_util::Stream;
13use futures_util::stream::BoxStream;
14
15use crate::guardrails::nli::{
16    HallucinationVerificationResult, HallucinationVerifierConfig, verify_claim_nli,
17};
18
19/// 文境界区切り文字。
20pub const SENTENCE_DELIMITERS: &[char] = &['。', '\n', '！', '？', '!', '?'];
21
22/// ストリーミング遮断時の置換メッセージ。
23pub const TERMINATION_REPLACEMENT_MESSAGE: &str =
24    "\n[安全基準に基づき生成を中断しました。事実整合性エラーを検知しました]";
25
26/// ストリーミング検証設定。
27#[derive(Debug, Clone, PartialEq)]
28pub struct StreamingVerifierConfig {
29    /// ストリーミング検証を有効にするか。
30    pub enabled: bool,
31    /// 文境界の最小文字数 (極端に短い助詞などの細切れ分割を抑止、デフォルト: 6)。
32    pub min_sentence_chars: usize,
33    /// エラー置換メッセージ。
34    pub replacement_message: String,
35    /// ハルシネーション検証設定。
36    pub verifier_config: HallucinationVerifierConfig,
37}
38
39impl Default for StreamingVerifierConfig {
40    fn default() -> Self {
41        Self {
42            enabled: true,
43            min_sentence_chars: 6,
44            replacement_message: TERMINATION_REPLACEMENT_MESSAGE.to_string(),
45            verifier_config: HallucinationVerifierConfig::default(),
46        }
47    }
48}
49
50/// 文境界バッファ。
51#[derive(Debug, Default)]
52pub struct SentenceChunkBuffer {
53    buffer: String,
54}
55
56impl SentenceChunkBuffer {
57    /// 新規バッファを生成する。
58    pub fn new() -> Self {
59        Self {
60            buffer: String::new(),
61        }
62    }
63
64    /// チャンク文字列を追加し、文境界に達した完結文を取り出す。
65    ///
66    /// # 引数
67    /// - `chunk`: 受信した文字列スライスのチャンク。
68    /// - `min_chars`: 最小文文字数。
69    ///
70    /// # 戻り値
71    /// 境界に達した完結文のベクタ (未完結の末尾は内部バッファに残る)。
72    pub fn push_and_drain_sentences(&mut self, chunk: &str, min_chars: usize) -> Vec<String> {
73        self.buffer.push_str(chunk);
74        let mut sentences = Vec::new();
75
76        while let Some(idx) = self.find_delimiter_index(min_chars) {
77            // 区切り文字の次のバイトインデックス
78            let delim_char = self.buffer[idx..].chars().next().unwrap();
79            let split_byte_idx = idx + delim_char.len_utf8();
80
81            let sentence = self.buffer[..split_byte_idx].to_string();
82            self.buffer = self.buffer[split_byte_idx..].to_string();
83            sentences.push(sentence);
84        }
85
86        sentences
87    }
88
89    /// バッファ内に残存する全テキストをフラッシュして取り出す。
90    pub fn flush(&mut self) -> Option<String> {
91        if self.buffer.is_empty() {
92            None
93        } else {
94            let remaining = std::mem::take(&mut self.buffer);
95            Some(remaining)
96        }
97    }
98
99    #[inline]
100    fn find_delimiter_index(&self, min_chars: usize) -> Option<usize> {
101        let mut char_count = 0;
102        for (byte_idx, c) in self.buffer.char_indices() {
103            char_count += 1;
104            if char_count >= min_chars && SENTENCE_DELIMITERS.contains(&c) {
105                return Some(byte_idx);
106            }
107        }
108        None
109    }
110}
111
112/// SSE の data ペイロードからテキスト本文を抽出する簡易ヘルパー。
113///
114/// 形式例: `data: {"choices":[{"delta":{"content":"こんにちは"}}]}\n\n`
115pub fn extract_content_from_sse_line(line: &str) -> Option<String> {
116    let trimmed = line.trim();
117    if !trimmed.starts_with("data:") {
118        return None;
119    }
120    let data_body = trimmed.strip_prefix("data:")?.trim();
121    if data_body == "[DONE]" {
122        return None;
123    }
124
125    if let Ok(v) = serde_json::from_str::<serde_json::Value>(data_body) {
126        if let Some(content) = v
127            .pointer("/choices/0/delta/content")
128            .and_then(|c| c.as_str())
129        {
130            return Some(content.to_string());
131        }
132        if let Some(text) = v.get("text").and_then(|t| t.as_str()) {
133            return Some(text.to_string());
134        }
135    }
136
137    None
138}
139
140/// 投機的ストリーム検証インターセプター。
141pub struct SpeculativeStreamInterceptor {
142    config: Arc<StreamingVerifierConfig>,
143    state_context: Arc<String>,
144}
145
146impl SpeculativeStreamInterceptor {
147    /// 新規インターセプターを構築する。
148    pub fn new(config: StreamingVerifierConfig, state_context: String) -> Self {
149        Self {
150            config: Arc::new(config),
151            state_context: Arc::new(state_context),
152        }
153    }
154
155    /// 文がハルシネーションであるかを同期または投機的に検査する。
156    pub fn verify_sentence(
157        &self,
158        sentence: &str,
159        mock_nli_logits: Option<[f64; 3]>,
160    ) -> HallucinationVerificationResult {
161        // デフォルトでは高い含意確率ロジット、指定があればそれを使用
162        let logits = mock_nli_logits.unwrap_or([3.0, -1.0, -3.0]);
163        verify_claim_nli(
164            &self.state_context,
165            sentence,
166            logits,
167            &self.config.verifier_config,
168        )
169    }
170
171    /// 遮断用 SSE エラーチャンクバイト列を構築する。
172    pub fn build_error_event_bytes(&self, reason: &str) -> Bytes {
173        let payload = serde_json::json!({
174            "error": "guardrail_violation",
175            "rule": "outbound_hallucination",
176            "message": self.config.replacement_message,
177            "reason": reason
178        });
179        let sse = format!(
180            "event: guardrail_violation\ndata: {}\n\ndata: [DONE]\n\n",
181            payload
182        );
183        Bytes::from(sse)
184    }
185
186    /// 設定への参照を取得する。
187    pub fn config(&self) -> &StreamingVerifierConfig {
188        &self.config
189    }
190
191    /// 前提コンテキスト文字列への参照を取得する。
192    pub fn state_context(&self) -> &str {
193        &self.state_context
194    }
195}
196
197/// SSE ストリームをインターセプトし、文境界ごとの投機的ハルシネーション検証とサーキットブレークを行う。
198///
199/// ハルシネーションを検知した場合は、安全な置換エラーイベントを注入してストリームを遮断する。
200pub fn intercept_sse_stream<S, E>(
201    stream: S,
202    interceptor: Arc<SpeculativeStreamInterceptor>,
203) -> BoxStream<'static, Result<Bytes, E>>
204where
205    S: Stream<Item = Result<Bytes, E>> + Send + 'static,
206    E: Send + 'static,
207{
208    use futures_util::StreamExt;
209
210    let mut buffer = SentenceChunkBuffer::new();
211    let mut terminated = false;
212
213    let stream = stream.flat_map(move |item| {
214        if terminated {
215            return futures_util::stream::empty().boxed();
216        }
217
218        match item {
219            Ok(bytes) => {
220                let text = String::from_utf8_lossy(&bytes);
221                let sentences =
222                    buffer.push_and_drain_sentences(&text, interceptor.config.min_sentence_chars);
223
224                for sentence in sentences {
225                    let verification = interceptor.verify_sentence(&sentence, None);
226                    if verification.is_hallucination {
227                        terminated = true;
228                        let err_reason = verification
229                            .reason
230                            .unwrap_or_else(|| "事実不整合".to_string());
231                        let err_bytes = interceptor.build_error_event_bytes(&err_reason);
232                        return futures_util::stream::iter(vec![Ok(err_bytes)]).boxed();
233                    }
234                }
235
236                futures_util::stream::iter(vec![Ok(bytes)]).boxed()
237            }
238            Err(e) => futures_util::stream::iter(vec![Err(e)]).boxed(),
239        }
240    });
241
242    stream.boxed()
243}
244
245#[cfg(test)]
246mod tests {
247    use super::*;
248
249    #[test]
250    fn test_sentence_chunk_buffer_drain() {
251        let mut buf = SentenceChunkBuffer::new();
252        let chunks = vec![
253            "本日の",
254            "契約金額は",
255            "1,000万円です。",
256            "次回は",
257            "来月です。",
258        ];
259
260        let mut sentences = Vec::new();
261        for c in chunks {
262            let drained = buf.push_and_drain_sentences(c, 5);
263            sentences.extend(drained);
264        }
265        assert_eq!(sentences.len(), 2);
266        assert_eq!(sentences[0], "本日の契約金額は1,000万円です。");
267        assert_eq!(sentences[1], "次回は来月です。");
268    }
269
270    #[test]
271    fn test_extract_content_from_sse() {
272        let line = r#"data: {"choices":[{"delta":{"content":"認証完了。"}}]}"#;
273        let content = extract_content_from_sse_line(line);
274        assert_eq!(content.as_deref(), Some("認証完了。"));
275    }
276
277    #[test]
278    fn test_interceptor_error_generation() {
279        let config = StreamingVerifierConfig::default();
280        let interceptor = SpeculativeStreamInterceptor::new(config, "前提コンテキスト".to_string());
281        let err_bytes = interceptor.build_error_event_bytes("テスト理由");
282        let s = String::from_utf8_lossy(&err_bytes);
283        assert!(s.contains("event: guardrail_violation"));
284        assert!(s.contains("data: [DONE]"));
285    }
286
287    #[tokio::test]
288    async fn test_intercept_sse_stream_circuit_breaker() {
289        use futures_util::StreamExt;
290
291        let config = StreamingVerifierConfig::default();
292        // 前提: 契約金は 1,000万円
293        let interceptor = Arc::new(SpeculativeStreamInterceptor::new(
294            config,
295            "前提: 契約金額は 1,000万円です。".to_string(),
296        ));
297
298        // 正常文のあとにハルシネーション文 (1,000億円) を流す
299        let chunks = vec![
300            Ok::<_, std::io::Error>(Bytes::from("data: 本件の確認を開始します。\n\n")),
301            Ok(Bytes::from(
302                "data: 契約金額は 1,000億円に確定しました。\n\n",
303            )),
304            Ok(Bytes::from("data: 完了しました。\n\n")),
305        ];
306        let raw_stream = futures_util::stream::iter(chunks);
307
308        let verified_stream = intercept_sse_stream(raw_stream, interceptor);
309        let items: Vec<Result<Bytes, std::io::Error>> = verified_stream.collect().await;
310
311        let texts: Vec<String> = items
312            .into_iter()
313            .map(|r| String::from_utf8_lossy(&r.unwrap()).to_string())
314            .collect();
315
316        // 1つ目は通過
317        assert!(texts[0].contains("本件の確認を開始します。"));
318        // 2つ目でハルシネーション (Entity Mismatch) を検知して遮断イベントへ置換
319        assert!(texts[1].contains("event: guardrail_violation"));
320        assert!(texts[1].contains("[安全基準に基づき生成を中断しました"));
321        // 3つ目以降はストリームが終了して送出されない
322        assert_eq!(texts.len(), 2);
323    }
324}
sokuto_server/guardrails/streaming.rs

sokuto_server/guardrails/
streaming.rs
