1use 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#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct ChatRequest {
18 pub message: String,
20
21 pub context: Option<ChatContext>,
23
24 pub workspace_id: Option<String>,
26
27 pub org_id: Option<String>,
29
30 pub user_id: Option<String>,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct ChatContext {
37 pub history: Vec<ChatMessage>,
39
40 #[serde(default)]
42 pub workspace_id: Option<String>,
43}
44
45#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ChatMessage {
48 pub role: String,
50
51 pub content: String,
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct ChatResponse {
58 pub intent: ChatIntent,
60
61 pub message: String,
63
64 pub data: Option<serde_json::Value>,
66
67 pub error: Option<String>,
69
70 pub tokens_used: Option<u64>,
72
73 pub cost_usd: Option<f64>,
75}
76
77#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
79#[serde(rename_all = "snake_case")]
80pub enum ChatIntent {
81 GenerateMock,
83
84 DebugTest,
86
87 GeneratePersona,
89
90 ContractDiff,
92
93 ApiCritique,
95
96 GenerateSystem,
98
99 SimulateBehavior,
101
102 General,
104
105 Unknown,
107}
108
109pub struct ChatOrchestrator {
111 #[allow(dead_code)]
113 llm_client: LlmClient,
114
115 config: IntelligentBehaviorConfig,
117
118 budget_manager: BudgetManager,
120 ai_mode: Option<AiMode>,
122 deterministic_config: Option<DeterministicModeConfig>,
123}
124
125impl ChatOrchestrator {
126 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 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 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 #[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 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 pub async fn process(&self, request: &ChatRequest) -> Result<ChatResponse> {
198 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 let intent = self.detect_intent(&message_with_context).await?;
207
208 match intent {
210 ChatIntent::GenerateMock => {
211 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 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 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 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 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 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 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 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 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 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 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 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 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 async fn detect_intent(&self, message: &str) -> Result<ChatIntent> {
506 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 Ok(ChatIntent::General)
555 }
556}