1use crate::{
7 ai_response::{expand_prompt_template, AiResponseConfig, RequestContext},
8 OpenApiSpec, Result,
9};
10use async_trait::async_trait;
11use chrono;
12use openapiv3::{Operation, ReferenceOr, Response, Responses, Schema};
13use rand::{rng, Rng};
14use serde_json::Value;
15use std::collections::HashMap;
16use uuid;
17
18#[async_trait]
23pub trait AiGenerator: Send + Sync {
24 async fn generate(&self, prompt: &str, config: &AiResponseConfig) -> Result<Value>;
33}
34
35pub struct ResponseGenerator;
37
38impl ResponseGenerator {
39 pub async fn generate_ai_response(
52 ai_config: &AiResponseConfig,
53 context: &RequestContext,
54 generator: Option<&dyn AiGenerator>,
55 ) -> Result<Value> {
56 let prompt_template = ai_config
58 .prompt
59 .as_ref()
60 .ok_or_else(|| crate::Error::generic("AI prompt is required"))?;
61
62 let expanded_prompt = expand_prompt_template(prompt_template, context);
63
64 tracing::info!("AI response generation requested with prompt: {}", expanded_prompt);
65
66 if let Some(gen) = generator {
68 tracing::debug!("Using provided AI generator for response");
69 return gen.generate(&expanded_prompt, ai_config).await;
70 }
71
72 tracing::warn!("No AI generator provided, returning placeholder response");
74 Ok(serde_json::json!({
75 "ai_response": "AI generation placeholder",
76 "note": "This endpoint is configured for AI-assisted responses, but no AI generator was provided",
77 "expanded_prompt": expanded_prompt,
78 "mode": format!("{:?}", ai_config.mode),
79 "temperature": ai_config.temperature,
80 "implementation_note": "Pass an AiGenerator implementation to ResponseGenerator::generate_ai_response to enable actual AI generation"
81 }))
82 }
83
84 pub fn generate_response(
86 spec: &OpenApiSpec,
87 operation: &Operation,
88 status_code: u16,
89 content_type: Option<&str>,
90 ) -> Result<Value> {
91 Self::generate_response_with_expansion(spec, operation, status_code, content_type, true)
92 }
93
94 pub fn generate_response_with_expansion(
96 spec: &OpenApiSpec,
97 operation: &Operation,
98 status_code: u16,
99 content_type: Option<&str>,
100 expand_tokens: bool,
101 ) -> Result<Value> {
102 let response = Self::find_response_for_status(&operation.responses, status_code);
104
105 match response {
106 Some(response_ref) => {
107 match response_ref {
108 ReferenceOr::Item(response) => {
109 Self::generate_from_response(spec, response, content_type, expand_tokens)
110 }
111 ReferenceOr::Reference { reference } => {
112 if let Some(resolved_response) = spec.get_response(reference) {
114 Self::generate_from_response(
115 spec,
116 resolved_response,
117 content_type,
118 expand_tokens,
119 )
120 } else {
121 Ok(Value::Object(serde_json::Map::new()))
123 }
124 }
125 }
126 }
127 None => {
128 Ok(Value::Object(serde_json::Map::new()))
130 }
131 }
132 }
133
134 fn find_response_for_status(
136 responses: &Responses,
137 status_code: u16,
138 ) -> Option<&ReferenceOr<Response>> {
139 if let Some(response) = responses.responses.get(&openapiv3::StatusCode::Code(status_code)) {
141 return Some(response);
142 }
143
144 if let Some(default_response) = &responses.default {
146 return Some(default_response);
147 }
148
149 None
150 }
151
152 fn generate_from_response(
154 spec: &OpenApiSpec,
155 response: &Response,
156 content_type: Option<&str>,
157 expand_tokens: bool,
158 ) -> Result<Value> {
159 if let Some(content_type) = content_type {
161 if let Some(media_type) = response.content.get(content_type) {
162 return Self::generate_from_media_type(spec, media_type, expand_tokens);
163 }
164 }
165
166 let preferred_types = ["application/json", "application/xml", "text/plain"];
168
169 for content_type in &preferred_types {
170 if let Some(media_type) = response.content.get(*content_type) {
171 return Self::generate_from_media_type(spec, media_type, expand_tokens);
172 }
173 }
174
175 if let Some((_, media_type)) = response.content.iter().next() {
177 return Self::generate_from_media_type(spec, media_type, expand_tokens);
178 }
179
180 Ok(Value::Object(serde_json::Map::new()))
182 }
183
184 fn generate_from_media_type(
186 spec: &OpenApiSpec,
187 media_type: &openapiv3::MediaType,
188 expand_tokens: bool,
189 ) -> Result<Value> {
190 if let Some(example) = &media_type.example {
192 tracing::debug!("Using explicit example from media type: {:?}", example);
193 if expand_tokens {
195 let expanded_example = Self::expand_templates(example);
196 return Ok(expanded_example);
197 } else {
198 return Ok(example.clone());
199 }
200 }
201
202 if !media_type.examples.is_empty() {
204 if let Some((_, example_ref)) = media_type.examples.iter().next() {
205 match example_ref {
206 ReferenceOr::Item(example) => {
207 if let Some(value) = &example.value {
208 tracing::debug!("Using example from examples map: {:?}", value);
209 if expand_tokens {
210 return Ok(Self::expand_templates(value));
211 } else {
212 return Ok(value.clone());
213 }
214 }
215 }
216 ReferenceOr::Reference { reference } => {
217 if let Some(example) = spec.get_example(reference) {
219 if let Some(value) = &example.value {
220 tracing::debug!("Using resolved example reference: {:?}", value);
221 if expand_tokens {
222 return Ok(Self::expand_templates(value));
223 } else {
224 return Ok(value.clone());
225 }
226 }
227 } else {
228 tracing::warn!("Example reference '{}' not found", reference);
229 }
230 }
231 }
232 }
233 }
234
235 match &media_type.schema {
237 Some(schema_ref) => {
238 match schema_ref {
239 ReferenceOr::Item(schema) => Ok(Self::generate_example_from_schema(schema)),
240 ReferenceOr::Reference { reference } => {
241 if let Some(schema) = spec.get_schema(reference) {
243 Ok(Self::generate_example_from_schema(&schema.schema))
244 } else {
245 Ok(Value::Object(serde_json::Map::new()))
247 }
248 }
249 }
250 }
251 None => {
252 Ok(Value::Object(serde_json::Map::new()))
254 }
255 }
256 }
257
258 fn generate_example_from_schema(schema: &Schema) -> Value {
260 match &schema.schema_kind {
261 openapiv3::SchemaKind::Type(openapiv3::Type::String(_)) => {
262 Value::String("example string".to_string())
264 }
265 openapiv3::SchemaKind::Type(openapiv3::Type::Integer(_)) => Value::Number(42.into()),
266 openapiv3::SchemaKind::Type(openapiv3::Type::Number(_)) => {
267 Value::Number(serde_json::Number::from_f64(std::f64::consts::PI).unwrap())
268 }
269 openapiv3::SchemaKind::Type(openapiv3::Type::Boolean(_)) => Value::Bool(true),
270 openapiv3::SchemaKind::Type(openapiv3::Type::Object(obj)) => {
271 let mut map = serde_json::Map::new();
272 for (prop_name, _) in &obj.properties {
273 let value = Self::generate_example_for_property(prop_name);
274 map.insert(prop_name.clone(), value);
275 }
276 Value::Object(map)
277 }
278 openapiv3::SchemaKind::Type(openapiv3::Type::Array(arr)) => match &arr.items {
279 Some(ReferenceOr::Item(item_schema)) => {
280 let example_item = Self::generate_example_from_schema(item_schema);
281 Value::Array(vec![example_item])
282 }
283 _ => Value::Array(vec![Value::String("item".to_string())]),
284 },
285 _ => Value::Object(serde_json::Map::new()),
286 }
287 }
288
289 fn generate_example_for_property(prop_name: &str) -> Value {
291 let prop_lower = prop_name.to_lowercase();
292
293 if prop_lower.contains("id") || prop_lower.contains("uuid") {
295 Value::String(uuid::Uuid::new_v4().to_string())
296 } else if prop_lower.contains("email") {
297 Value::String(format!("user{}@example.com", rng().random_range(1000..=9999)))
298 } else if prop_lower.contains("name") || prop_lower.contains("title") {
299 let names = ["John Doe", "Jane Smith", "Bob Johnson", "Alice Brown"];
300 Value::String(names[rng().random_range(0..names.len())].to_string())
301 } else if prop_lower.contains("phone") || prop_lower.contains("mobile") {
302 Value::String(format!("+1-555-{:04}", rng().random_range(1000..=9999)))
303 } else if prop_lower.contains("address") || prop_lower.contains("street") {
304 let streets = ["123 Main St", "456 Oak Ave", "789 Pine Rd", "321 Elm St"];
305 Value::String(streets[rng().random_range(0..streets.len())].to_string())
306 } else if prop_lower.contains("city") {
307 let cities = ["New York", "London", "Tokyo", "Paris", "Sydney"];
308 Value::String(cities[rng().random_range(0..cities.len())].to_string())
309 } else if prop_lower.contains("country") {
310 let countries = ["USA", "UK", "Japan", "France", "Australia"];
311 Value::String(countries[rng().random_range(0..countries.len())].to_string())
312 } else if prop_lower.contains("company") || prop_lower.contains("organization") {
313 let companies = ["Acme Corp", "Tech Solutions", "Global Inc", "Innovate Ltd"];
314 Value::String(companies[rng().random_range(0..companies.len())].to_string())
315 } else if prop_lower.contains("url") || prop_lower.contains("website") {
316 Value::String("https://example.com".to_string())
317 } else if prop_lower.contains("age") {
318 Value::Number((18 + rng().random_range(0..60)).into())
319 } else if prop_lower.contains("count") || prop_lower.contains("quantity") {
320 Value::Number((1 + rng().random_range(0..100)).into())
321 } else if prop_lower.contains("price")
322 || prop_lower.contains("amount")
323 || prop_lower.contains("cost")
324 {
325 Value::Number(
326 serde_json::Number::from_f64(
327 (rng().random::<f64>() * 1000.0 * 100.0).round() / 100.0,
328 )
329 .unwrap(),
330 )
331 } else if prop_lower.contains("active")
332 || prop_lower.contains("enabled")
333 || prop_lower.contains("is_")
334 {
335 Value::Bool(rng().random_bool(0.5))
336 } else if prop_lower.contains("date") || prop_lower.contains("time") {
337 Value::String(chrono::Utc::now().to_rfc3339())
338 } else if prop_lower.contains("description") || prop_lower.contains("comment") {
339 Value::String("This is a sample description text.".to_string())
340 } else {
341 Value::String(format!("example {}", prop_name))
342 }
343 }
344
345 pub fn generate_from_examples(
347 response: &Response,
348 content_type: Option<&str>,
349 ) -> Result<Option<Value>> {
350 use openapiv3::ReferenceOr;
351
352 if let Some(content_type) = content_type {
354 if let Some(media_type) = response.content.get(content_type) {
355 if let Some(example) = &media_type.example {
357 return Ok(Some(example.clone()));
358 }
359
360 for (_, example_ref) in &media_type.examples {
362 if let ReferenceOr::Item(example) = example_ref {
363 if let Some(value) = &example.value {
364 return Ok(Some(value.clone()));
365 }
366 }
367 }
369 }
370 }
371
372 for (_, media_type) in &response.content {
374 if let Some(example) = &media_type.example {
376 return Ok(Some(example.clone()));
377 }
378
379 for (_, example_ref) in &media_type.examples {
381 if let ReferenceOr::Item(example) = example_ref {
382 if let Some(value) = &example.value {
383 return Ok(Some(value.clone()));
384 }
385 }
386 }
388 }
389
390 Ok(None)
391 }
392
393 fn expand_templates(value: &Value) -> Value {
395 match value {
396 Value::String(s) => {
397 let expanded = s
398 .replace("{{now}}", &chrono::Utc::now().to_rfc3339())
399 .replace("{{uuid}}", &uuid::Uuid::new_v4().to_string());
400 Value::String(expanded)
401 }
402 Value::Object(map) => {
403 let mut new_map = serde_json::Map::new();
404 for (key, val) in map {
405 new_map.insert(key.clone(), Self::expand_templates(val));
406 }
407 Value::Object(new_map)
408 }
409 Value::Array(arr) => {
410 let new_arr: Vec<Value> = arr.iter().map(Self::expand_templates).collect();
411 Value::Array(new_arr)
412 }
413 _ => value.clone(),
414 }
415 }
416}
417
418#[derive(Debug, Clone)]
420pub struct MockResponse {
421 pub status_code: u16,
423 pub headers: HashMap<String, String>,
425 pub body: Option<Value>,
427}
428
429impl MockResponse {
430 pub fn new(status_code: u16) -> Self {
432 Self {
433 status_code,
434 headers: HashMap::new(),
435 body: None,
436 }
437 }
438
439 pub fn with_header(mut self, name: String, value: String) -> Self {
441 self.headers.insert(name, value);
442 self
443 }
444
445 pub fn with_body(mut self, body: Value) -> Self {
447 self.body = Some(body);
448 self
449 }
450}
451
452#[derive(Debug, Clone)]
454pub struct OpenApiSecurityRequirement {
455 pub scheme: String,
457 pub scopes: Vec<String>,
459}
460
461impl OpenApiSecurityRequirement {
462 pub fn new(scheme: String, scopes: Vec<String>) -> Self {
464 Self { scheme, scopes }
465 }
466}
467
468#[derive(Debug, Clone)]
470pub struct OpenApiOperation {
471 pub method: String,
473 pub path: String,
475 pub operation: openapiv3::Operation,
477}
478
479impl OpenApiOperation {
480 pub fn new(method: String, path: String, operation: openapiv3::Operation) -> Self {
482 Self {
483 method,
484 path,
485 operation,
486 }
487 }
488}