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 Self::generate_response_with_expansion_and_mode(
103 spec,
104 operation,
105 status_code,
106 content_type,
107 expand_tokens,
108 None,
109 None,
110 )
111 }
112
113 pub fn generate_response_with_expansion_and_mode(
115 spec: &OpenApiSpec,
116 operation: &Operation,
117 status_code: u16,
118 content_type: Option<&str>,
119 expand_tokens: bool,
120 selection_mode: Option<crate::openapi::response_selection::ResponseSelectionMode>,
121 selector: Option<&crate::openapi::response_selection::ResponseSelector>,
122 ) -> Result<Value> {
123 Self::generate_response_with_scenario_and_mode(
124 spec,
125 operation,
126 status_code,
127 content_type,
128 expand_tokens,
129 None, selection_mode,
131 selector,
132 )
133 }
134
135 pub fn generate_response_with_scenario(
161 spec: &OpenApiSpec,
162 operation: &Operation,
163 status_code: u16,
164 content_type: Option<&str>,
165 expand_tokens: bool,
166 scenario: Option<&str>,
167 ) -> Result<Value> {
168 Self::generate_response_with_scenario_and_mode(
169 spec,
170 operation,
171 status_code,
172 content_type,
173 expand_tokens,
174 scenario,
175 None,
176 None,
177 )
178 }
179
180 pub fn generate_response_with_scenario_and_mode(
182 spec: &OpenApiSpec,
183 operation: &Operation,
184 status_code: u16,
185 content_type: Option<&str>,
186 expand_tokens: bool,
187 scenario: Option<&str>,
188 selection_mode: Option<crate::openapi::response_selection::ResponseSelectionMode>,
189 selector: Option<&crate::openapi::response_selection::ResponseSelector>,
190 ) -> Result<Value> {
191 let response = Self::find_response_for_status(&operation.responses, status_code);
193
194 match response {
195 Some(response_ref) => {
196 match response_ref {
197 ReferenceOr::Item(response) => {
198 Self::generate_from_response_with_scenario_and_mode(
199 spec,
200 response,
201 content_type,
202 expand_tokens,
203 scenario,
204 selection_mode,
205 selector,
206 )
207 }
208 ReferenceOr::Reference { reference } => {
209 if let Some(resolved_response) = spec.get_response(reference) {
211 Self::generate_from_response_with_scenario_and_mode(
212 spec,
213 resolved_response,
214 content_type,
215 expand_tokens,
216 scenario,
217 selection_mode,
218 selector,
219 )
220 } else {
221 Ok(Value::Object(serde_json::Map::new()))
223 }
224 }
225 }
226 }
227 None => {
228 Ok(Value::Object(serde_json::Map::new()))
230 }
231 }
232 }
233
234 fn find_response_for_status(
236 responses: &Responses,
237 status_code: u16,
238 ) -> Option<&ReferenceOr<Response>> {
239 if let Some(response) = responses.responses.get(&openapiv3::StatusCode::Code(status_code)) {
241 return Some(response);
242 }
243
244 if let Some(default_response) = &responses.default {
246 return Some(default_response);
247 }
248
249 None
250 }
251
252 fn generate_from_response(
254 spec: &OpenApiSpec,
255 response: &Response,
256 content_type: Option<&str>,
257 expand_tokens: bool,
258 ) -> Result<Value> {
259 Self::generate_from_response_with_scenario(
260 spec,
261 response,
262 content_type,
263 expand_tokens,
264 None,
265 )
266 }
267
268 fn generate_from_response_with_scenario(
270 spec: &OpenApiSpec,
271 response: &Response,
272 content_type: Option<&str>,
273 expand_tokens: bool,
274 scenario: Option<&str>,
275 ) -> Result<Value> {
276 Self::generate_from_response_with_scenario_and_mode(
277 spec,
278 response,
279 content_type,
280 expand_tokens,
281 scenario,
282 None,
283 None,
284 )
285 }
286
287 fn generate_from_response_with_scenario_and_mode(
289 spec: &OpenApiSpec,
290 response: &Response,
291 content_type: Option<&str>,
292 expand_tokens: bool,
293 scenario: Option<&str>,
294 selection_mode: Option<crate::openapi::response_selection::ResponseSelectionMode>,
295 selector: Option<&crate::openapi::response_selection::ResponseSelector>,
296 ) -> Result<Value> {
297 if let Some(content_type) = content_type {
299 if let Some(media_type) = response.content.get(content_type) {
300 return Self::generate_from_media_type_with_scenario_and_mode(
301 spec,
302 media_type,
303 expand_tokens,
304 scenario,
305 selection_mode,
306 selector,
307 );
308 }
309 }
310
311 let preferred_types = ["application/json", "application/xml", "text/plain"];
313
314 for content_type in &preferred_types {
315 if let Some(media_type) = response.content.get(*content_type) {
316 return Self::generate_from_media_type_with_scenario_and_mode(
317 spec,
318 media_type,
319 expand_tokens,
320 scenario,
321 selection_mode,
322 selector,
323 );
324 }
325 }
326
327 if let Some((_, media_type)) = response.content.iter().next() {
329 return Self::generate_from_media_type_with_scenario_and_mode(
330 spec,
331 media_type,
332 expand_tokens,
333 scenario,
334 selection_mode,
335 selector,
336 );
337 }
338
339 Ok(Value::Object(serde_json::Map::new()))
341 }
342
343 fn generate_from_media_type(
345 spec: &OpenApiSpec,
346 media_type: &openapiv3::MediaType,
347 expand_tokens: bool,
348 ) -> Result<Value> {
349 Self::generate_from_media_type_with_scenario(spec, media_type, expand_tokens, None)
350 }
351
352 fn generate_from_media_type_with_scenario(
354 spec: &OpenApiSpec,
355 media_type: &openapiv3::MediaType,
356 expand_tokens: bool,
357 scenario: Option<&str>,
358 ) -> Result<Value> {
359 Self::generate_from_media_type_with_scenario_and_mode(
360 spec,
361 media_type,
362 expand_tokens,
363 scenario,
364 None,
365 None,
366 )
367 }
368
369 fn generate_from_media_type_with_scenario_and_mode(
371 spec: &OpenApiSpec,
372 media_type: &openapiv3::MediaType,
373 expand_tokens: bool,
374 scenario: Option<&str>,
375 selection_mode: Option<crate::openapi::response_selection::ResponseSelectionMode>,
376 selector: Option<&crate::openapi::response_selection::ResponseSelector>,
377 ) -> Result<Value> {
378 if let Some(example) = &media_type.example {
380 tracing::debug!("Using explicit example from media type: {:?}", example);
381 if expand_tokens {
383 let expanded_example = Self::expand_templates(example);
384 return Ok(expanded_example);
385 } else {
386 return Ok(example.clone());
387 }
388 }
389
390 if !media_type.examples.is_empty() {
392 use crate::openapi::response_selection::{ResponseSelectionMode, ResponseSelector};
393
394 if let Some(scenario_name) = scenario {
396 if let Some(example_ref) = media_type.examples.get(scenario_name) {
397 tracing::debug!("Using scenario '{}' from examples map", scenario_name);
398 return Self::extract_example_value(spec, example_ref, expand_tokens);
399 } else {
400 tracing::warn!(
401 "Scenario '{}' not found in examples, falling back based on selection mode",
402 scenario_name
403 );
404 }
405 }
406
407 let mode = selection_mode.unwrap_or(ResponseSelectionMode::First);
409
410 let example_names: Vec<String> = media_type.examples.keys().cloned().collect();
412
413 if example_names.is_empty() {
414 } else if mode == ResponseSelectionMode::Scenario && scenario.is_some() {
416 } else {
418 let selected_index = if let Some(sel) = selector {
420 sel.select(&example_names)
421 } else {
422 let temp_selector = ResponseSelector::new(mode);
424 temp_selector.select(&example_names)
425 };
426
427 if let Some(example_name) = example_names.get(selected_index) {
428 if let Some(example_ref) = media_type.examples.get(example_name) {
429 tracing::debug!(
430 "Using example '{}' from examples map (mode: {:?}, index: {})",
431 example_name,
432 mode,
433 selected_index
434 );
435 return Self::extract_example_value(spec, example_ref, expand_tokens);
436 }
437 }
438 }
439
440 if let Some((example_name, example_ref)) = media_type.examples.iter().next() {
442 tracing::debug!(
443 "Using first example '{}' from examples map as fallback",
444 example_name
445 );
446 return Self::extract_example_value(spec, example_ref, expand_tokens);
447 }
448 }
449
450 if let Some(schema_ref) = &media_type.schema {
452 Ok(Self::generate_example_from_schema_ref(spec, schema_ref))
453 } else {
454 Ok(Value::Object(serde_json::Map::new()))
455 }
456 }
457
458 fn extract_example_value(
460 spec: &OpenApiSpec,
461 example_ref: &ReferenceOr<openapiv3::Example>,
462 expand_tokens: bool,
463 ) -> Result<Value> {
464 match example_ref {
465 ReferenceOr::Item(example) => {
466 if let Some(value) = &example.value {
467 tracing::debug!("Using example from examples map: {:?}", value);
468 if expand_tokens {
469 return Ok(Self::expand_templates(value));
470 } else {
471 return Ok(value.clone());
472 }
473 }
474 }
475 ReferenceOr::Reference { reference } => {
476 if let Some(example) = spec.get_example(reference) {
478 if let Some(value) = &example.value {
479 tracing::debug!("Using resolved example reference: {:?}", value);
480 if expand_tokens {
481 return Ok(Self::expand_templates(value));
482 } else {
483 return Ok(value.clone());
484 }
485 }
486 } else {
487 tracing::warn!("Example reference '{}' not found", reference);
488 }
489 }
490 }
491 Ok(Value::Object(serde_json::Map::new()))
492 }
493
494 fn generate_example_from_schema_ref(
495 spec: &OpenApiSpec,
496 schema_ref: &ReferenceOr<Schema>,
497 ) -> Value {
498 match schema_ref {
499 ReferenceOr::Item(schema) => Self::generate_example_from_schema(spec, schema),
500 ReferenceOr::Reference { reference } => spec
501 .get_schema(reference)
502 .map(|schema| Self::generate_example_from_schema(spec, &schema.schema))
503 .unwrap_or_else(|| Value::Object(serde_json::Map::new())),
504 }
505 }
506
507 fn generate_example_from_schema(spec: &OpenApiSpec, schema: &Schema) -> Value {
514 if let Some(example) = schema.schema_data.example.as_ref() {
517 tracing::debug!("Using schema-level example: {:?}", example);
518 return example.clone();
519 }
520
521 match &schema.schema_kind {
525 openapiv3::SchemaKind::Type(openapiv3::Type::String(_)) => {
526 Value::String("example string".to_string())
528 }
529 openapiv3::SchemaKind::Type(openapiv3::Type::Integer(_)) => Value::Number(42.into()),
530 openapiv3::SchemaKind::Type(openapiv3::Type::Number(_)) => {
531 Value::Number(serde_json::Number::from_f64(std::f64::consts::PI).unwrap())
532 }
533 openapiv3::SchemaKind::Type(openapiv3::Type::Boolean(_)) => Value::Bool(true),
534 openapiv3::SchemaKind::Type(openapiv3::Type::Object(obj)) => {
535 let mut map = serde_json::Map::new();
536 for (prop_name, prop_schema) in &obj.properties {
537 let value = match prop_schema {
538 ReferenceOr::Item(prop_schema) => {
539 if let Some(prop_example) = prop_schema.schema_data.example.as_ref() {
541 tracing::debug!(
542 "Using example for property '{}': {:?}",
543 prop_name,
544 prop_example
545 );
546 prop_example.clone()
547 } else {
548 Self::generate_example_from_schema(spec, prop_schema.as_ref())
549 }
550 }
551 ReferenceOr::Reference { reference } => {
552 if let Some(resolved_schema) = spec.get_schema(reference) {
554 if let Some(ref_example) =
555 resolved_schema.schema.schema_data.example.as_ref()
556 {
557 tracing::debug!(
558 "Using example from referenced schema '{}': {:?}",
559 reference,
560 ref_example
561 );
562 ref_example.clone()
563 } else {
564 Self::generate_example_from_schema(
565 spec,
566 &resolved_schema.schema,
567 )
568 }
569 } else {
570 Self::generate_example_for_property(prop_name)
571 }
572 }
573 };
574 let value = match value {
575 Value::Null => Self::generate_example_for_property(prop_name),
576 Value::Object(ref obj) if obj.is_empty() => {
577 Self::generate_example_for_property(prop_name)
578 }
579 _ => value,
580 };
581 map.insert(prop_name.clone(), value);
582 }
583 Value::Object(map)
584 }
585 openapiv3::SchemaKind::Type(openapiv3::Type::Array(arr)) => {
586 match &arr.items {
592 Some(item_schema) => {
593 let example_item = match item_schema {
594 ReferenceOr::Item(item_schema) => {
595 Self::generate_example_from_schema(spec, item_schema.as_ref())
598 }
599 ReferenceOr::Reference { reference } => {
600 if let Some(resolved_schema) = spec.get_schema(reference) {
603 Self::generate_example_from_schema(
604 spec,
605 &resolved_schema.schema,
606 )
607 } else {
608 Value::Object(serde_json::Map::new())
609 }
610 }
611 };
612 Value::Array(vec![example_item])
613 }
614 None => Value::Array(vec![Value::String("item".to_string())]),
615 }
616 }
617 _ => Value::Object(serde_json::Map::new()),
618 }
619 }
620
621 fn generate_example_for_property(prop_name: &str) -> Value {
623 let prop_lower = prop_name.to_lowercase();
624
625 if prop_lower.contains("id") || prop_lower.contains("uuid") {
627 Value::String(uuid::Uuid::new_v4().to_string())
628 } else if prop_lower.contains("email") {
629 Value::String(format!("user{}@example.com", rng().random_range(1000..=9999)))
630 } else if prop_lower.contains("name") || prop_lower.contains("title") {
631 let names = ["John Doe", "Jane Smith", "Bob Johnson", "Alice Brown"];
632 Value::String(names[rng().random_range(0..names.len())].to_string())
633 } else if prop_lower.contains("phone") || prop_lower.contains("mobile") {
634 Value::String(format!("+1-555-{:04}", rng().random_range(1000..=9999)))
635 } else if prop_lower.contains("address") || prop_lower.contains("street") {
636 let streets = ["123 Main St", "456 Oak Ave", "789 Pine Rd", "321 Elm St"];
637 Value::String(streets[rng().random_range(0..streets.len())].to_string())
638 } else if prop_lower.contains("city") {
639 let cities = ["New York", "London", "Tokyo", "Paris", "Sydney"];
640 Value::String(cities[rng().random_range(0..cities.len())].to_string())
641 } else if prop_lower.contains("country") {
642 let countries = ["USA", "UK", "Japan", "France", "Australia"];
643 Value::String(countries[rng().random_range(0..countries.len())].to_string())
644 } else if prop_lower.contains("company") || prop_lower.contains("organization") {
645 let companies = ["Acme Corp", "Tech Solutions", "Global Inc", "Innovate Ltd"];
646 Value::String(companies[rng().random_range(0..companies.len())].to_string())
647 } else if prop_lower.contains("url") || prop_lower.contains("website") {
648 Value::String("https://example.com".to_string())
649 } else if prop_lower.contains("age") {
650 Value::Number((18 + rng().random_range(0..60)).into())
651 } else if prop_lower.contains("count") || prop_lower.contains("quantity") {
652 Value::Number((1 + rng().random_range(0..100)).into())
653 } else if prop_lower.contains("price")
654 || prop_lower.contains("amount")
655 || prop_lower.contains("cost")
656 {
657 Value::Number(
658 serde_json::Number::from_f64(
659 (rng().random::<f64>() * 1000.0 * 100.0).round() / 100.0,
660 )
661 .unwrap(),
662 )
663 } else if prop_lower.contains("active")
664 || prop_lower.contains("enabled")
665 || prop_lower.contains("is_")
666 {
667 Value::Bool(rng().random_bool(0.5))
668 } else if prop_lower.contains("date") || prop_lower.contains("time") {
669 Value::String(chrono::Utc::now().to_rfc3339())
670 } else if prop_lower.contains("description") || prop_lower.contains("comment") {
671 Value::String("This is a sample description text.".to_string())
672 } else {
673 Value::String(format!("example {}", prop_name))
674 }
675 }
676
677 pub fn generate_from_examples(
679 response: &Response,
680 content_type: Option<&str>,
681 ) -> Result<Option<Value>> {
682 use openapiv3::ReferenceOr;
683
684 if let Some(content_type) = content_type {
686 if let Some(media_type) = response.content.get(content_type) {
687 if let Some(example) = &media_type.example {
689 return Ok(Some(example.clone()));
690 }
691
692 for (_, example_ref) in &media_type.examples {
694 if let ReferenceOr::Item(example) = example_ref {
695 if let Some(value) = &example.value {
696 return Ok(Some(value.clone()));
697 }
698 }
699 }
701 }
702 }
703
704 for (_, media_type) in &response.content {
706 if let Some(example) = &media_type.example {
708 return Ok(Some(example.clone()));
709 }
710
711 for (_, example_ref) in &media_type.examples {
713 if let ReferenceOr::Item(example) = example_ref {
714 if let Some(value) = &example.value {
715 return Ok(Some(value.clone()));
716 }
717 }
718 }
720 }
721
722 Ok(None)
723 }
724
725 fn expand_templates(value: &Value) -> Value {
727 match value {
728 Value::String(s) => {
729 let expanded = s
730 .replace("{{now}}", &chrono::Utc::now().to_rfc3339())
731 .replace("{{uuid}}", &uuid::Uuid::new_v4().to_string());
732 Value::String(expanded)
733 }
734 Value::Object(map) => {
735 let mut new_map = serde_json::Map::new();
736 for (key, val) in map {
737 new_map.insert(key.clone(), Self::expand_templates(val));
738 }
739 Value::Object(new_map)
740 }
741 Value::Array(arr) => {
742 let new_arr: Vec<Value> = arr.iter().map(Self::expand_templates).collect();
743 Value::Array(new_arr)
744 }
745 _ => value.clone(),
746 }
747 }
748}
749
750#[cfg(test)]
751mod tests {
752 use super::*;
753 use openapiv3::ReferenceOr;
754
755 #[test]
756 fn generates_example_using_referenced_schemas() {
757 let yaml = r#"
758openapi: 3.0.3
759info:
760 title: Test API
761 version: "1.0.0"
762paths:
763 /apiaries:
764 get:
765 responses:
766 '200':
767 description: ok
768 content:
769 application/json:
770 schema:
771 $ref: '#/components/schemas/Apiary'
772components:
773 schemas:
774 Apiary:
775 type: object
776 properties:
777 id:
778 type: string
779 hive:
780 $ref: '#/components/schemas/Hive'
781 Hive:
782 type: object
783 properties:
784 name:
785 type: string
786 active:
787 type: boolean
788 "#;
789
790 let spec = OpenApiSpec::from_string(yaml, Some("yaml")).expect("load spec");
791 let path_item = spec
792 .spec
793 .paths
794 .paths
795 .get("/apiaries")
796 .and_then(ReferenceOr::as_item)
797 .expect("path item");
798 let operation = path_item.get.as_ref().expect("GET operation");
799
800 let response =
801 ResponseGenerator::generate_response(&spec, operation, 200, Some("application/json"))
802 .expect("generate response");
803
804 let obj = response.as_object().expect("response object");
805 assert!(obj.contains_key("id"));
806 let hive = obj.get("hive").and_then(|value| value.as_object()).expect("hive object");
807 assert!(hive.contains_key("name"));
808 assert!(hive.contains_key("active"));
809 }
810}
811
812#[derive(Debug, Clone)]
814pub struct MockResponse {
815 pub status_code: u16,
817 pub headers: HashMap<String, String>,
819 pub body: Option<Value>,
821}
822
823impl MockResponse {
824 pub fn new(status_code: u16) -> Self {
826 Self {
827 status_code,
828 headers: HashMap::new(),
829 body: None,
830 }
831 }
832
833 pub fn with_header(mut self, name: String, value: String) -> Self {
835 self.headers.insert(name, value);
836 self
837 }
838
839 pub fn with_body(mut self, body: Value) -> Self {
841 self.body = Some(body);
842 self
843 }
844}
845
846#[derive(Debug, Clone)]
848pub struct OpenApiSecurityRequirement {
849 pub scheme: String,
851 pub scopes: Vec<String>,
853}
854
855impl OpenApiSecurityRequirement {
856 pub fn new(scheme: String, scopes: Vec<String>) -> Self {
858 Self { scheme, scopes }
859 }
860}
861
862#[derive(Debug, Clone)]
864pub struct OpenApiOperation {
865 pub method: String,
867 pub path: String,
869 pub operation: openapiv3::Operation,
871}
872
873impl OpenApiOperation {
874 pub fn new(method: String, path: String, operation: openapiv3::Operation) -> Self {
876 Self {
877 method,
878 path,
879 operation,
880 }
881 }
882}