atman_runtime/
injection_classifier.rs

1use std::sync::Arc;
2use std::time::Duration;
3
4use regex::Regex;
5
6use crate::injection::InjectionLevel;
7use crate::tool::BoxFut;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum ClassifierSource {
11    Prefix,
12    Rule,
13    Llm,
14    Default,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct Classification {
19    pub level: InjectionLevel,
20    pub redirect_target: Option<String>,
21    pub source: ClassifierSource,
22}
23
24pub trait InjectionClassifier: Send + Sync {
25    fn classify<'a>(&'a self, text: &'a str) -> BoxFut<'a, Classification>;
26    fn kind(&self) -> &'static str;
27}
28
29pub struct RuleClassifier {
30    l4: Regex,
31    l3: Regex,
32    l2: Regex,
33}
34
35impl Default for RuleClassifier {
36    fn default() -> Self {
37        Self {
38            l4: Regex::new(r"^(?i)\s*(停|停下|停止|stop|abort|halt|kill|终止)[\s!.。！]*$")
39                .unwrap(),
40            l3: Regex::new(r"(?i)^\s*(换成|切换到|redirect to|switch to)\s+(\S+)").unwrap(),
41            l2: Regex::new(r"(别用|不要|不能|改成|错了|应该|请用|do not|don't use|stop using)")
42                .unwrap(),
43        }
44    }
45}
46
47impl RuleClassifier {
48    fn classify_sync(&self, text: &str) -> Classification {
49        let trimmed = text.trim();
50        if self.l4.is_match(trimmed) {
51            return Classification {
52                level: InjectionLevel::L4HardStop,
53                redirect_target: None,
54                source: ClassifierSource::Rule,
55            };
56        }
57        if let Some(caps) = self.l3.captures(trimmed) {
58            let target = caps.get(2).map(|m| m.as_str().to_string());
59            return Classification {
60                level: InjectionLevel::L3Redirect,
61                redirect_target: target,
62                source: ClassifierSource::Rule,
63            };
64        }
65        if self.l2.is_match(trimmed) {
66            return Classification {
67                level: InjectionLevel::L2CourseCorrect,
68                redirect_target: None,
69                source: ClassifierSource::Rule,
70            };
71        }
72        Classification {
73            level: InjectionLevel::L1Nudge,
74            redirect_target: None,
75            source: ClassifierSource::Default,
76        }
77    }
78}
79
80impl InjectionClassifier for RuleClassifier {
81    fn classify<'a>(&'a self, text: &'a str) -> BoxFut<'a, Classification> {
82        let out = self.classify_sync(text);
83        Box::pin(async move { out })
84    }
85    fn kind(&self) -> &'static str {
86        "rule"
87    }
88}
89
90pub struct LlmClassifier {
91    provider: Arc<dyn crate::provider::Provider>,
92    model: String,
93}
94
95impl LlmClassifier {
96    pub fn new(provider: Arc<dyn crate::provider::Provider>, model: impl Into<String>) -> Self {
97        Self {
98            provider,
99            model: model.into(),
100        }
101    }
102}
103
104impl InjectionClassifier for LlmClassifier {
105    fn classify<'a>(&'a self, text: &'a str) -> BoxFut<'a, Classification> {
106        Box::pin(async move {
107            let prompt = format!(
108                "You are classifying a user interruption sent to a running AI agent. Reply STRICTLY as one line of JSON, no prose.\n\n\
109                 Levels:\n\
110                 - L1: minor nudge (add info, remind, small hint)\n\
111                 - L2: course-correct (wrong approach mid-stream, forbidden pattern)\n\
112                 - L3: redirect (switch flows entirely, name a target if user gave one)\n\
113                 - L4: hard stop (kill flow immediately)\n\n\
114                 Reply shape: {{\"level\": \"L1\"|\"L2\"|\"L3\"|\"L4\", \"redirect_target\": null | \"<flow_name>\"}}\n\n\
115                 User interruption: {text}"
116            );
117            let req = crate::provider::LlmRequest {
118                model: self.model.clone(),
119                messages: vec![crate::provider::user_text_message(prompt)],
120                system: None,
121                input: crate::value::Value::Unit,
122                schema: None,
123                cache_prompt: false,
124                tools: Vec::new(),
125                thinking_enabled: false,
126                stall_timeout_secs: 0,
127            };
128            let am = match self.provider.call(req).await {
129                Ok(am) => am,
130                Err(_) => return default_l1(ClassifierSource::Default),
131            };
132            let body = am.message.text_concat();
133            let Some(parsed) =
134                extract_json(&body).and_then(|s| serde_json::from_str::<ClassifyJson>(&s).ok())
135            else {
136                return default_l1(ClassifierSource::Default);
137            };
138            Classification {
139                level: parse_level(&parsed.level),
140                redirect_target: parsed.redirect_target,
141                source: ClassifierSource::Llm,
142            }
143        })
144    }
145    fn kind(&self) -> &'static str {
146        "llm"
147    }
148}
149
150#[derive(serde::Deserialize)]
151struct ClassifyJson {
152    level: String,
153    #[serde(default)]
154    redirect_target: Option<String>,
155}
156
157fn parse_level(s: &str) -> InjectionLevel {
158    match s.trim().to_ascii_lowercase().as_str() {
159        "l4" => InjectionLevel::L4HardStop,
160        "l3" => InjectionLevel::L3Redirect,
161        "l2" => InjectionLevel::L2CourseCorrect,
162        _ => InjectionLevel::L1Nudge,
163    }
164}
165
166fn extract_json(body: &str) -> Option<String> {
167    let start = body.find('{')?;
168    let end = body.rfind('}')?;
169    if end < start {
170        return None;
171    }
172    Some(body[start..=end].to_string())
173}
174
175fn default_l1(source: ClassifierSource) -> Classification {
176    Classification {
177        level: InjectionLevel::L1Nudge,
178        redirect_target: None,
179        source,
180    }
181}
182
183pub struct ComposedClassifier {
184    rule: RuleClassifier,
185    llm: Option<Arc<dyn InjectionClassifier>>,
186    llm_timeout: Duration,
187}
188
189impl ComposedClassifier {
190    pub fn new(rule: RuleClassifier) -> Self {
191        Self {
192            rule,
193            llm: None,
194            llm_timeout: Duration::from_secs(3),
195        }
196    }
197
198    pub fn with_llm(mut self, llm: Arc<dyn InjectionClassifier>, timeout: Duration) -> Self {
199        self.llm = Some(llm);
200        self.llm_timeout = timeout;
201        self
202    }
203}
204
205impl InjectionClassifier for ComposedClassifier {
206    fn classify<'a>(&'a self, text: &'a str) -> BoxFut<'a, Classification> {
207        Box::pin(async move {
208            let rule_hit = self.rule.classify_sync(text);
209            if rule_hit.level != InjectionLevel::L1Nudge {
210                return rule_hit;
211            }
212            let Some(llm) = self.llm.as_ref() else {
213                return rule_hit;
214            };
215            match tokio::time::timeout(self.llm_timeout, llm.classify(text)).await {
216                Ok(cls) => cls,
217                Err(_) => rule_hit,
218            }
219        })
220    }
221    fn kind(&self) -> &'static str {
222        "composed"
223    }
224}
225
226pub fn source_tag(source: ClassifierSource) -> &'static str {
227    match source {
228        ClassifierSource::Prefix => "prefix",
229        ClassifierSource::Rule => "rule",
230        ClassifierSource::Llm => "llm",
231        ClassifierSource::Default => "default",
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238
239    #[tokio::test]
240    async fn rule_matches_l4_stop_variants() {
241        let c = RuleClassifier::default();
242        for text in ["停", "停下!", "stop", "STOP", "abort.", "kill"] {
243            let out = c.classify(text).await;
244            assert_eq!(out.level, InjectionLevel::L4HardStop, "text: {text}");
245            assert_eq!(out.source, ClassifierSource::Rule);
246        }
247    }
248
249    #[tokio::test]
250    async fn rule_matches_l3_redirect_and_extracts_target() {
251        let c = RuleClassifier::default();
252        let out = c.classify("换成 review_code").await;
253        assert_eq!(out.level, InjectionLevel::L3Redirect);
254        assert_eq!(out.redirect_target.as_deref(), Some("review_code"));
255        let out = c.classify("switch to hello").await;
256        assert_eq!(out.level, InjectionLevel::L3Redirect);
257        assert_eq!(out.redirect_target.as_deref(), Some("hello"));
258    }
259
260    #[tokio::test]
261    async fn rule_matches_l2_course_correct_hints() {
262        let c = RuleClassifier::default();
263        for text in [
264            "别用 as any",
265            "不要写 unsafe",
266            "改成 async",
267            "错了 应该用 tokio",
268            "do not use blocking",
269        ] {
270            let out = c.classify(text).await;
271            assert_eq!(out.level, InjectionLevel::L2CourseCorrect, "text: {text}");
272        }
273    }
274
275    #[tokio::test]
276    async fn rule_falls_back_to_l1_nudge_default() {
277        let c = RuleClassifier::default();
278        let out = c.classify("记得加一句测试").await;
279        assert_eq!(out.level, InjectionLevel::L1Nudge);
280        assert_eq!(out.source, ClassifierSource::Default);
281    }
282
283    #[tokio::test]
284    async fn composed_returns_rule_hit_when_matched() {
285        let c = ComposedClassifier::new(RuleClassifier::default());
286        let out = c.classify("停").await;
287        assert_eq!(out.level, InjectionLevel::L4HardStop);
288    }
289}
atman_runtime/injection_classifier.rs

atman_runtime/
injection_classifier.rs
