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                prompt_cache_key: None,
125                tools: Vec::new(),
126                reasoning: crate::provider::ReasoningSelection::ProviderDefault,
127                stall_timeout_secs: 0,
128            };
129            let am = match self.provider.call(req).await {
130                Ok(am) => am,
131                Err(_) => return default_l1(ClassifierSource::Default),
132            };
133            let body = am.message.text_concat();
134            let Some(parsed) =
135                extract_json(&body).and_then(|s| serde_json::from_str::<ClassifyJson>(&s).ok())
136            else {
137                return default_l1(ClassifierSource::Default);
138            };
139            Classification {
140                level: parse_level(&parsed.level),
141                redirect_target: parsed.redirect_target,
142                source: ClassifierSource::Llm,
143            }
144        })
145    }
146    fn kind(&self) -> &'static str {
147        "llm"
148    }
149}
150
151#[derive(serde::Deserialize)]
152struct ClassifyJson {
153    level: String,
154    #[serde(default)]
155    redirect_target: Option<String>,
156}
157
158fn parse_level(s: &str) -> InjectionLevel {
159    match s.trim().to_ascii_lowercase().as_str() {
160        "l4" => InjectionLevel::L4HardStop,
161        "l3" => InjectionLevel::L3Redirect,
162        "l2" => InjectionLevel::L2CourseCorrect,
163        _ => InjectionLevel::L1Nudge,
164    }
165}
166
167fn extract_json(body: &str) -> Option<String> {
168    let start = body.find('{')?;
169    let end = body.rfind('}')?;
170    if end < start {
171        return None;
172    }
173    Some(body[start..=end].to_string())
174}
175
176fn default_l1(source: ClassifierSource) -> Classification {
177    Classification {
178        level: InjectionLevel::L1Nudge,
179        redirect_target: None,
180        source,
181    }
182}
183
184pub struct ComposedClassifier {
185    rule: RuleClassifier,
186    llm: Option<Arc<dyn InjectionClassifier>>,
187    llm_timeout: Duration,
188}
189
190impl ComposedClassifier {
191    pub fn new(rule: RuleClassifier) -> Self {
192        Self {
193            rule,
194            llm: None,
195            llm_timeout: Duration::from_secs(3),
196        }
197    }
198
199    pub fn with_llm(mut self, llm: Arc<dyn InjectionClassifier>, timeout: Duration) -> Self {
200        self.llm = Some(llm);
201        self.llm_timeout = timeout;
202        self
203    }
204}
205
206impl InjectionClassifier for ComposedClassifier {
207    fn classify<'a>(&'a self, text: &'a str) -> BoxFut<'a, Classification> {
208        Box::pin(async move {
209            let rule_hit = self.rule.classify_sync(text);
210            if rule_hit.level != InjectionLevel::L1Nudge {
211                return rule_hit;
212            }
213            let Some(llm) = self.llm.as_ref() else {
214                return rule_hit;
215            };
216            match tokio::time::timeout(self.llm_timeout, llm.classify(text)).await {
217                Ok(cls) => cls,
218                Err(_) => rule_hit,
219            }
220        })
221    }
222    fn kind(&self) -> &'static str {
223        "composed"
224    }
225}
226
227pub fn source_tag(source: ClassifierSource) -> &'static str {
228    match source {
229        ClassifierSource::Prefix => "prefix",
230        ClassifierSource::Rule => "rule",
231        ClassifierSource::Llm => "llm",
232        ClassifierSource::Default => "default",
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    #[tokio::test]
241    async fn rule_matches_l4_stop_variants() {
242        let c = RuleClassifier::default();
243        for text in ["停", "停下!", "stop", "STOP", "abort.", "kill"] {
244            let out = c.classify(text).await;
245            assert_eq!(out.level, InjectionLevel::L4HardStop, "text: {text}");
246            assert_eq!(out.source, ClassifierSource::Rule);
247        }
248    }
249
250    #[tokio::test]
251    async fn rule_matches_l3_redirect_and_extracts_target() {
252        let c = RuleClassifier::default();
253        let out = c.classify("换成 review_code").await;
254        assert_eq!(out.level, InjectionLevel::L3Redirect);
255        assert_eq!(out.redirect_target.as_deref(), Some("review_code"));
256        let out = c.classify("switch to hello").await;
257        assert_eq!(out.level, InjectionLevel::L3Redirect);
258        assert_eq!(out.redirect_target.as_deref(), Some("hello"));
259    }
260
261    #[tokio::test]
262    async fn rule_matches_l2_course_correct_hints() {
263        let c = RuleClassifier::default();
264        for text in [
265            "别用 as any",
266            "不要写 unsafe",
267            "改成 async",
268            "错了 应该用 tokio",
269            "do not use blocking",
270        ] {
271            let out = c.classify(text).await;
272            assert_eq!(out.level, InjectionLevel::L2CourseCorrect, "text: {text}");
273        }
274    }
275
276    #[tokio::test]
277    async fn rule_falls_back_to_l1_nudge_default() {
278        let c = RuleClassifier::default();
279        let out = c.classify("记得加一句测试").await;
280        assert_eq!(out.level, InjectionLevel::L1Nudge);
281        assert_eq!(out.source, ClassifierSource::Default);
282    }
283
284    #[tokio::test]
285    async fn composed_returns_rule_hit_when_matched() {
286        let c = ComposedClassifier::new(RuleClassifier::default());
287        let out = c.classify("停").await;
288        assert_eq!(out.level, InjectionLevel::L4HardStop);
289    }
290}
atman_runtime/injection_classifier.rs

atman_runtime/
injection_classifier.rs
