ultrafast_gateway/plugins/
content_filtering.rs

1use crate::gateway_error::GatewayError;
2use crate::plugins::PluginConfig;
3use axum::{body::Body, http::Request, response::Response};
4use std::collections::HashSet;
5use ultrafast_models_sdk::models::{ChatRequest, EmbeddingRequest, ImageRequest};
6
7#[derive(Clone, Debug)]
8pub struct ContentFilteringPlugin {
9    name: String,
10    enabled: bool,
11    blocked_words: HashSet<String>,
12    blocked_patterns: Vec<String>,
13    max_input_length: usize,
14}
15
16impl ContentFilteringPlugin {
17    pub fn new(config: &PluginConfig) -> Result<Self, GatewayError> {
18        let mut blocked_words = HashSet::new();
19
20        // Add default blocked words
21        blocked_words.insert("hate".to_string());
22        blocked_words.insert("violence".to_string());
23        blocked_words.insert("discrimination".to_string());
24
25        let blocked_patterns = vec![r"\b(hate|violence|discrimination)\b".to_string()];
26
27        let max_input_length = config
28            .config
29            .get("max_input_length")
30            .and_then(|v| v.as_u64())
31            .unwrap_or(10000) as usize;
32
33        Ok(Self {
34            name: config.name.clone(),
35            enabled: config.enabled,
36            blocked_words,
37            blocked_patterns,
38            max_input_length,
39        })
40    }
41
42    fn check_content(&self, content: &str) -> Result<(), GatewayError> {
43        // Check input length
44        if content.len() > self.max_input_length {
45            return Err(GatewayError::InvalidRequest {
46                message: format!(
47                    "Input too long: {} characters (max: {})",
48                    content.len(),
49                    self.max_input_length
50                ),
51            });
52        }
53
54        // Check for blocked words
55        let content_lower = content.to_lowercase();
56        for word in &self.blocked_words {
57            if content_lower.contains(word) {
58                return Err(GatewayError::ContentFiltered {
59                    message: format!("Content contains blocked word: {word}"),
60                });
61            }
62        }
63
64        // Check for blocked patterns (simplified regex check)
65        for pattern in &self.blocked_patterns {
66            if content_lower.contains(&pattern.replace(r"\b", "")) {
67                return Err(GatewayError::ContentFiltered {
68                    message: format!("Content matches blocked pattern: {pattern}"),
69                });
70            }
71        }
72
73        Ok(())
74    }
75
76    fn extract_content_from_request(&self, request: &Request<Body>) -> Option<String> {
77        // Extract content from request based on path and method
78        let path = request.uri().path();
79        let _method = request.method().as_str();
80
81        // For chat completions, we would extract the messages content
82        // For embeddings, we would extract the input text
83        // For image generation, we would extract the prompt
84
85        // Since we can't easily access the body here in the middleware,
86        // we'll check the Content-Type header and URL path to determine
87        // if this is a request type we should filter
88        if path.contains("/chat/completions")
89            || path.contains("/completions")
90            || path.contains("/embeddings")
91            || path.contains("/images/generations")
92        {
93            // In a real implementation, you would:
94            // 1. Read the request body (need to be careful about consuming it)
95            // 2. Parse the JSON
96            // 3. Extract the relevant text fields
97            // 4. Return the concatenated content
98
99            // For now, we'll extract content from headers and query parameters
100            let mut content_parts = Vec::new();
101
102            // Extract from query parameters
103            if let Some(query) = request.uri().query() {
104                // Parse query parameters manually
105                for param in query.split('&') {
106                    if let Some((key, value)) = param.split_once('=') {
107                        if key == "prompt" || key == "input" || key == "text" {
108                            content_parts.push(value.to_string());
109                        }
110                    }
111                }
112            }
113
114            // Extract from headers that might contain content
115            if let Some(content_type) = request.headers().get("content-type") {
116                if let Ok(content_type_str) = content_type.to_str() {
117                    if content_type_str.contains("application/json") {
118                        // This would be a JSON request, content would be in body
119                        content_parts.push("JSON request body".to_string());
120                    }
121                }
122            }
123
124            // Extract from user agent (might contain suspicious content)
125            if let Some(user_agent) = request.headers().get("user-agent") {
126                if let Ok(ua_str) = user_agent.to_str() {
127                    if ua_str.len() > 100 {
128                        content_parts.push(format!("Long user agent: {}", &ua_str[..50]));
129                    }
130                }
131            }
132
133            if content_parts.is_empty() {
134                None
135            } else {
136                Some(content_parts.join(" | "))
137            }
138        } else {
139            None
140        }
141    }
142}
143
144impl ContentFilteringPlugin {
145    pub fn name(&self) -> &str {
146        &self.name
147    }
148
149    pub fn enabled(&self) -> bool {
150        self.enabled
151    }
152
153    pub async fn before_request(&self, request: &mut Request<Body>) -> Result<(), GatewayError> {
154        if let Some(content) = self.extract_content_from_request(request) {
155            self.check_content(&content)?;
156        }
157        Ok(())
158    }
159
160    pub async fn after_response(&self, _response: &mut Response<Body>) -> Result<(), GatewayError> {
161        // Could check response content here
162        Ok(())
163    }
164
165    pub async fn on_error(&self, _error: &GatewayError) -> Result<(), GatewayError> {
166        // Could log filtering events here
167        Ok(())
168    }
169
170    // Handler-level content filtering methods
171    pub fn filter_chat_request(&self, request: &ChatRequest) -> Result<(), GatewayError> {
172        // Check system message
173        if let Some(system_msg) = request
174            .messages
175            .iter()
176            .find(|m| m.role == ultrafast_models_sdk::models::Role::System)
177        {
178            self.check_content(&system_msg.content)?;
179        }
180
181        // Check user messages
182        for message in &request.messages {
183            if message.role == ultrafast_models_sdk::models::Role::User {
184                self.check_content(&message.content)?;
185            }
186        }
187
188        Ok(())
189    }
190
191    pub fn filter_embedding_request(&self, request: &EmbeddingRequest) -> Result<(), GatewayError> {
192        match &request.input {
193            ultrafast_models_sdk::models::EmbeddingInput::String(text) => {
194                self.check_content(text)?;
195            }
196            ultrafast_models_sdk::models::EmbeddingInput::StringArray(texts) => {
197                for text in texts {
198                    self.check_content(text)?;
199                }
200            }
201            ultrafast_models_sdk::models::EmbeddingInput::TokenArray(_) => {
202                // Token arrays don't contain readable text to filter
203            }
204            ultrafast_models_sdk::models::EmbeddingInput::TokenArrayArray(_) => {
205                // Token arrays don't contain readable text to filter
206            }
207        }
208        Ok(())
209    }
210
211    pub fn filter_image_request(&self, request: &ImageRequest) -> Result<(), GatewayError> {
212        self.check_content(&request.prompt)?;
213
214        // Note: ImageRequest doesn't have negative_prompt field in the current model
215        // If needed, this can be added to the ImageRequest struct in models.rs
216
217        Ok(())
218    }
219}
ultrafast_gateway/plugins/content_filtering.rs

ultrafast_gateway/plugins/
content_filtering.rs
