mockforge_intelligence/handlers/
spec_generation.rs

1//! AI-powered spec generation (`POST /__mockforge/ai/generate-spec`).
2//!
3//! Moved from `mockforge_http::management::ai_gen::generate_ai_spec`
4//! under #656. The original took `State<ManagementState>` but never
5//! read it, so this version drops the extractor — axum allows handlers
6//! to skip the state argument even on routers that carry state. The
7//! data-faker / stub-503 contract is preserved via this crate's mirror
8//! `data-faker` feature flag (controlled by `mockforge-http`'s flag of
9//! the same name).
10//!
11//! Only foreign dep is `mockforge_data::rag::*`, already an
12//! unconditional dep of this crate, so the move is cycle-safe.
13
14use axum::{
15    http::StatusCode,
16    response::{IntoResponse, Json},
17};
18use serde::Deserialize;
19
20/// Request for AI-powered API specification generation
21#[derive(Debug, Deserialize)]
22pub struct GenerateSpecRequest {
23    /// Natural language description of the API to generate
24    pub query: String,
25    /// Type of specification to generate: "openapi", "graphql", or "asyncapi"
26    pub spec_type: String,
27    /// Optional API version (e.g., "3.0.0" for OpenAPI)
28    pub api_version: Option<String>,
29}
30
31/// Generate API specification from natural language using AI
32#[cfg(feature = "data-faker")]
33pub async fn generate_ai_spec(Json(request): Json<GenerateSpecRequest>) -> impl IntoResponse {
34    use mockforge_data::rag::{
35        config::{LlmProvider, RagConfig},
36        engine::RagEngine,
37        storage::DocumentStorage,
38    };
39    use std::sync::Arc;
40
41    // Build RAG config from environment variables
42    let api_key = std::env::var("MOCKFORGE_RAG_API_KEY")
43        .ok()
44        .or_else(|| std::env::var("OPENAI_API_KEY").ok());
45
46    // Check if RAG is configured - require API key
47    if api_key.is_none() {
48        return (
49            StatusCode::SERVICE_UNAVAILABLE,
50            Json(serde_json::json!({
51                "error": "AI service not configured",
52                "message": "Please provide an API key via MOCKFORGE_RAG_API_KEY or OPENAI_API_KEY"
53            })),
54        )
55            .into_response();
56    }
57
58    // Build RAG configuration
59    let provider_str = std::env::var("MOCKFORGE_RAG_PROVIDER")
60        .unwrap_or_else(|_| "openai".to_string())
61        .to_lowercase();
62
63    let provider = match provider_str.as_str() {
64        "openai" => LlmProvider::OpenAI,
65        "anthropic" => LlmProvider::Anthropic,
66        "ollama" => LlmProvider::Ollama,
67        "openai-compatible" | "openai_compatible" => LlmProvider::OpenAICompatible,
68        _ => LlmProvider::OpenAI,
69    };
70
71    let api_endpoint =
72        std::env::var("MOCKFORGE_RAG_API_ENDPOINT").unwrap_or_else(|_| match provider {
73            LlmProvider::OpenAI => "https://api.openai.com/v1".to_string(),
74            LlmProvider::Anthropic => "https://api.anthropic.com/v1".to_string(),
75            LlmProvider::Ollama => "http://localhost:11434/api".to_string(),
76            LlmProvider::OpenAICompatible => "http://localhost:8000/v1".to_string(),
77        });
78
79    let model = std::env::var("MOCKFORGE_RAG_MODEL").unwrap_or_else(|_| match provider {
80        LlmProvider::OpenAI => "gpt-3.5-turbo".to_string(),
81        LlmProvider::Anthropic => "claude-3-sonnet-20240229".to_string(),
82        LlmProvider::Ollama => "llama2".to_string(),
83        LlmProvider::OpenAICompatible => "gpt-3.5-turbo".to_string(),
84    });
85
86    // Build RagConfig using struct literal with defaults
87    let rag_config = RagConfig {
88        provider,
89        api_endpoint,
90        api_key,
91        model,
92        max_tokens: std::env::var("MOCKFORGE_RAG_MAX_TOKENS")
93            .unwrap_or_else(|_| "4096".to_string())
94            .parse()
95            .unwrap_or(4096),
96        temperature: std::env::var("MOCKFORGE_RAG_TEMPERATURE")
97            .unwrap_or_else(|_| "0.3".to_string())
98            .parse()
99            .unwrap_or(0.3), // Lower temperature for more structured output
100        timeout_secs: std::env::var("MOCKFORGE_RAG_TIMEOUT")
101            .unwrap_or_else(|_| "60".to_string())
102            .parse()
103            .unwrap_or(60),
104        max_context_length: std::env::var("MOCKFORGE_RAG_CONTEXT_WINDOW")
105            .unwrap_or_else(|_| "4000".to_string())
106            .parse()
107            .unwrap_or(4000),
108        ..Default::default()
109    };
110
111    // Build the prompt for spec generation
112    let spec_type_label = match request.spec_type.as_str() {
113        "openapi" => "OpenAPI 3.0",
114        "graphql" => "GraphQL",
115        "asyncapi" => "AsyncAPI",
116        _ => "OpenAPI 3.0",
117    };
118
119    let api_version = request.api_version.as_deref().unwrap_or("3.0.0");
120
121    let prompt = format!(
122        r#"You are an expert API architect. Generate a complete {} specification based on the following user requirements.
123
124User Requirements:
125{}
126
127Instructions:
1281. Generate a complete, valid {} specification
1292. Include all paths, operations, request/response schemas, and components
1303. Use realistic field names and data types
1314. Include proper descriptions and examples
1325. Follow {} best practices
1336. Return ONLY the specification, no additional explanation
1347. For OpenAPI, use version {}
135
136Return the specification in {} format."#,
137        spec_type_label,
138        request.query,
139        spec_type_label,
140        spec_type_label,
141        api_version,
142        if request.spec_type == "graphql" {
143            "GraphQL SDL"
144        } else {
145            "YAML"
146        }
147    );
148
149    // Create in-memory storage for RAG engine
150    use mockforge_data::rag::storage::InMemoryStorage;
151    let storage: Arc<dyn DocumentStorage> = Arc::new(InMemoryStorage::new());
152
153    // Create RAG engine
154    let mut rag_engine = match RagEngine::new(rag_config.clone(), storage) {
155        Ok(engine) => engine,
156        Err(e) => {
157            return (
158                StatusCode::INTERNAL_SERVER_ERROR,
159                Json(serde_json::json!({
160                    "error": "Failed to initialize RAG engine",
161                    "message": e.to_string()
162                })),
163            )
164                .into_response();
165        }
166    };
167
168    // Generate using RAG engine
169    match rag_engine.generate(&prompt, None).await {
170        Ok(generated_text) => {
171            // Try to extract just the YAML/JSON/SDL content if LLM added explanation
172            let spec = if request.spec_type == "graphql" {
173                // For GraphQL, extract SDL
174                extract_graphql_schema(&generated_text)
175            } else {
176                // For OpenAPI/AsyncAPI, extract YAML
177                extract_yaml_spec(&generated_text)
178            };
179
180            Json(serde_json::json!({
181                "success": true,
182                "spec": spec,
183                "spec_type": request.spec_type,
184            }))
185            .into_response()
186        }
187        Err(e) => (
188            StatusCode::INTERNAL_SERVER_ERROR,
189            Json(serde_json::json!({
190                "error": "AI generation failed",
191                "message": e.to_string()
192            })),
193        )
194            .into_response(),
195    }
196}
197
198#[cfg(not(feature = "data-faker"))]
199pub async fn generate_ai_spec(Json(_request): Json<GenerateSpecRequest>) -> impl IntoResponse {
200    (
201        StatusCode::NOT_IMPLEMENTED,
202        Json(serde_json::json!({
203            "error": "AI features not enabled",
204            "message": "Please enable the 'data-faker' feature to use AI-powered specification generation"
205        })),
206    )
207        .into_response()
208}
209
210#[cfg(feature = "data-faker")]
211fn extract_yaml_spec(text: &str) -> String {
212    // Try to find YAML code blocks
213    if let Some(start) = text.find("```yaml") {
214        let yaml_start = text[start + 7..].trim_start();
215        if let Some(end) = yaml_start.find("```") {
216            return yaml_start[..end].trim().to_string();
217        }
218    }
219    if let Some(start) = text.find("```") {
220        let content_start = text[start + 3..].trim_start();
221        if let Some(end) = content_start.find("```") {
222            return content_start[..end].trim().to_string();
223        }
224    }
225
226    // Check if it starts with openapi: or asyncapi:
227    if text.trim_start().starts_with("openapi:") || text.trim_start().starts_with("asyncapi:") {
228        return text.trim().to_string();
229    }
230
231    // Return as-is if no code blocks found
232    text.trim().to_string()
233}
234
235/// Extract GraphQL schema from text content
236#[cfg(feature = "data-faker")]
237fn extract_graphql_schema(text: &str) -> String {
238    // Try to find GraphQL code blocks
239    if let Some(start) = text.find("```graphql") {
240        let schema_start = text[start + 10..].trim_start();
241        if let Some(end) = schema_start.find("```") {
242            return schema_start[..end].trim().to_string();
243        }
244    }
245    if let Some(start) = text.find("```") {
246        let content_start = text[start + 3..].trim_start();
247        if let Some(end) = content_start.find("```") {
248            return content_start[..end].trim().to_string();
249        }
250    }
251
252    // Check if it looks like GraphQL SDL (starts with type, schema, etc.)
253    if text.trim_start().starts_with("type ") || text.trim_start().starts_with("schema ") {
254        return text.trim().to_string();
255    }
256
257    text.trim().to_string()
258}
mockforge_intelligence/handlers/spec_generation.rs

mockforge_intelligence/handlers/
spec_generation.rs
