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

mockforge_http/
ai_handler.rs
