veil_engine/
stream.rs

1//! Restoration for streamed text, where a placeholder can be split across chunks.
2
3use crate::{Engine, Report, ScopeId};
4
5/// Restores one block of streamed model text.
6///
7/// Text is released as soon as it cannot be part of a placeholder. A tail that
8/// could still become one (for the default format, an open `[` followed only
9/// by placeholder characters) is held until it completes, stops matching,
10/// grows past the longest possible placeholder, or [`StreamRestorer::finish`]
11/// is called.
12///
13/// It holds no reference to the engine, so it can live inside a long-running
14/// response stream; the engine and scope are passed to each call.
15#[derive(Default)]
16pub struct StreamRestorer {
17    pending: String,
18    /// The last character released, as the model wrote it.
19    last: Option<char>,
20    report: Report,
21}
22
23impl StreamRestorer {
24    /// Adds a chunk and returns the restored text that is safe to send now.
25    pub fn push(&mut self, engine: &Engine, scope: &ScopeId, chunk: &str) -> String {
26        self.pending.push_str(chunk);
27        let rules = engine.rules();
28        let hold_from = match rules.fuzzy_restore {
29            true => rules.placeholder.hold_from_fuzzy(&self.pending),
30            false => rules.placeholder.hold_from(&self.pending),
31        };
32        let ready: String = self.pending.drain(..hold_from).collect();
33        self.restore(engine, scope, &ready)
34    }
35
36    /// Flushes any held text at the end of the block.
37    pub fn finish(&mut self, engine: &Engine, scope: &ScopeId) -> String {
38        let rest = std::mem::take(&mut self.pending);
39        self.restore(engine, scope, &rest)
40    }
41
42    /// True while text is being held back.
43    pub fn is_holding(&self) -> bool {
44        !self.pending.is_empty()
45    }
46
47    /// Counts so far for this block.
48    pub fn report(&self) -> &Report {
49        &self.report
50    }
51
52    fn restore(&mut self, engine: &Engine, scope: &ScopeId, text: &str) -> String {
53        if text.is_empty() {
54            return String::new();
55        }
56        let r = engine.restore_after(scope, self.last, text, str::to_owned);
57        self.last = text.chars().next_back();
58        self.report.merge(&r.report);
59        r.text
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use proptest::prelude::*;
66
67    use super::StreamRestorer;
68    use crate::{Engine, Options, ScopeId};
69
70    fn setup() -> (Engine, ScopeId) {
71        let engine = Engine::default();
72        let scope = ScopeId::profile("default");
73        engine.redact_text(&scope, "ana@example.com bob@example.org").unwrap();
74        (engine, scope)
75    }
76
77    fn stream(engine: &Engine, scope: &ScopeId, chunks: &[&str]) -> String {
78        let mut r = StreamRestorer::default();
79        let mut out: String = chunks.iter().map(|c| r.push(engine, scope, c)).collect();
80        out.push_str(&r.finish(engine, scope));
81        out
82    }
83
84    #[test]
85    fn placeholder_split_across_chunks() {
86        let (e, s) = setup();
87        assert_eq!(stream(&e, &s, &["Hi [EM", "AIL_", "1], bye"]), "Hi ana@example.com, bye");
88    }
89
90    #[test]
91    fn releases_text_before_the_open_bracket_immediately() {
92        let (e, s) = setup();
93        let mut r = StreamRestorer::default();
94        assert_eq!(r.push(&e, &s, "Hello there [EMA"), "Hello there ");
95        assert!(r.is_holding());
96        assert_eq!(r.push(&e, &s, "IL_2]!"), "bob@example.org!");
97        assert!(!r.is_holding());
98    }
99
100    #[test]
101    fn non_placeholder_brackets_are_not_held() {
102        let (e, s) = setup();
103        let mut r = StreamRestorer::default();
104        assert_eq!(r.push(&e, &s, "a [link](x) and [b"), "a [link](x) and [b");
105        assert_eq!(r.push(&e, &s, "[EMAIL_"), "");
106        assert_eq!(r.finish(&e, &s), "[EMAIL_");
107    }
108
109    #[test]
110    fn long_brackets_are_released() {
111        let (e, s) = setup();
112        let mut r = StreamRestorer::default();
113        let long = format!("[{}", "A".repeat(60));
114        assert_eq!(r.push(&e, &s, &long), long);
115    }
116
117    #[test]
118    fn custom_format_streams() {
119        let options = Options { placeholder_format: "<<{TYPE}:{N}>>".into(), ..Options::default() };
120        let e = Engine::builder().options(options).build().unwrap();
121        let s = ScopeId::profile("default");
122        assert_eq!(e.redact_text(&s, "ana@example.com").unwrap().text, "<<EMAIL:1>>");
123        assert_eq!(stream(&e, &s, &["x <", "<EMAIL:", "1>", "> y"]), "x ana@example.com y");
124    }
125
126    fn fuzzy_setup() -> (Engine, ScopeId) {
127        let options = Options { fuzzy_restore: true, ..Options::default() };
128        let engine = Engine::builder().options(options).build().unwrap();
129        let scope = ScopeId::profile("default");
130        engine.redact_text(&scope, "ana@example.com bob@example.org").unwrap();
131        (engine, scope)
132    }
133
134    #[test]
135    fn fuzzy_variants_split_across_chunks() {
136        let (e, s) = fuzzy_setup();
137        assert_eq!(
138            stream(&e, &s, &["Hi [ em", "ail 1 ] and EMA", "IL_2", ". Bye"]),
139            "Hi ana@example.com and bob@example.org. Bye"
140        );
141        assert_eq!(stream(&e, &s, &[r"\[EMAIL\", r"_2\]"]), "bob@example.org");
142        // Plain words are not held for long: only a trailing uppercase word.
143        let mut r = StreamRestorer::default();
144        assert_eq!(r.push(&e, &s, "Hello there. Nice"), "Hello there. Nice");
145        assert_eq!(r.push(&e, &s, " day, OK"), " day, ");
146    }
147
148    fn model_text() -> impl Strategy<Value = String> {
149        let piece = prop_oneof![
150            Just("[EMAIL_1]".to_owned()),
151            Just("[EMAIL_2]".to_owned()),
152            Just("[EMAIL_7]".to_owned()),
153            Just("[".to_owned()),
154            Just("]".to_owned()),
155            Just("[EMAIL".to_owned()),
156            Just("[ email 1 ]".to_owned()),
157            Just("EMAIL_2".to_owned()),
158            Just("\\[EMAIL\\_1\\]".to_owned()),
159            "[a-zA-Z0-9 _\\[\\]\\n,.é\\\\-]{0,12}",
160        ];
161        prop::collection::vec(piece, 0..20).prop_map(|v| v.concat())
162    }
163
164    proptest! {
165        #[test]
166        fn streaming_matches_whole_text_restore(text in model_text(), cuts in prop::collection::vec(any::<prop::sample::Index>(), 0..12)) {
167            let (e, s) = setup();
168            let mut bounds: Vec<usize> = cuts.iter().map(|i| i.index(text.len() + 1)).filter(|&b| text.is_char_boundary(b)).collect();
169            bounds.sort_unstable();
170            bounds.dedup();
171            let mut chunks = Vec::new();
172            let mut last = 0;
173            for b in bounds { chunks.push(&text[last..b]); last = b; }
174            chunks.push(&text[last..]);
175            let streamed = stream(&e, &s, &chunks);
176            let whole = e.restore_text(&s, &text).text;
177            prop_assert_eq!(streamed, whole);
178            let (e, s) = fuzzy_setup();
179            let streamed = stream(&e, &s, &chunks);
180            let whole = e.restore_text(&s, &text).text;
181            prop_assert_eq!(streamed, whole);
182        }
183    }
184}
veil_engine/stream.rs

veil_engine/
stream.rs
