Skip to main content

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")]
13#[derive(Default)]
14pub enum AiResponseMode {
15    /// Static response (no AI)
16    #[default]
17    Static,
18    /// Generate response using LLM
19    Intelligent,
20    /// Use static template enhanced with LLM
21    Hybrid,
22}
23
24/// Configuration for AI-assisted response generation per endpoint
25/// This is parsed from the `x-mockforge-ai` OpenAPI extension
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct AiResponseConfig {
28    /// Whether AI response generation is enabled
29    #[serde(default)]
30    pub enabled: bool,
31
32    /// Response generation mode
33    #[serde(default)]
34    pub mode: AiResponseMode,
35
36    /// Prompt template for LLM generation
37    /// Supports template variables: {{body.field}}, {{path.param}}, {{query.param}}, {{headers.name}}
38    pub prompt: Option<String>,
39
40    /// Additional context for generation
41    pub context: Option<String>,
42
43    /// Temperature for LLM (0.0 to 2.0)
44    #[serde(default = "default_temperature")]
45    pub temperature: f32,
46
47    /// Max tokens for LLM response
48    #[serde(default = "default_max_tokens")]
49    pub max_tokens: usize,
50
51    /// Schema that the response should conform to (JSON Schema)
52    pub schema: Option<Value>,
53
54    /// Enable caching for identical requests
55    #[serde(default = "default_true")]
56    pub cache_enabled: bool,
57}
58
59fn default_temperature() -> f32 {
60    0.7
61}
62
63fn default_max_tokens() -> usize {
64    1024
65}
66
67fn default_true() -> bool {
68    true
69}
70
71impl Default for AiResponseConfig {
72    fn default() -> Self {
73        Self {
74            enabled: false,
75            mode: AiResponseMode::Static,
76            prompt: None,
77            context: None,
78            temperature: default_temperature(),
79            max_tokens: default_max_tokens(),
80            schema: None,
81            cache_enabled: true,
82        }
83    }
84}
85
86impl AiResponseConfig {
87    /// Create a new AI response configuration
88    pub fn new(enabled: bool, mode: AiResponseMode, prompt: String) -> Self {
89        Self {
90            enabled,
91            mode,
92            prompt: Some(prompt),
93            ..Default::default()
94        }
95    }
96
97    /// Check if AI generation is enabled and configured
98    pub fn is_active(&self) -> bool {
99        self.enabled && self.mode != AiResponseMode::Static && self.prompt.is_some()
100    }
101}
102
103/// Request context for prompt template expansion
104#[derive(Debug, Clone, Default)]
105pub struct RequestContext {
106    /// HTTP method (GET, POST, etc.)
107    pub method: String,
108    /// Request path
109    pub path: String,
110    /// Path parameters
111    pub path_params: HashMap<String, Value>,
112    /// Query parameters
113    pub query_params: HashMap<String, Value>,
114    /// Request headers
115    pub headers: HashMap<String, Value>,
116    /// Request body (if JSON)
117    pub body: Option<Value>,
118    /// Multipart form fields (for multipart/form-data requests)
119    pub multipart_fields: HashMap<String, Value>,
120    /// Multipart file uploads (filename -> file path)
121    pub multipart_files: HashMap<String, String>,
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    /// Set multipart form fields
159    pub fn with_multipart_fields(mut self, fields: HashMap<String, Value>) -> Self {
160        self.multipart_fields = fields;
161        self
162    }
163
164    /// Set multipart file uploads
165    pub fn with_multipart_files(mut self, files: HashMap<String, String>) -> Self {
166        self.multipart_files = files;
167        self
168    }
169}
170
171/// Expand template variables in a prompt string using request context
172///
173/// **Note**: This function has been moved to `mockforge-template-expansion` crate.
174/// Use `mockforge_template_expansion::expand_prompt_template` for the full implementation.
175///
176/// This backward-compatible version performs basic template expansion for common patterns.
177/// It handles `{{method}}`, `{{path}}`, `{{query.*}}`, `{{path.*}}`, `{{headers.*}}`, and `{{body.*}}`.
178///
179/// For new code, import directly from the template expansion crate:
180/// ```rust
181/// use mockforge_template_expansion::expand_prompt_template;
182/// ```
183#[deprecated(note = "Use mockforge_template_expansion::expand_prompt_template instead")]
184pub fn expand_prompt_template(template: &str, context: &RequestContext) -> String {
185    let mut result = template.to_string();
186
187    // Normalize {{request.*}} prefix to standard format
188    result = result
189        .replace("{{request.query.", "{{query.")
190        .replace("{{request.path.", "{{path.")
191        .replace("{{request.headers.", "{{headers.")
192        .replace("{{request.body.", "{{body.")
193        .replace("{{request.method}}", "{{method}}")
194        .replace("{{request.path}}", "{{path}}");
195
196    // Replace {{method}}
197    result = result.replace("{{method}}", &context.method);
198
199    // Replace {{path}}
200    result = result.replace("{{path}}", &context.path);
201
202    // Replace {{path.*}} variables
203    for (key, val) in &context.path_params {
204        let placeholder = format!("{{{{path.{key}}}}}");
205        let replacement = value_to_string(val);
206        result = result.replace(&placeholder, &replacement);
207    }
208
209    // Replace {{query.*}} variables
210    for (key, val) in &context.query_params {
211        let placeholder = format!("{{{{query.{key}}}}}");
212        let replacement = value_to_string(val);
213        result = result.replace(&placeholder, &replacement);
214    }
215
216    // Replace {{headers.*}} variables
217    for (key, val) in &context.headers {
218        let placeholder = format!("{{{{headers.{key}}}}}");
219        let replacement = value_to_string(val);
220        result = result.replace(&placeholder, &replacement);
221    }
222
223    // Replace {{body.*}} variables
224    if let Some(body) = &context.body {
225        if let Some(obj) = body.as_object() {
226            for (key, val) in obj {
227                let placeholder = format!("{{{{body.{key}}}}}");
228                let replacement = value_to_string(val);
229                result = result.replace(&placeholder, &replacement);
230            }
231        }
232    }
233
234    // Replace {{multipart.*}} variables
235    for (key, val) in &context.multipart_fields {
236        let placeholder = format!("{{{{multipart.{key}}}}}");
237        let replacement = value_to_string(val);
238        result = result.replace(&placeholder, &replacement);
239    }
240
241    result
242}
243
244/// Helper to convert a JSON value to a string representation
245fn value_to_string(val: &Value) -> String {
246    match val {
247        Value::String(s) => s.clone(),
248        Value::Number(n) => n.to_string(),
249        Value::Bool(b) => b.to_string(),
250        Value::Null => "null".to_string(),
251        _ => serde_json::to_string(val).unwrap_or_default(),
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258    use mockforge_template_expansion::{
259        expand_prompt_template, RequestContext as TemplateRequestContext,
260    };
261    use serde_json::json;
262
263    // Helper to convert core RequestContext to template expansion RequestContext
264    fn to_template_context(context: &RequestContext) -> TemplateRequestContext {
265        TemplateRequestContext {
266            method: context.method.clone(),
267            path: context.path.clone(),
268            path_params: context.path_params.clone(),
269            query_params: context.query_params.clone(),
270            headers: context.headers.clone(),
271            body: context.body.clone(),
272            multipart_fields: context.multipart_fields.clone(),
273            multipart_files: context.multipart_files.clone(),
274        }
275    }
276
277    #[test]
278    fn test_ai_response_config_default() {
279        let config = AiResponseConfig::default();
280        assert!(!config.enabled);
281        assert_eq!(config.mode, AiResponseMode::Static);
282        assert!(!config.is_active());
283    }
284
285    #[test]
286    fn test_ai_response_config_is_active() {
287        let config =
288            AiResponseConfig::new(true, AiResponseMode::Intelligent, "Test prompt".to_string());
289        assert!(config.is_active());
290
291        let config_disabled = AiResponseConfig {
292            enabled: false,
293            mode: AiResponseMode::Intelligent,
294            prompt: Some("Test".to_string()),
295            ..Default::default()
296        };
297        assert!(!config_disabled.is_active());
298    }
299
300    #[test]
301    fn test_request_context_builder() {
302        let mut path_params = HashMap::new();
303        path_params.insert("id".to_string(), json!("123"));
304
305        let context = RequestContext::new("POST".to_string(), "/users/123".to_string())
306            .with_path_params(path_params)
307            .with_body(json!({"name": "John"}));
308
309        assert_eq!(context.method, "POST");
310        assert_eq!(context.path, "/users/123");
311        assert_eq!(context.path_params.get("id"), Some(&json!("123")));
312        assert_eq!(context.body, Some(json!({"name": "John"})));
313    }
314
315    #[test]
316    fn test_expand_prompt_template_basic() {
317        let context = RequestContext::new("GET".to_string(), "/users".to_string());
318        let template = "Method: {{method}}, Path: {{path}}";
319        let template_context = to_template_context(&context);
320        let expanded = expand_prompt_template(template, &template_context);
321        assert_eq!(expanded, "Method: GET, Path: /users");
322    }
323
324    #[test]
325    fn test_expand_prompt_template_body() {
326        let body = json!({
327            "message": "Hello",
328            "user": "Alice"
329        });
330        let context = RequestContext::new("POST".to_string(), "/chat".to_string()).with_body(body);
331
332        let template = "User {{body.user}} says: {{body.message}}";
333        let template_context = to_template_context(&context);
334        let expanded = expand_prompt_template(template, &template_context);
335        assert_eq!(expanded, "User Alice says: Hello");
336    }
337
338    #[test]
339    fn test_expand_prompt_template_path_params() {
340        let mut path_params = HashMap::new();
341        path_params.insert("id".to_string(), json!("456"));
342        path_params.insert("name".to_string(), json!("test"));
343
344        let context = RequestContext::new("GET".to_string(), "/users/456".to_string())
345            .with_path_params(path_params);
346
347        let template = "Get user {{path.id}} with name {{path.name}}";
348        let template_context = to_template_context(&context);
349        let expanded = expand_prompt_template(template, &template_context);
350        assert_eq!(expanded, "Get user 456 with name test");
351    }
352
353    #[test]
354    fn test_expand_prompt_template_query_params() {
355        let mut query_params = HashMap::new();
356        query_params.insert("search".to_string(), json!("term"));
357        query_params.insert("limit".to_string(), json!(10));
358
359        let context = RequestContext::new("GET".to_string(), "/search".to_string())
360            .with_query_params(query_params);
361
362        let template = "Search for {{query.search}} with limit {{query.limit}}";
363        let template_context = to_template_context(&context);
364        let expanded = expand_prompt_template(template, &template_context);
365        assert_eq!(expanded, "Search for term with limit 10");
366    }
367
368    #[test]
369    fn test_expand_prompt_template_headers() {
370        let mut headers = HashMap::new();
371        headers.insert("user-agent".to_string(), json!("TestClient/1.0"));
372
373        let context =
374            RequestContext::new("GET".to_string(), "/api".to_string()).with_headers(headers);
375
376        let template = "Request from {{headers.user-agent}}";
377        let template_context = to_template_context(&context);
378        let expanded = expand_prompt_template(template, &template_context);
379        assert_eq!(expanded, "Request from TestClient/1.0");
380    }
381
382    #[test]
383    fn test_expand_prompt_template_complex() {
384        let mut path_params = HashMap::new();
385        path_params.insert("id".to_string(), json!("789"));
386
387        let mut query_params = HashMap::new();
388        query_params.insert("format".to_string(), json!("json"));
389
390        let body = json!({"action": "update", "value": 42});
391
392        let context = RequestContext::new("PUT".to_string(), "/api/items/789".to_string())
393            .with_path_params(path_params)
394            .with_query_params(query_params)
395            .with_body(body);
396
397        let template = "{{method}} item {{path.id}} with action {{body.action}} and value {{body.value}} in format {{query.format}}";
398        let template_context = to_template_context(&context);
399        let expanded = expand_prompt_template(template, &template_context);
400        assert_eq!(expanded, "PUT item 789 with action update and value 42 in format json");
401    }
402}