aimds_detection/
lib.rs

1//! AIMDS Detection Layer
2//!
3//! This crate provides pattern matching, sanitization, and scheduling
4//! for detecting potential threats in AI model inputs.
5
6pub mod pattern_matcher;
7pub mod sanitizer;
8pub mod scheduler;
9
10pub use pattern_matcher::PatternMatcher;
11pub use sanitizer::{Sanitizer, PiiMatch, PiiType};
12pub use scheduler::{DetectionScheduler, ThreatPriority};
13
14use aimds_core::{DetectionResult, PromptInput, Result};
15
16/// Main detection service that coordinates all detection components
17pub struct DetectionService {
18    pattern_matcher: PatternMatcher,
19    sanitizer: Sanitizer,
20    scheduler: DetectionScheduler,
21}
22
23impl DetectionService {
24    /// Create a new detection service
25    pub fn new() -> Result<Self> {
26        Ok(Self {
27            pattern_matcher: PatternMatcher::new()?,
28            sanitizer: Sanitizer::new(),
29            scheduler: DetectionScheduler::new()?,
30        })
31    }
32
33    /// Process a prompt input through all detection layers
34    pub async fn detect(&self, input: &PromptInput) -> Result<DetectionResult> {
35        // Schedule the detection task
36        self.scheduler.schedule_detection(input.id).await?;
37
38        // Pattern matching
39        let detection = self.pattern_matcher.match_patterns(&input.content).await?;
40
41        // Sanitization
42        let _sanitized = self.sanitizer.sanitize(&input.content).await?;
43
44        Ok(detection)
45    }
46}
47
48impl Default for DetectionService {
49    fn default() -> Self {
50        Self::new().expect("Failed to create detection service")
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[tokio::test]
59    async fn test_detection_service() {
60        let service = DetectionService::new().unwrap();
61        let input = PromptInput::new("Test prompt".to_string());
62
63        let result = service.detect(&input).await;
64        assert!(result.is_ok());
65    }
66}
aimds_detection/lib.rs

aimds_detection/
lib.rs
