Skip to main content

mockforge_intelligence/ai_studio/
chat_orchestrator.rs

1//! Chat orchestrator for routing natural language commands
2//!
3//! This module provides the main entry point for processing natural language
4//! commands and routing them to appropriate handlers based on intent detection.
5
6use crate::ai_studio::budget_manager::{BudgetConfig, BudgetManager};
7use crate::ai_studio::config::{AiMode, DeterministicModeConfig};
8use crate::ai_studio::debug_analyzer::DebugRequest;
9use crate::intelligent_behavior::{
10    config::IntelligentBehaviorConfig, llm_client::LlmClient, LlmUsage,
11};
12use mockforge_foundation::Result;
13use serde::{Deserialize, Serialize};
14
15/// Chat request from user
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct ChatRequest {
18    /// User's message/command
19    pub message: String,
20
21    /// Optional conversation context
22    pub context: Option<ChatContext>,
23
24    /// Optional workspace ID for context
25    pub workspace_id: Option<String>,
26
27    /// Optional organization ID for org-level controls
28    pub org_id: Option<String>,
29
30    /// Optional user ID for audit logging
31    pub user_id: Option<String>,
32}
33
34/// Chat context for multi-turn conversations
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct ChatContext {
37    /// Conversation history
38    pub history: Vec<ChatMessage>,
39
40    /// Optional workspace ID
41    #[serde(default)]
42    pub workspace_id: Option<String>,
43}
44
45/// Chat message in conversation history
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ChatMessage {
48    /// Role (user or assistant)
49    pub role: String,
50
51    /// Message content
52    pub content: String,
53}
54
55/// Chat response from orchestrator
56#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ChatResponse {
58    /// Detected intent
59    pub intent: ChatIntent,
60
61    /// Response message
62    pub message: String,
63
64    /// Optional structured data (e.g., generated spec, persona, etc.)
65    pub data: Option<serde_json::Value>,
66
67    /// Optional error message
68    pub error: Option<String>,
69
70    /// Token usage for this request
71    pub tokens_used: Option<u64>,
72
73    /// Estimated cost in USD
74    pub cost_usd: Option<f64>,
75}
76
77/// Detected intent from user message
78#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
79#[serde(rename_all = "snake_case")]
80pub enum ChatIntent {
81    /// Generate a mock API
82    GenerateMock,
83
84    /// Debug a test failure
85    DebugTest,
86
87    /// Generate or modify a persona
88    GeneratePersona,
89
90    /// Run contract diff analysis
91    ContractDiff,
92
93    /// Critique API architecture
94    ApiCritique,
95
96    /// Generate entire system from description
97    GenerateSystem,
98
99    /// Simulate user behavior
100    SimulateBehavior,
101
102    /// General question/chat
103    General,
104
105    /// Unknown intent
106    Unknown,
107}
108
109/// Chat orchestrator that routes commands to appropriate handlers
110pub struct ChatOrchestrator {
111    /// LLM client for intent detection and processing
112    #[allow(dead_code)]
113    llm_client: LlmClient,
114
115    /// Configuration
116    config: IntelligentBehaviorConfig,
117
118    /// Budget manager for tracking usage
119    budget_manager: BudgetManager,
120    /// Explicit authoring policy supplied by the caller.
121    ai_mode: Option<AiMode>,
122    deterministic_config: Option<DeterministicModeConfig>,
123}
124
125impl ChatOrchestrator {
126    /// Create a new chat orchestrator
127    pub fn new(config: IntelligentBehaviorConfig) -> Self {
128        let llm_client = LlmClient::new(config.behavior_model.clone());
129        let budget_config = BudgetConfig::default();
130        let budget_manager = BudgetManager::new(budget_config);
131        Self {
132            llm_client,
133            config,
134            budget_manager,
135            ai_mode: None,
136            deterministic_config: None,
137        }
138    }
139
140    /// Apply an explicit policy for mock and persona authoring.
141    /// Snapshot lookup happens before provider initialization. The caller owns
142    /// workspace configuration lookup and authorization; chat JSON cannot select
143    /// a filesystem directory or bypass hosted AI controls through this builder.
144    pub fn with_generation_policy(
145        mut self,
146        ai_mode: Option<AiMode>,
147        deterministic_config: Option<DeterministicModeConfig>,
148    ) -> Self {
149        self.ai_mode = ai_mode;
150        self.deterministic_config = deterministic_config;
151        self
152    }
153
154    /// Helper to calculate cost from usage
155    fn calculate_cost(&self, usage: &LlmUsage) -> f64 {
156        let provider = &self.config.behavior_model.llm_provider;
157        let model = &self.config.behavior_model.model;
158        BudgetManager::calculate_cost(provider, model, usage.total_tokens)
159    }
160
161    /// Helper to track usage and return token/cost info
162    #[allow(dead_code)]
163    async fn track_usage(
164        &self,
165        org_id: Option<&str>,
166        workspace_id: &str,
167        user_id: Option<&str>,
168        usage: &LlmUsage,
169    ) -> Result<(Option<u64>, Option<f64>)> {
170        self.track_usage_with_feature(org_id, workspace_id, user_id, usage, None).await
171    }
172
173    /// Helper to track usage with feature information
174    async fn track_usage_with_feature(
175        &self,
176        org_id: Option<&str>,
177        workspace_id: &str,
178        user_id: Option<&str>,
179        usage: &LlmUsage,
180        feature: Option<crate::ai_studio::budget_manager::AiFeature>,
181    ) -> Result<(Option<u64>, Option<f64>)> {
182        let cost = self.calculate_cost(usage);
183        self.budget_manager
184            .record_usage_with_feature(
185                org_id,
186                workspace_id,
187                user_id,
188                usage.total_tokens,
189                cost,
190                feature,
191            )
192            .await?;
193        Ok((Some(usage.total_tokens), Some(cost)))
194    }
195
196    /// Process a chat request and return response
197    pub async fn process(&self, request: &ChatRequest) -> Result<ChatResponse> {
198        // Build message with context if available
199        let message_with_context = if let Some(context) = &request.context {
200            self.build_contextual_message(&request.message, context)
201        } else {
202            request.message.clone()
203        };
204
205        // Detect intent from message
206        let intent = self.detect_intent(&message_with_context).await?;
207
208        // Route to appropriate handler based on intent
209        match intent {
210            ChatIntent::GenerateMock => {
211                // Use MockGenerator to generate mock from message
212                use crate::ai_studio::nl_mock_generator::MockGenerator;
213                let generator = MockGenerator::with_config(self.config.clone());
214                match generator
215                    .generate(
216                        &request.message,
217                        request.workspace_id.as_deref(),
218                        self.ai_mode,
219                        self.deterministic_config.as_ref(),
220                    )
221                    .await
222                {
223                    Ok(result) => {
224                        // Estimate tokens (MockGenerator uses LLM internally, but doesn't expose usage)
225                        // For now, estimate based on message length and response size
226                        let estimated_tokens =
227                            (request.message.len() + result.message.len()) as u64 / 4;
228                        let usage = LlmUsage::new(estimated_tokens / 2, estimated_tokens / 2);
229                        let (tokens, cost) = if result.replayed {
230                            (Some(0), Some(0.0))
231                        } else {
232                            self.track_usage_with_feature(
233                                request.org_id.as_deref(),
234                                &request.workspace_id.clone().unwrap_or_default(),
235                                request.user_id.as_deref(),
236                                &usage,
237                                Some(crate::ai_studio::budget_manager::AiFeature::MockAi),
238                            )
239                            .await
240                            .unwrap_or((None, None))
241                        };
242
243                        Ok(ChatResponse {
244                            intent: ChatIntent::GenerateMock,
245                            message: result.message,
246                            data: result.spec.map(|s| {
247                                serde_json::json!({
248                                    "spec": s,
249                                    "type": "openapi_spec"
250                                })
251                            }),
252                            error: None,
253                            tokens_used: tokens,
254                            cost_usd: cost,
255                        })
256                    }
257                    Err(e) => Ok(ChatResponse {
258                        intent: ChatIntent::GenerateMock,
259                        message: format!("Failed to generate mock: {}", e),
260                        data: None,
261                        error: Some(e.to_string()),
262                        tokens_used: None,
263                        cost_usd: None,
264                    }),
265                }
266            }
267            ChatIntent::DebugTest => {
268                // Use DebugAnalyzer to analyze test failure
269                use crate::ai_studio::debug_analyzer::DebugAnalyzer;
270                let analyzer = DebugAnalyzer::with_config(self.config.clone());
271                let debug_request = DebugRequest {
272                    test_logs: request.message.clone(),
273                    test_name: None,
274                    workspace_id: request.workspace_id.clone(),
275                };
276                match analyzer.analyze(&debug_request).await {
277                    Ok(result) => {
278                        // Estimate tokens (DebugAnalyzer uses LLM internally)
279                        let estimated_tokens =
280                            (request.message.len() + result.root_cause.len()) as u64 / 4;
281                        let usage = LlmUsage::new(estimated_tokens / 2, estimated_tokens / 2);
282                        let (tokens, cost) = self
283                            .track_usage_with_feature(
284                                request.org_id.as_deref(),
285                                &request.workspace_id.clone().unwrap_or_default(),
286                                request.user_id.as_deref(),
287                                &usage,
288                                Some(crate::ai_studio::budget_manager::AiFeature::DebugAnalysis),
289                            )
290                            .await
291                            .unwrap_or((None, None));
292                        Ok(ChatResponse {
293                            intent: ChatIntent::DebugTest,
294                            message: format!("Root cause: {}\n\nFound {} suggestions and {} related configurations.",
295                                result.root_cause, result.suggestions.len(), result.related_configs.len()),
296                            data: Some(serde_json::json!({
297                                "root_cause": result.root_cause,
298                                "suggestions": result.suggestions,
299                                "related_configs": result.related_configs,
300                                "type": "debug_analysis"
301                            })),
302                            error: None,
303                            tokens_used: tokens,
304                            cost_usd: cost,
305                        })
306                    }
307                    Err(e) => Ok(ChatResponse {
308                        intent: ChatIntent::DebugTest,
309                        message: format!("Failed to analyze test failure: {}", e),
310                        data: None,
311                        error: Some(e.to_string()),
312                        tokens_used: None,
313                        cost_usd: None,
314                    }),
315                }
316            }
317            ChatIntent::GeneratePersona => {
318                // Use PersonaGenerator to generate persona from message
319                use crate::ai_studio::persona_generator::{
320                    PersonaGenerationRequest, PersonaGenerator,
321                };
322                let generator = PersonaGenerator::with_config(self.config.clone());
323                let persona_request = PersonaGenerationRequest {
324                    description: request.message.clone(),
325                    base_persona_id: None,
326                    workspace_id: request.workspace_id.clone(),
327                };
328                match generator
329                    .generate(&persona_request, self.ai_mode, self.deterministic_config.as_ref())
330                    .await
331                {
332                    Ok(result) => {
333                        // Estimate tokens (PersonaGenerator uses LLM internally)
334                        let estimated_tokens =
335                            (request.message.len() + result.message.len()) as u64 / 4;
336                        let usage = LlmUsage::new(estimated_tokens / 2, estimated_tokens / 2);
337                        let (tokens, cost) = if result.replayed {
338                            (Some(0), Some(0.0))
339                        } else {
340                            self.track_usage_with_feature(
341                                request.org_id.as_deref(),
342                                &request.workspace_id.clone().unwrap_or_default(),
343                                request.user_id.as_deref(),
344                                &usage,
345                                Some(
346                                    crate::ai_studio::budget_manager::AiFeature::PersonaGeneration,
347                                ),
348                            )
349                            .await
350                            .unwrap_or((None, None))
351                        };
352
353                        Ok(ChatResponse {
354                            intent: ChatIntent::GeneratePersona,
355                            message: result.message,
356                            data: result.persona.map(|p| {
357                                serde_json::json!({
358                                    "persona": p,
359                                    "type": "persona"
360                                })
361                            }),
362                            error: None,
363                            tokens_used: tokens,
364                            cost_usd: cost,
365                        })
366                    }
367                    Err(e) => Ok(ChatResponse {
368                        intent: ChatIntent::GeneratePersona,
369                        message: format!("Failed to generate persona: {}", e),
370                        data: None,
371                        error: Some(e.to_string()),
372                        tokens_used: None,
373                        cost_usd: None,
374                    }),
375                }
376            }
377            ChatIntent::ContractDiff => {
378                // Use ContractDiffHandler to process the query
379                use crate::ai_studio::contract_diff_handler::ContractDiffHandler;
380                let handler = ContractDiffHandler::new().map_err(|e| {
381                    mockforge_foundation::Error::io_with_context(
382                        "ContractDiffHandler",
383                        e.to_string(),
384                    )
385                })?;
386
387                // For now, we don't have direct access to specs/requests in the orchestrator
388                // The handler will provide guidance on how to use contract diff
389                match handler.analyze_from_query(&request.message, None, None).await {
390                    Ok(query_result) => {
391                        let mut message = query_result.summary.clone();
392                        if let Some(link) = &query_result.link_to_viewer {
393                            message.push_str(&format!("\n\nView details: {}", link));
394                        }
395
396                        Ok(ChatResponse {
397                            intent: ChatIntent::ContractDiff,
398                            message,
399                            data: Some(serde_json::json!({
400                                "type": "contract_diff_query",
401                                "intent": query_result.intent,
402                                "result": query_result.result,
403                                "breaking_changes": query_result.breaking_changes,
404                                "link_to_viewer": query_result.link_to_viewer,
405                            })),
406                            error: None,
407                            tokens_used: None,
408                            cost_usd: None,
409                        })
410                    }
411                    Err(e) => Ok(ChatResponse {
412                        intent: ChatIntent::ContractDiff,
413                        message: format!("I can help with contract diff analysis! Try asking:\n- \"Analyze the last captured request\"\n- \"Show me breaking changes\"\n- \"Compare contract versions\"\n\nError: {}", e),
414                        data: Some(serde_json::json!({
415                            "type": "contract_diff_info",
416                            "endpoints": {
417                                "analyze": "/api/v1/contract-diff/analyze",
418                                "capture": "/api/v1/contract-diff/capture",
419                                "compare": "/api/v1/contract-diff/compare"
420                            }
421                        })),
422                        error: Some(e.to_string()),
423                        tokens_used: None,
424                        cost_usd: None,
425                    }),
426                }
427            }
428            ChatIntent::ApiCritique => {
429                // Guide user to use API Critique feature
430                Ok(ChatResponse {
431                    intent: ChatIntent::ApiCritique,
432                    message: "I can help you critique your API architecture! Please use the 'API Critique' tab in AI Studio, or provide your API schema (OpenAPI, GraphQL, or Protobuf) for analysis.".to_string(),
433                    data: Some(serde_json::json!({
434                        "type": "api_critique_info",
435                        "endpoint": "/api/v1/ai-studio/api-critique",
436                        "description": "Analyzes API schemas for anti-patterns, redundancy, naming issues, tone, and restructuring recommendations"
437                    })),
438                    error: None,
439                    tokens_used: None,
440                    cost_usd: None,
441                })
442            }
443            ChatIntent::GenerateSystem => {
444                // Guide user to use System Generator feature
445                Ok(ChatResponse {
446                    intent: ChatIntent::GenerateSystem,
447                    message: format!("I can generate a complete backend system from your description! Use the 'System Designer' tab in AI Studio, or describe your system here. Example: \"{}\"", request.message),
448                    data: Some(serde_json::json!({
449                        "type": "system_generator_info",
450                        "endpoint": "/api/v1/ai-studio/generate-system",
451                        "description": "Generates complete backend systems including OpenAPI specs, personas, lifecycles, WebSocket topics, chaos profiles, CI templates, and more"
452                    })),
453                    error: None,
454                    tokens_used: None,
455                    cost_usd: None,
456                })
457            }
458            ChatIntent::SimulateBehavior => {
459                // Guide user to use Behavioral Simulator feature
460                Ok(ChatResponse {
461                    intent: ChatIntent::SimulateBehavior,
462                    message: "I can simulate user behavior as narrative agents! Use the 'AI User Simulator' tab in AI Studio to create agents, attach them to personas, and simulate multi-step interactions.".to_string(),
463                    data: Some(serde_json::json!({
464                        "type": "behavioral_simulator_info",
465                        "endpoints": {
466                            "create_agent": "/api/v1/ai-studio/simulate-behavior/create-agent",
467                            "simulate": "/api/v1/ai-studio/simulate-behavior"
468                        },
469                        "description": "Models users as narrative agents that react to app state, form intentions, respond to errors, and trigger multi-step interactions"
470                    })),
471                    error: None,
472                    tokens_used: None,
473                    cost_usd: None,
474                })
475            }
476            ChatIntent::General | ChatIntent::Unknown => {
477                // General chat response
478                Ok(ChatResponse {
479                    intent: ChatIntent::General,
480                    message: "I'm here to help! You can ask me to generate mocks, debug tests, create personas, analyze contracts, critique APIs, generate entire systems, or simulate user behavior.".to_string(),
481                    data: None,
482                    error: None,
483                    tokens_used: None,
484                    cost_usd: None,
485                })
486            }
487        }
488    }
489
490    /// Build contextual message from conversation history
491    fn build_contextual_message(&self, current_message: &str, context: &ChatContext) -> String {
492        if context.history.is_empty() {
493            return current_message.to_string();
494        }
495
496        let mut contextual = String::from("Previous conversation:\n");
497        for msg in &context.history {
498            contextual.push_str(&format!("{}: {}\n", msg.role, msg.content));
499        }
500        contextual.push_str(&format!("\nCurrent message: {}", current_message));
501        contextual
502    }
503
504    /// Detect intent from user message using LLM
505    async fn detect_intent(&self, message: &str) -> Result<ChatIntent> {
506        // Use simple keyword matching for now (can be enhanced with LLM)
507        let message_lower = message.to_lowercase();
508
509        if message_lower.contains("create")
510            && (message_lower.contains("api") || message_lower.contains("mock"))
511        {
512            return Ok(ChatIntent::GenerateMock);
513        }
514
515        if message_lower.contains("debug")
516            || message_lower.contains("test") && message_lower.contains("fail")
517        {
518            return Ok(ChatIntent::DebugTest);
519        }
520
521        if message_lower.contains("persona") {
522            return Ok(ChatIntent::GeneratePersona);
523        }
524
525        if message_lower.contains("contract") || message_lower.contains("diff") {
526            return Ok(ChatIntent::ContractDiff);
527        }
528
529        if message_lower.contains("critique")
530            || message_lower.contains("review api")
531            || (message_lower.contains("analyze") && message_lower.contains("api"))
532        {
533            return Ok(ChatIntent::ApiCritique);
534        }
535
536        if message_lower.contains("generate system")
537            || message_lower.contains("build backend")
538            || message_lower.contains("system design")
539            || message_lower.contains("entire system")
540            || (message_lower.contains("i'm building") && message_lower.contains("app"))
541        {
542            return Ok(ChatIntent::GenerateSystem);
543        }
544
545        if message_lower.contains("simulate")
546            || message_lower.contains("user behavior")
547            || message_lower.contains("behavioral")
548            || message_lower.contains("narrative agent")
549        {
550            return Ok(ChatIntent::SimulateBehavior);
551        }
552
553        // Default to general for now
554        Ok(ChatIntent::General)
555    }
556}