lc_agents/hooks/
content_filter.rs

1// lc-agents/src/hooks/content_filter.rs
2//! ContentFilterHook — filters or replaces sensitive words in stream output.
3//!
4//! Scans each streaming token for sensitive words and either drops the token
5//! or replaces the sensitive content with a placeholder.
6
7use async_trait::async_trait;
8
9use super::{AgentHook, StreamAction};
10
11/// A hook that filters sensitive words from streaming output.
12///
13/// When `on_stream_chunk` is invoked, it checks the token for any of the
14/// configured sensitive words. If found, the token is either filtered (dropped)
15/// or replaced with a placeholder.
16///
17/// # Example
18///
19/// ```rust,ignore
20/// use lc_agents::hooks::ContentFilterHook;
21///
22/// let hook = ContentFilterHook::new(vec!["secret".to_string(), "password".to_string()]);
23/// let executor = AgentExecutor::new(agent, tools).hook(hook);
24/// ```
25pub struct ContentFilterHook {
26    /// Words to filter from the stream.
27    sensitive_words: Vec<String>,
28    /// Placeholder to replace sensitive words with.
29    placeholder: String,
30    /// If true, drop the entire token if it contains a sensitive word.
31    /// If false, replace the sensitive word with the placeholder.
32    drop_token: bool,
33}
34
35impl ContentFilterHook {
36    /// Creates a new ContentFilterHook with the given sensitive words.
37    pub fn new(sensitive_words: Vec<String>) -> Self {
38        Self {
39            sensitive_words,
40            placeholder: "[REDACTED]".to_string(),
41            drop_token: false,
42        }
43    }
44
45    /// Sets the placeholder for replaced words.
46    pub fn with_placeholder(mut self, placeholder: impl Into<String>) -> Self {
47        self.placeholder = placeholder.into();
48        self
49    }
50
51    /// Sets whether to drop entire tokens containing sensitive words.
52    pub fn with_drop_token(mut self, drop: bool) -> Self {
53        self.drop_token = drop;
54        self
55    }
56
57    /// Checks if the text contains any sensitive word.
58    fn contains_sensitive(&self, text: &str) -> bool {
59        let text_lower = text.to_lowercase();
60        self.sensitive_words
61            .iter()
62            .any(|word| text_lower.contains(&word.to_lowercase()))
63    }
64
65    /// Replaces sensitive words in the text with the placeholder.
66    fn replace_sensitive(&self, text: &str) -> String {
67        let mut result = text.to_string();
68        for word in &self.sensitive_words {
69            // Case-insensitive replacement
70            let lower = text.to_lowercase();
71            let mut start = 0;
72            while let Some(pos) = lower[start..].find(&word.to_lowercase()) {
73                let actual_pos = start + pos;
74                let end = actual_pos + word.len();
75                result = format!(
76                    "{}{}{}",
77                    &result[..actual_pos],
78                    self.placeholder,
79                    &result[end..]
80                );
81                start = actual_pos + self.placeholder.len();
82                // Re-check the modified string
83                let new_lower = result.to_lowercase();
84                if start >= new_lower.len() {
85                    break;
86                }
87                if !new_lower[start..].contains(&word.to_lowercase()) {
88                    break;
89                }
90            }
91        }
92        result
93    }
94}
95
96#[async_trait]
97impl AgentHook for ContentFilterHook {
98    fn on_stream_chunk(&self, chunk: &str) -> StreamAction {
99        if !self.contains_sensitive(chunk) {
100            return StreamAction::Forward(chunk.to_string());
101        }
102
103        if self.drop_token {
104            StreamAction::Filter
105        } else {
106            StreamAction::Replace(self.replace_sensitive(chunk))
107        }
108    }
109}
lc_agents/hooks/content_filter.rs

lc_agents/hooks/
content_filter.rs
