mockforge_http/
ai_handler.rs

1//! AI-powered response handler for HTTP requests
2//!
3//! This module integrates intelligent mock generation and data drift simulation
4//! into the HTTP request handling pipeline.
5
6use mockforge_core::Result;
7use mockforge_data::{
8    DataDriftConfig, DataDriftEngine, IntelligentMockConfig, IntelligentMockGenerator,
9};
10use serde_json::Value;
11use std::sync::Arc;
12use tokio::sync::RwLock;
13use tracing::{debug, warn};
14
15/// AI response handler that combines intelligent generation and drift simulation
16pub struct AiResponseHandler {
17    /// Intelligent mock generator (if configured)
18    intelligent_generator: Option<IntelligentMockGenerator>,
19    /// Data drift engine (if configured)
20    drift_engine: Option<Arc<RwLock<DataDriftEngine>>>,
21}
22
23impl AiResponseHandler {
24    /// Create a new AI response handler
25    pub fn new(
26        intelligent_config: Option<IntelligentMockConfig>,
27        drift_config: Option<DataDriftConfig>,
28    ) -> Result<Self> {
29        debug!("Creating AI response handler");
30
31        // Initialize intelligent generator if configured
32        let intelligent_generator = if let Some(config) = intelligent_config {
33            debug!("Initializing intelligent mock generator with mode: {:?}", config.mode);
34            Some(IntelligentMockGenerator::new(config).map_err(|e| {
35                mockforge_core::Error::Config {
36                    message: format!("Failed to initialize intelligent generator: {}", e),
37                }
38            })?)
39        } else {
40            None
41        };
42
43        // Initialize drift engine if configured
44        let drift_engine = if let Some(config) = drift_config {
45            debug!("Initializing data drift engine");
46            let engine =
47                DataDriftEngine::new(config).map_err(|e| mockforge_core::Error::Config {
48                    message: format!("Failed to initialize drift engine: {}", e),
49                })?;
50            Some(Arc::new(RwLock::new(engine)))
51        } else {
52            None
53        };
54
55        Ok(Self {
56            intelligent_generator,
57            drift_engine,
58        })
59    }
60
61    /// Check if this handler has any AI features enabled
62    pub fn is_enabled(&self) -> bool {
63        self.intelligent_generator.is_some() || self.drift_engine.is_some()
64    }
65
66    /// Generate a response using configured AI features
67    pub async fn generate_response(&mut self, base_response: Option<Value>) -> Result<Value> {
68        debug!("Generating AI-powered response");
69
70        // Step 1: Generate or use base response
71        let mut response = if let Some(generator) = &mut self.intelligent_generator {
72            match generator.generate().await {
73                Ok(resp) => {
74                    debug!("Intelligent generation successful");
75                    resp
76                }
77                Err(e) => {
78                    warn!("Intelligent generation failed: {}, using fallback", e);
79                    base_response.unwrap_or_else(|| serde_json::json!({}))
80                }
81            }
82        } else if let Some(base) = base_response {
83            base
84        } else {
85            serde_json::json!({})
86        };
87
88        // Step 2: Apply drift if configured
89        if let Some(drift_engine) = &self.drift_engine {
90            match drift_engine.read().await.apply_drift(response.clone()).await {
91                Ok(drifted) => {
92                    debug!("Data drift applied successfully");
93                    response = drifted;
94                }
95                Err(e) => {
96                    warn!("Data drift failed: {}, using non-drifted response", e);
97                }
98            }
99        }
100
101        Ok(response)
102    }
103
104    /// Reset drift state (useful for testing or specific scenarios)
105    pub async fn reset_drift(&self) {
106        if let Some(drift_engine) = &self.drift_engine {
107            drift_engine.read().await.reset().await;
108            debug!("Drift state reset");
109        }
110    }
111
112    /// Get drift request count
113    pub async fn drift_request_count(&self) -> u64 {
114        if let Some(drift_engine) = &self.drift_engine {
115            drift_engine.read().await.request_count().await
116        } else {
117            0
118        }
119    }
120}
121
122/// Helper function to create an AI handler from optional configs
123pub fn create_ai_handler(
124    intelligent_config: Option<IntelligentMockConfig>,
125    drift_config: Option<DataDriftConfig>,
126) -> Result<Option<AiResponseHandler>> {
127    if intelligent_config.is_some() || drift_config.is_some() {
128        Ok(Some(AiResponseHandler::new(intelligent_config, drift_config)?))
129    } else {
130        Ok(None)
131    }
132}
133
134/// Process a response body with AI features if configured
135///
136/// This is a helper function that checks if a response has AI configuration
137/// and applies intelligent generation or drift if present.
138///
139/// # Arguments
140/// * `response_body` - The base response body (as JSON string or Value)
141/// * `intelligent_config` - Optional intelligent mock configuration (from MockResponse.intelligent)
142/// * `drift_config` - Optional drift configuration (from MockResponse.drift)
143///
144/// # Returns
145/// The processed response body as a JSON Value
146pub async fn process_response_with_ai(
147    response_body: Option<Value>,
148    intelligent_config: Option<Value>,
149    drift_config: Option<Value>,
150) -> Result<Value> {
151    // Parse configs if present
152    let intelligent: Option<IntelligentMockConfig> =
153        intelligent_config.and_then(|v| serde_json::from_value(v).ok());
154
155    let drift: Option<DataDriftConfig> = drift_config.and_then(|v| serde_json::from_value(v).ok());
156
157    // If no AI config, return original response
158    if intelligent.is_none() && drift.is_none() {
159        return Ok(response_body.unwrap_or_else(|| serde_json::json!({})));
160    }
161
162    // Create AI handler and generate response
163    let mut handler = AiResponseHandler::new(intelligent, drift)?;
164    handler.generate_response(response_body).await
165}
166
167/// Configuration for AI-powered responses (to be added to MockResponse)
168#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
169pub struct AiResponseConfig {
170    /// Intelligent mock configuration
171    #[serde(skip_serializing_if = "Option::is_none")]
172    pub intelligent: Option<IntelligentMockConfig>,
173
174    /// Data drift configuration
175    #[serde(skip_serializing_if = "Option::is_none")]
176    pub drift: Option<DataDriftConfig>,
177}
178
179impl AiResponseConfig {
180    /// Check if any AI features are configured
181    pub fn is_enabled(&self) -> bool {
182        self.intelligent.is_some() || self.drift.is_some()
183    }
184
185    /// Create an AI handler from this configuration
186    pub fn create_handler(&self) -> Result<Option<AiResponseHandler>> {
187        create_ai_handler(self.intelligent.clone(), self.drift.clone())
188    }
189}
190
191#[cfg(test)]
192mod tests {
193    use super::*;
194    use mockforge_data::drift::{DriftRule, DriftStrategy};
195    use mockforge_data::ResponseMode;
196
197    #[test]
198    fn test_ai_handler_creation_intelligent_only() {
199        let config = IntelligentMockConfig::new(ResponseMode::Intelligent)
200            .with_prompt("Test prompt".to_string());
201
202        let result = AiResponseHandler::new(Some(config), None);
203        assert!(result.is_ok());
204
205        let handler = result.unwrap();
206        assert!(handler.is_enabled());
207        assert!(handler.intelligent_generator.is_some());
208        assert!(handler.drift_engine.is_none());
209    }
210
211    #[test]
212    fn test_ai_handler_creation_drift_only() {
213        let rule = DriftRule::new("field".to_string(), DriftStrategy::Linear).with_rate(1.0);
214        let drift_config = DataDriftConfig::new().with_rule(rule);
215
216        let result = AiResponseHandler::new(None, Some(drift_config));
217        assert!(result.is_ok());
218
219        let handler = result.unwrap();
220        assert!(handler.is_enabled());
221        assert!(handler.intelligent_generator.is_none());
222        assert!(handler.drift_engine.is_some());
223    }
224
225    #[test]
226    fn test_ai_handler_creation_both() {
227        let intelligent_config =
228            IntelligentMockConfig::new(ResponseMode::Intelligent).with_prompt("Test".to_string());
229        let rule = DriftRule::new("field".to_string(), DriftStrategy::Linear);
230        let drift_config = DataDriftConfig::new().with_rule(rule);
231
232        let result = AiResponseHandler::new(Some(intelligent_config), Some(drift_config));
233        assert!(result.is_ok());
234
235        let handler = result.unwrap();
236        assert!(handler.is_enabled());
237        assert!(handler.intelligent_generator.is_some());
238        assert!(handler.drift_engine.is_some());
239    }
240
241    #[test]
242    fn test_ai_handler_creation_neither() {
243        let result = AiResponseHandler::new(None, None);
244        assert!(result.is_ok());
245
246        let handler = result.unwrap();
247        assert!(!handler.is_enabled());
248    }
249
250    #[test]
251    fn test_ai_response_config_is_enabled() {
252        let config = AiResponseConfig {
253            intelligent: Some(IntelligentMockConfig::new(ResponseMode::Intelligent)),
254            drift: None,
255        };
256        assert!(config.is_enabled());
257
258        let config = AiResponseConfig {
259            intelligent: None,
260            drift: None,
261        };
262        assert!(!config.is_enabled());
263    }
264
265    #[tokio::test]
266    async fn test_generate_response_with_base() {
267        let mut handler = AiResponseHandler::new(None, None).unwrap();
268        let base = serde_json::json!({"test": "value"});
269
270        let result = handler.generate_response(Some(base.clone())).await;
271        assert!(result.is_ok());
272        assert_eq!(result.unwrap(), base);
273    }
274}
mockforge_http/ai_handler.rs

mockforge_http/
ai_handler.rs
