atman_runtime/
safety.rs

1use std::sync::Arc;
2
3use crate::error::RuntimeError;
4use crate::tool::BoxFut;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum ScanVerdict {
8    Pass,
9    Warn(Vec<String>),
10    Deny(Vec<String>),
11}
12
13impl ScanVerdict {
14    pub fn is_pass(&self) -> bool {
15        matches!(self, ScanVerdict::Pass)
16    }
17
18    pub fn is_deny(&self) -> bool {
19        matches!(self, ScanVerdict::Deny(_))
20    }
21
22    pub fn categories(&self) -> &[String] {
23        match self {
24            ScanVerdict::Pass => &[],
25            ScanVerdict::Warn(c) | ScanVerdict::Deny(c) => c,
26        }
27    }
28}
29
30pub trait SafetyClassifier: Send + Sync {
31    fn scan<'a>(&'a self, text: &'a str) -> BoxFut<'a, Result<ScanVerdict, RuntimeError>>;
32    fn kind(&self) -> &'static str;
33}
34
35pub struct NoopClassifier;
36
37impl SafetyClassifier for NoopClassifier {
38    fn scan<'a>(&'a self, _text: &'a str) -> BoxFut<'a, Result<ScanVerdict, RuntimeError>> {
39        Box::pin(async { Ok(ScanVerdict::Pass) })
40    }
41
42    fn kind(&self) -> &'static str {
43        "noop"
44    }
45}
46
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, documented::DocumentedVariants)]
48pub enum SafetyMode {
49    #[default]
50    /// Log a warning but allow the content through.
51    Warn,
52    /// Block the content entirely.
53    Deny,
54}
55
56pub struct OpenAiModerationClassifier {
57    api_key: String,
58    base_url: String,
59    client: reqwest::Client,
60    model: String,
61    timeout: std::time::Duration,
62}
63
64impl OpenAiModerationClassifier {
65    pub fn new(api_key: impl Into<String>) -> Self {
66        Self {
67            api_key: api_key.into(),
68            base_url: "https://api.openai.com/v1".to_string(),
69            client: reqwest::Client::new(),
70            model: "omni-moderation-latest".to_string(),
71            timeout: std::time::Duration::from_secs(5),
72        }
73    }
74
75    pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
76        self.base_url = url.into();
77        self
78    }
79
80    pub fn with_model(mut self, model: impl Into<String>) -> Self {
81        self.model = model.into();
82        self
83    }
84
85    pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
86        self.timeout = timeout;
87        self
88    }
89}
90
91impl SafetyClassifier for OpenAiModerationClassifier {
92    fn scan<'a>(&'a self, text: &'a str) -> BoxFut<'a, Result<ScanVerdict, RuntimeError>> {
93        Box::pin(async move {
94            let body = serde_json::json!({ "input": text, "model": self.model });
95            let call = self
96                .client
97                .post(format!("{}/moderations", self.base_url))
98                .bearer_auth(&self.api_key)
99                .json(&body)
100                .send();
101            let resp = tokio::time::timeout(self.timeout, call)
102                .await
103                .map_err(|_| {
104                    RuntimeError::ToolFailed(format!(
105                        "safety: openai moderation timed out after {}s",
106                        self.timeout.as_secs()
107                    ))
108                })?;
109            let resp = resp
110                .map_err(|e| RuntimeError::ToolFailed(format!("safety: openai moderation: {e}")))?;
111            let status = resp.status();
112            if !status.is_success() {
113                let body = resp.text().await.unwrap_or_default();
114                return Err(RuntimeError::ToolFailed(format!(
115                    "safety: openai moderation http {status}: {body}"
116                )));
117            }
118            let doc: serde_json::Value = resp
119                .json()
120                .await
121                .map_err(|e| RuntimeError::ToolFailed(format!("safety: parse moderation: {e}")))?;
122            let Some(result) = doc.get("results").and_then(|r| r.get(0)) else {
123                return Ok(ScanVerdict::Pass);
124            };
125            let flagged = result
126                .get("flagged")
127                .and_then(|f| f.as_bool())
128                .unwrap_or(false);
129            if !flagged {
130                return Ok(ScanVerdict::Pass);
131            }
132            let categories = result
133                .get("categories")
134                .and_then(|c| c.as_object())
135                .map(|obj| {
136                    obj.iter()
137                        .filter(|(_, v)| v.as_bool().unwrap_or(false))
138                        .map(|(k, _)| k.clone())
139                        .collect::<Vec<_>>()
140                })
141                .unwrap_or_default();
142            Ok(ScanVerdict::Deny(categories))
143        })
144    }
145
146    fn kind(&self) -> &'static str {
147        "openai-moderation"
148    }
149}
150
151/// Configuration for safety checks on LLM input/output.
152#[derive(documented::Documented, documented::DocumentedFields)]
153pub struct SafetyConfig {
154    /// Whether safety checks are active.
155    pub enabled: bool,
156    /// Action to take when unsafe content is detected.
157    pub mode: SafetyMode,
158    /// If true, attempt to auto-rewrite unsafe content instead of blocking.
159    pub auto_rewrite: bool,
160    /// The classifier implementation used to scan content.
161    pub classifier: Arc<dyn SafetyClassifier>,
162}
163
164impl Clone for SafetyConfig {
165    fn clone(&self) -> Self {
166        Self {
167            enabled: self.enabled,
168            mode: self.mode,
169            auto_rewrite: self.auto_rewrite,
170            classifier: self.classifier.clone(),
171        }
172    }
173}
174
175impl SafetyConfig {
176    pub fn noop() -> Self {
177        Self {
178            enabled: false,
179            mode: SafetyMode::Warn,
180            auto_rewrite: false,
181            classifier: Arc::new(NoopClassifier),
182        }
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189    use wiremock::matchers::{header, method, path};
190    use wiremock::{Mock, MockServer, ResponseTemplate};
191
192    #[tokio::test]
193    async fn noop_classifier_always_passes() {
194        let c = NoopClassifier;
195        assert_eq!(c.scan("anything").await.unwrap(), ScanVerdict::Pass);
196        assert_eq!(c.kind(), "noop");
197    }
198
199    #[tokio::test]
200    async fn openai_moderation_passes_when_flagged_is_false() {
201        let server = MockServer::start().await;
202        Mock::given(method("POST"))
203            .and(path("/moderations"))
204            .and(header("authorization", "Bearer test-key"))
205            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
206                "results": [{
207                    "flagged": false,
208                    "categories": {"hate": false, "violence": false}
209                }]
210            })))
211            .mount(&server)
212            .await;
213        let c = OpenAiModerationClassifier::new("test-key").with_base_url(server.uri());
214        let out = c.scan("hello world").await.unwrap();
215        assert_eq!(out, ScanVerdict::Pass);
216    }
217
218    #[tokio::test]
219    async fn openai_moderation_returns_deny_with_flagged_categories() {
220        let server = MockServer::start().await;
221        Mock::given(method("POST"))
222            .and(path("/moderations"))
223            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
224                "results": [{
225                    "flagged": true,
226                    "categories": {"hate": true, "violence": false, "self-harm": true}
227                }]
228            })))
229            .mount(&server)
230            .await;
231        let c = OpenAiModerationClassifier::new("k").with_base_url(server.uri());
232        let out = c.scan("bad").await.unwrap();
233        let mut cats: Vec<String> = match out {
234            ScanVerdict::Deny(c) => c,
235            other => panic!("expected Deny, got {other:?}"),
236        };
237        cats.sort();
238        assert_eq!(cats, vec!["hate".to_string(), "self-harm".to_string()]);
239    }
240
241    #[tokio::test]
242    async fn openai_moderation_bubbles_http_error() {
243        let server = MockServer::start().await;
244        Mock::given(method("POST"))
245            .and(path("/moderations"))
246            .respond_with(ResponseTemplate::new(401).set_body_string("bad key"))
247            .mount(&server)
248            .await;
249        let c = OpenAiModerationClassifier::new("k").with_base_url(server.uri());
250        let err = c.scan("hi").await.unwrap_err();
251        let msg = format!("{err}");
252        assert!(msg.contains("401"), "err: {msg}");
253    }
254}
atman_runtime/safety.rs

atman_runtime/
safety.rs
