mockforge_core/
ai_response.rs

1//! AI-assisted response generation for dynamic mock endpoints
2//!
3//! This module provides configuration and utilities for generating
4//! dynamic mock responses using LLMs based on request context.
5
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::collections::HashMap;
9
10/// AI response generation mode
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
12#[serde(rename_all = "lowercase")]
13pub enum AiResponseMode {
14    /// Static response (no AI)
15    Static,
16    /// Generate response using LLM
17    Intelligent,
18    /// Use static template enhanced with LLM
19    Hybrid,
20}
21
22impl Default for AiResponseMode {
23    fn default() -> Self {
24        Self::Static
25    }
26}
27
28/// Configuration for AI-assisted response generation per endpoint
29/// This is parsed from the `x-mockforge-ai` OpenAPI extension
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct AiResponseConfig {
32    /// Whether AI response generation is enabled
33    #[serde(default)]
34    pub enabled: bool,
35
36    /// Response generation mode
37    #[serde(default)]
38    pub mode: AiResponseMode,
39
40    /// Prompt template for LLM generation
41    /// Supports template variables: {{body.field}}, {{path.param}}, {{query.param}}, {{headers.name}}
42    pub prompt: Option<String>,
43
44    /// Additional context for generation
45    pub context: Option<String>,
46
47    /// Temperature for LLM (0.0 to 2.0)
48    #[serde(default = "default_temperature")]
49    pub temperature: f32,
50
51    /// Max tokens for LLM response
52    #[serde(default = "default_max_tokens")]
53    pub max_tokens: usize,
54
55    /// Schema that the response should conform to (JSON Schema)
56    pub schema: Option<Value>,
57
58    /// Enable caching for identical requests
59    #[serde(default = "default_true")]
60    pub cache_enabled: bool,
61}
62
63fn default_temperature() -> f32 {
64    0.7
65}
66
67fn default_max_tokens() -> usize {
68    1024
69}
70
71fn default_true() -> bool {
72    true
73}
74
75impl Default for AiResponseConfig {
76    fn default() -> Self {
77        Self {
78            enabled: false,
79            mode: AiResponseMode::Static,
80            prompt: None,
81            context: None,
82            temperature: default_temperature(),
83            max_tokens: default_max_tokens(),
84            schema: None,
85            cache_enabled: true,
86        }
87    }
88}
89
90impl AiResponseConfig {
91    /// Create a new AI response configuration
92    pub fn new(enabled: bool, mode: AiResponseMode, prompt: String) -> Self {
93        Self {
94            enabled,
95            mode,
96            prompt: Some(prompt),
97            ..Default::default()
98        }
99    }
100
101    /// Check if AI generation is enabled and configured
102    pub fn is_active(&self) -> bool {
103        self.enabled && self.mode != AiResponseMode::Static && self.prompt.is_some()
104    }
105}
106
107/// Request context for prompt template expansion
108#[derive(Debug, Clone, Default)]
109pub struct RequestContext {
110    /// HTTP method (GET, POST, etc.)
111    pub method: String,
112    /// Request path
113    pub path: String,
114    /// Path parameters
115    pub path_params: HashMap<String, Value>,
116    /// Query parameters
117    pub query_params: HashMap<String, Value>,
118    /// Request headers
119    pub headers: HashMap<String, Value>,
120    /// Request body (if JSON)
121    pub body: Option<Value>,
122}
123
124impl RequestContext {
125    /// Create a new request context
126    pub fn new(method: String, path: String) -> Self {
127        Self {
128            method,
129            path,
130            ..Default::default()
131        }
132    }
133
134    /// Set path parameters
135    pub fn with_path_params(mut self, params: HashMap<String, Value>) -> Self {
136        self.path_params = params;
137        self
138    }
139
140    /// Set query parameters
141    pub fn with_query_params(mut self, params: HashMap<String, Value>) -> Self {
142        self.query_params = params;
143        self
144    }
145
146    /// Set headers
147    pub fn with_headers(mut self, headers: HashMap<String, Value>) -> Self {
148        self.headers = headers;
149        self
150    }
151
152    /// Set body
153    pub fn with_body(mut self, body: Value) -> Self {
154        self.body = Some(body);
155        self
156    }
157}
158
159/// Expand template variables in a prompt string using request context
160pub fn expand_prompt_template(template: &str, context: &RequestContext) -> String {
161    let mut result = template.to_string();
162
163    // Replace {{method}}
164    result = result.replace("{{method}}", &context.method);
165
166    // Replace {{path}}
167    result = result.replace("{{path}}", &context.path);
168
169    // Replace {{body.*}} variables
170    if let Some(body) = &context.body {
171        result = expand_json_variables(&result, "body", body);
172    }
173
174    // Replace {{path.*}} variables
175    result = expand_map_variables(&result, "path", &context.path_params);
176
177    // Replace {{query.*}} variables
178    result = expand_map_variables(&result, "query", &context.query_params);
179
180    // Replace {{headers.*}} variables
181    result = expand_map_variables(&result, "headers", &context.headers);
182
183    result
184}
185
186/// Expand template variables from a JSON value
187fn expand_json_variables(template: &str, prefix: &str, value: &Value) -> String {
188    let mut result = template.to_string();
189
190    // Handle object fields
191    if let Some(obj) = value.as_object() {
192        for (key, val) in obj {
193            let placeholder = format!("{{{{{}.{}}}}}", prefix, key);
194            let replacement = match val {
195                Value::String(s) => s.clone(),
196                Value::Number(n) => n.to_string(),
197                Value::Bool(b) => b.to_string(),
198                Value::Null => "null".to_string(),
199                _ => serde_json::to_string(val).unwrap_or_default(),
200            };
201            result = result.replace(&placeholder, &replacement);
202        }
203    }
204
205    result
206}
207
208/// Expand template variables from a HashMap
209fn expand_map_variables(template: &str, prefix: &str, map: &HashMap<String, Value>) -> String {
210    let mut result = template.to_string();
211
212    for (key, val) in map {
213        let placeholder = format!("{{{{{}.{}}}}}", prefix, key);
214        let replacement = match val {
215            Value::String(s) => s.clone(),
216            Value::Number(n) => n.to_string(),
217            Value::Bool(b) => b.to_string(),
218            Value::Null => "null".to_string(),
219            _ => serde_json::to_string(val).unwrap_or_default(),
220        };
221        result = result.replace(&placeholder, &replacement);
222    }
223
224    result
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230    use serde_json::json;
231
232    #[test]
233    fn test_ai_response_config_default() {
234        let config = AiResponseConfig::default();
235        assert!(!config.enabled);
236        assert_eq!(config.mode, AiResponseMode::Static);
237        assert!(!config.is_active());
238    }
239
240    #[test]
241    fn test_ai_response_config_is_active() {
242        let config =
243            AiResponseConfig::new(true, AiResponseMode::Intelligent, "Test prompt".to_string());
244        assert!(config.is_active());
245
246        let config_disabled = AiResponseConfig {
247            enabled: false,
248            mode: AiResponseMode::Intelligent,
249            prompt: Some("Test".to_string()),
250            ..Default::default()
251        };
252        assert!(!config_disabled.is_active());
253    }
254
255    #[test]
256    fn test_request_context_builder() {
257        let mut path_params = HashMap::new();
258        path_params.insert("id".to_string(), json!("123"));
259
260        let context = RequestContext::new("POST".to_string(), "/users/123".to_string())
261            .with_path_params(path_params)
262            .with_body(json!({"name": "John"}));
263
264        assert_eq!(context.method, "POST");
265        assert_eq!(context.path, "/users/123");
266        assert_eq!(context.path_params.get("id"), Some(&json!("123")));
267        assert_eq!(context.body, Some(json!({"name": "John"})));
268    }
269
270    #[test]
271    fn test_expand_prompt_template_basic() {
272        let context = RequestContext::new("GET".to_string(), "/users".to_string());
273        let template = "Method: {{method}}, Path: {{path}}";
274        let expanded = expand_prompt_template(template, &context);
275        assert_eq!(expanded, "Method: GET, Path: /users");
276    }
277
278    #[test]
279    fn test_expand_prompt_template_body() {
280        let body = json!({
281            "message": "Hello",
282            "user": "Alice"
283        });
284        let context = RequestContext::new("POST".to_string(), "/chat".to_string()).with_body(body);
285
286        let template = "User {{body.user}} says: {{body.message}}";
287        let expanded = expand_prompt_template(template, &context);
288        assert_eq!(expanded, "User Alice says: Hello");
289    }
290
291    #[test]
292    fn test_expand_prompt_template_path_params() {
293        let mut path_params = HashMap::new();
294        path_params.insert("id".to_string(), json!("456"));
295        path_params.insert("name".to_string(), json!("test"));
296
297        let context = RequestContext::new("GET".to_string(), "/users/456".to_string())
298            .with_path_params(path_params);
299
300        let template = "Get user {{path.id}} with name {{path.name}}";
301        let expanded = expand_prompt_template(template, &context);
302        assert_eq!(expanded, "Get user 456 with name test");
303    }
304
305    #[test]
306    fn test_expand_prompt_template_query_params() {
307        let mut query_params = HashMap::new();
308        query_params.insert("search".to_string(), json!("term"));
309        query_params.insert("limit".to_string(), json!(10));
310
311        let context = RequestContext::new("GET".to_string(), "/search".to_string())
312            .with_query_params(query_params);
313
314        let template = "Search for {{query.search}} with limit {{query.limit}}";
315        let expanded = expand_prompt_template(template, &context);
316        assert_eq!(expanded, "Search for term with limit 10");
317    }
318
319    #[test]
320    fn test_expand_prompt_template_headers() {
321        let mut headers = HashMap::new();
322        headers.insert("user-agent".to_string(), json!("TestClient/1.0"));
323
324        let context =
325            RequestContext::new("GET".to_string(), "/api".to_string()).with_headers(headers);
326
327        let template = "Request from {{headers.user-agent}}";
328        let expanded = expand_prompt_template(template, &context);
329        assert_eq!(expanded, "Request from TestClient/1.0");
330    }
331
332    #[test]
333    fn test_expand_prompt_template_complex() {
334        let mut path_params = HashMap::new();
335        path_params.insert("id".to_string(), json!("789"));
336
337        let mut query_params = HashMap::new();
338        query_params.insert("format".to_string(), json!("json"));
339
340        let body = json!({"action": "update", "value": 42});
341
342        let context = RequestContext::new("PUT".to_string(), "/api/items/789".to_string())
343            .with_path_params(path_params)
344            .with_query_params(query_params)
345            .with_body(body);
346
347        let template = "{{method}} item {{path.id}} with action {{body.action}} and value {{body.value}} in format {{query.format}}";
348        let expanded = expand_prompt_template(template, &context);
349        assert_eq!(expanded, "PUT item 789 with action update and value 42 in format json");
350    }
351}