abu_agent/middleware/
privacy.rs

1use abu_base::chat::AssistantMessage;
2use regex::Regex;
3use std::convert::Infallible;
4
5use super::{LlmOutMiddleware, MiddlewareFlow};
6
7pub struct PrivacyScrubberMiddleware {
8    phone_regex: Regex,
9    email_regex: Regex,
10}
11
12impl PrivacyScrubberMiddleware {
13    pub fn new() -> Self {
14        Self {
15            // TODO
16            phone_regex: Regex::new(r"1[3-9]\d{9}").unwrap(),
17            email_regex: Regex::new(r"[a-zA-Z0-9_-]+@[a-zA-Z0-9_-]+(\.[a-zA-Z0-9_-]+)+").unwrap(),
18        }
19    }
20}
21
22#[async_trait::async_trait]
23impl LlmOutMiddleware for PrivacyScrubberMiddleware {
24    type Error = Infallible;
25
26    async fn intercept(&self, ai_message: &mut AssistantMessage) -> Result<MiddlewareFlow, Self::Error> {
27        let scrubbed = self.phone_regex.replace_all(&ai_message.content, "***-****-****");
28        let scrubbed = self.email_regex.replace_all(&scrubbed, "[EMAIL PROTECTED]");
29
30        ai_message.content = scrubbed.to_string();
31
32        Ok(MiddlewareFlow::Continue)
33    }
34}
abu_agent/middleware/privacy.rs

abu_agent/middleware/
privacy.rs
