use mockforge_core::Result;
use mockforge_data::{
DataDriftConfig, DataDriftEngine, IntelligentMockConfig, IntelligentMockGenerator,
};
use serde_json::Value;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, warn};
pub struct AiResponseHandler {
intelligent_generator: Option<IntelligentMockGenerator>,
drift_engine: Option<Arc<RwLock<DataDriftEngine>>>,
}
impl AiResponseHandler {
pub fn new(
intelligent_config: Option<IntelligentMockConfig>,
drift_config: Option<DataDriftConfig>,
) -> Result<Self> {
debug!("Creating AI response handler");
let intelligent_generator = if let Some(config) = intelligent_config {
debug!("Initializing intelligent mock generator with mode: {:?}", config.mode);
Some(IntelligentMockGenerator::new(config).map_err(|e| {
mockforge_core::Error::Config {
message: format!("Failed to initialize intelligent generator: {}", e),
}
})?)
} else {
None
};
let drift_engine = if let Some(config) = drift_config {
debug!("Initializing data drift engine");
let engine =
DataDriftEngine::new(config).map_err(|e| mockforge_core::Error::Config {
message: format!("Failed to initialize drift engine: {}", e),
})?;
Some(Arc::new(RwLock::new(engine)))
} else {
None
};
Ok(Self {
intelligent_generator,
drift_engine,
})
}
pub fn is_enabled(&self) -> bool {
self.intelligent_generator.is_some() || self.drift_engine.is_some()
}
pub async fn generate_response(&mut self, base_response: Option<Value>) -> Result<Value> {
debug!("Generating AI-powered response");
let mut response = if let Some(generator) = &mut self.intelligent_generator {
match generator.generate().await {
Ok(resp) => {
debug!("Intelligent generation successful");
resp
}
Err(e) => {
warn!("Intelligent generation failed: {}, using fallback", e);
base_response.unwrap_or_else(|| serde_json::json!({}))
}
}
} else if let Some(base) = base_response {
base
} else {
serde_json::json!({})
};
if let Some(drift_engine) = &self.drift_engine {
match drift_engine.read().await.apply_drift(response.clone()).await {
Ok(drifted) => {
debug!("Data drift applied successfully");
response = drifted;
}
Err(e) => {
warn!("Data drift failed: {}, using non-drifted response", e);
}
}
}
Ok(response)
}
pub async fn reset_drift(&self) {
if let Some(drift_engine) = &self.drift_engine {
drift_engine.read().await.reset().await;
debug!("Drift state reset");
}
}
pub async fn drift_request_count(&self) -> u64 {
if let Some(drift_engine) = &self.drift_engine {
drift_engine.read().await.request_count().await
} else {
0
}
}
}
pub fn create_ai_handler(
intelligent_config: Option<IntelligentMockConfig>,
drift_config: Option<DataDriftConfig>,
) -> Result<Option<AiResponseHandler>> {
if intelligent_config.is_some() || drift_config.is_some() {
Ok(Some(AiResponseHandler::new(intelligent_config, drift_config)?))
} else {
Ok(None)
}
}
pub async fn process_response_with_ai(
response_body: Option<Value>,
intelligent_config: Option<Value>,
drift_config: Option<Value>,
) -> Result<Value> {
let intelligent: Option<IntelligentMockConfig> =
intelligent_config.and_then(|v| serde_json::from_value(v).ok());
let drift: Option<DataDriftConfig> = drift_config.and_then(|v| serde_json::from_value(v).ok());
if intelligent.is_none() && drift.is_none() {
return Ok(response_body.unwrap_or_else(|| serde_json::json!({})));
}
let mut handler = AiResponseHandler::new(intelligent, drift)?;
handler.generate_response(response_body).await
}
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
pub struct AiResponseConfig {
#[serde(skip_serializing_if = "Option::is_none")]
pub intelligent: Option<IntelligentMockConfig>,
#[serde(skip_serializing_if = "Option::is_none")]
pub drift: Option<DataDriftConfig>,
}
impl AiResponseConfig {
pub fn is_enabled(&self) -> bool {
self.intelligent.is_some() || self.drift.is_some()
}
pub fn create_handler(&self) -> Result<Option<AiResponseHandler>> {
create_ai_handler(self.intelligent.clone(), self.drift.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
use mockforge_data::drift::{DriftRule, DriftStrategy};
use mockforge_data::ResponseMode;
#[test]
fn test_ai_handler_creation_intelligent_only() {
let config = IntelligentMockConfig::new(ResponseMode::Intelligent)
.with_prompt("Test prompt".to_string());
let result = AiResponseHandler::new(Some(config), None);
assert!(result.is_ok());
let handler = result.unwrap();
assert!(handler.is_enabled());
assert!(handler.intelligent_generator.is_some());
assert!(handler.drift_engine.is_none());
}
#[test]
fn test_ai_handler_creation_drift_only() {
let rule = DriftRule::new("field".to_string(), DriftStrategy::Linear).with_rate(1.0);
let drift_config = DataDriftConfig::new().with_rule(rule);
let result = AiResponseHandler::new(None, Some(drift_config));
assert!(result.is_ok());
let handler = result.unwrap();
assert!(handler.is_enabled());
assert!(handler.intelligent_generator.is_none());
assert!(handler.drift_engine.is_some());
}
#[test]
fn test_ai_handler_creation_both() {
let intelligent_config =
IntelligentMockConfig::new(ResponseMode::Intelligent).with_prompt("Test".to_string());
let rule = DriftRule::new("field".to_string(), DriftStrategy::Linear);
let drift_config = DataDriftConfig::new().with_rule(rule);
let result = AiResponseHandler::new(Some(intelligent_config), Some(drift_config));
assert!(result.is_ok());
let handler = result.unwrap();
assert!(handler.is_enabled());
assert!(handler.intelligent_generator.is_some());
assert!(handler.drift_engine.is_some());
}
#[test]
fn test_ai_handler_creation_neither() {
let result = AiResponseHandler::new(None, None);
assert!(result.is_ok());
let handler = result.unwrap();
assert!(!handler.is_enabled());
}
#[test]
fn test_ai_response_config_is_enabled() {
let config = AiResponseConfig {
intelligent: Some(IntelligentMockConfig::new(ResponseMode::Intelligent)),
drift: None,
};
assert!(config.is_enabled());
let config = AiResponseConfig {
intelligent: None,
drift: None,
};
assert!(!config.is_enabled());
}
#[tokio::test]
async fn test_generate_response_with_base() {
let mut handler = AiResponseHandler::new(None, None).unwrap();
let base = serde_json::json!({"test": "value"});
let result = handler.generate_response(Some(base.clone())).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), base);
}
}