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)]
48pub enum SafetyMode {
49 #[default]
50 Warn,
51 Deny,
52}
53
54pub struct OpenAiModerationClassifier {
55 api_key: String,
56 base_url: String,
57 client: reqwest::Client,
58 model: String,
59 timeout: std::time::Duration,
60}
61
62impl OpenAiModerationClassifier {
63 pub fn new(api_key: impl Into<String>) -> Self {
64 Self {
65 api_key: api_key.into(),
66 base_url: "https://api.openai.com/v1".to_string(),
67 client: reqwest::Client::new(),
68 model: "omni-moderation-latest".to_string(),
69 timeout: std::time::Duration::from_secs(5),
70 }
71 }
72
73 pub fn with_base_url(mut self, url: impl Into<String>) -> Self {
74 self.base_url = url.into();
75 self
76 }
77
78 pub fn with_model(mut self, model: impl Into<String>) -> Self {
79 self.model = model.into();
80 self
81 }
82
83 pub fn with_timeout(mut self, timeout: std::time::Duration) -> Self {
84 self.timeout = timeout;
85 self
86 }
87}
88
89impl SafetyClassifier for OpenAiModerationClassifier {
90 fn scan<'a>(&'a self, text: &'a str) -> BoxFut<'a, Result<ScanVerdict, RuntimeError>> {
91 Box::pin(async move {
92 let body = serde_json::json!({ "input": text, "model": self.model });
93 let call = self
94 .client
95 .post(format!("{}/moderations", self.base_url))
96 .bearer_auth(&self.api_key)
97 .json(&body)
98 .send();
99 let resp = tokio::time::timeout(self.timeout, call)
100 .await
101 .map_err(|_| {
102 RuntimeError::ToolFailed(format!(
103 "safety: openai moderation timed out after {}s",
104 self.timeout.as_secs()
105 ))
106 })?;
107 let resp = resp
108 .map_err(|e| RuntimeError::ToolFailed(format!("safety: openai moderation: {e}")))?;
109 let status = resp.status();
110 if !status.is_success() {
111 let body = resp.text().await.unwrap_or_default();
112 return Err(RuntimeError::ToolFailed(format!(
113 "safety: openai moderation http {status}: {body}"
114 )));
115 }
116 let doc: serde_json::Value = resp
117 .json()
118 .await
119 .map_err(|e| RuntimeError::ToolFailed(format!("safety: parse moderation: {e}")))?;
120 let Some(result) = doc.get("results").and_then(|r| r.get(0)) else {
121 return Ok(ScanVerdict::Pass);
122 };
123 let flagged = result
124 .get("flagged")
125 .and_then(|f| f.as_bool())
126 .unwrap_or(false);
127 if !flagged {
128 return Ok(ScanVerdict::Pass);
129 }
130 let categories = result
131 .get("categories")
132 .and_then(|c| c.as_object())
133 .map(|obj| {
134 obj.iter()
135 .filter(|(_, v)| v.as_bool().unwrap_or(false))
136 .map(|(k, _)| k.clone())
137 .collect::<Vec<_>>()
138 })
139 .unwrap_or_default();
140 Ok(ScanVerdict::Deny(categories))
141 })
142 }
143
144 fn kind(&self) -> &'static str {
145 "openai-moderation"
146 }
147}
148
149pub struct SafetyConfig {
150 pub enabled: bool,
151 pub mode: SafetyMode,
152 pub auto_rewrite: bool,
153 pub classifier: Arc<dyn SafetyClassifier>,
154}
155
156impl Clone for SafetyConfig {
157 fn clone(&self) -> Self {
158 Self {
159 enabled: self.enabled,
160 mode: self.mode,
161 auto_rewrite: self.auto_rewrite,
162 classifier: self.classifier.clone(),
163 }
164 }
165}
166
167impl SafetyConfig {
168 pub fn noop() -> Self {
169 Self {
170 enabled: false,
171 mode: SafetyMode::Warn,
172 auto_rewrite: false,
173 classifier: Arc::new(NoopClassifier),
174 }
175 }
176}
177
178#[cfg(test)]
179mod tests {
180 use super::*;
181 use wiremock::matchers::{header, method, path};
182 use wiremock::{Mock, MockServer, ResponseTemplate};
183
184 #[tokio::test]
185 async fn noop_classifier_always_passes() {
186 let c = NoopClassifier;
187 assert_eq!(c.scan("anything").await.unwrap(), ScanVerdict::Pass);
188 assert_eq!(c.kind(), "noop");
189 }
190
191 #[tokio::test]
192 async fn openai_moderation_passes_when_flagged_is_false() {
193 let server = MockServer::start().await;
194 Mock::given(method("POST"))
195 .and(path("/moderations"))
196 .and(header("authorization", "Bearer test-key"))
197 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
198 "results": [{
199 "flagged": false,
200 "categories": {"hate": false, "violence": false}
201 }]
202 })))
203 .mount(&server)
204 .await;
205 let c = OpenAiModerationClassifier::new("test-key").with_base_url(server.uri());
206 let out = c.scan("hello world").await.unwrap();
207 assert_eq!(out, ScanVerdict::Pass);
208 }
209
210 #[tokio::test]
211 async fn openai_moderation_returns_deny_with_flagged_categories() {
212 let server = MockServer::start().await;
213 Mock::given(method("POST"))
214 .and(path("/moderations"))
215 .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
216 "results": [{
217 "flagged": true,
218 "categories": {"hate": true, "violence": false, "self-harm": true}
219 }]
220 })))
221 .mount(&server)
222 .await;
223 let c = OpenAiModerationClassifier::new("k").with_base_url(server.uri());
224 let out = c.scan("bad").await.unwrap();
225 let mut cats: Vec<String> = match out {
226 ScanVerdict::Deny(c) => c,
227 other => panic!("expected Deny, got {other:?}"),
228 };
229 cats.sort();
230 assert_eq!(cats, vec!["hate".to_string(), "self-harm".to_string()]);
231 }
232
233 #[tokio::test]
234 async fn openai_moderation_bubbles_http_error() {
235 let server = MockServer::start().await;
236 Mock::given(method("POST"))
237 .and(path("/moderations"))
238 .respond_with(ResponseTemplate::new(401).set_body_string("bad key"))
239 .mount(&server)
240 .await;
241 let c = OpenAiModerationClassifier::new("k").with_base_url(server.uri());
242 let err = c.scan("hi").await.unwrap_err();
243 let msg = format!("{err}");
244 assert!(msg.contains("401"), "err: {msg}");
245 }
246}