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

mockforge_intelligence/ai_studio/
chat_orchestrator.rs
