use axum::{
extract::{Path, State},
http::StatusCode,
response::Json,
routing::post,
Router,
};
use mockforge_core::{
ai_studio::{
api_critique::{ApiCritiqueEngine, CritiqueRequest},
artifact_freezer::{ArtifactFreezer, FreezeMetadata, FreezeRequest},
behavioral_simulator::{BehavioralSimulator, CreateAgentRequest, SimulateBehaviorRequest},
config::DeterministicModeConfig,
system_generator::{SystemGenerationRequest, SystemGenerator},
},
intelligent_behavior::IntelligentBehaviorConfig,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};
use tracing::{error, info, warn};
#[derive(Clone)]
pub struct AiStudioState {
pub critique_engine: Arc<ApiCritiqueEngine>,
pub system_generator: Arc<SystemGenerator>,
pub behavioral_simulator: Arc<Mutex<BehavioralSimulator>>,
pub artifact_freezer: Arc<ArtifactFreezer>,
pub config: IntelligentBehaviorConfig,
pub deterministic_config: Option<DeterministicModeConfig>,
pub workspace_id: Option<String>,
pub system_storage:
Arc<RwLock<HashMap<String, mockforge_core::ai_studio::system_generator::GeneratedSystem>>>,
}
impl AiStudioState {
pub fn new(config: IntelligentBehaviorConfig) -> Self {
let critique_engine = Arc::new(ApiCritiqueEngine::new(config.clone()));
let system_generator = Arc::new(SystemGenerator::new(config.clone()));
let behavioral_simulator = Arc::new(Mutex::new(BehavioralSimulator::new(config.clone())));
let artifact_freezer = Arc::new(ArtifactFreezer::new());
Self {
critique_engine,
system_generator,
behavioral_simulator,
artifact_freezer,
config,
deterministic_config: None,
workspace_id: None,
system_storage: Arc::new(RwLock::new(HashMap::new())),
}
}
pub fn with_workspace(
mut self,
workspace_id: String,
deterministic_config: Option<DeterministicModeConfig>,
) -> Self {
self.workspace_id = Some(workspace_id);
self.deterministic_config = deterministic_config;
self
}
}
#[derive(Debug, Deserialize, Serialize)]
pub struct ApiCritiqueRequest {
pub schema: serde_json::Value,
pub schema_type: String,
#[serde(default)]
pub focus_areas: Vec<String>,
pub workspace_id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct ApiCritiqueResponse {
pub critique: mockforge_core::ai_studio::api_critique::ApiCritique,
#[serde(skip_serializing_if = "Option::is_none")]
pub artifact_id: Option<String>,
pub frozen: bool,
}
pub async fn api_critique_handler(
State(state): State<AiStudioState>,
Json(request): Json<ApiCritiqueRequest>,
) -> Result<Json<ApiCritiqueResponse>, StatusCode> {
info!("API critique request received for schema type: {}", request.schema_type);
let critique_request = CritiqueRequest {
schema: request.schema,
schema_type: request.schema_type,
focus_areas: request.focus_areas,
workspace_id: request.workspace_id.or_else(|| state.workspace_id.clone()),
};
let critique = match state.critique_engine.analyze(&critique_request).await {
Ok(c) => c,
Err(e) => {
error!("Failed to generate API critique: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
let critique_json = match serde_json::to_value(&critique) {
Ok(v) => v,
Err(e) => {
error!("Failed to serialize critique: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
let freeze_request = FreezeRequest {
artifact_type: "api_critique".to_string(),
content: critique_json,
format: "json".to_string(),
path: None,
metadata: Some(FreezeMetadata {
llm_provider: Some(state.config.behavior_model.llm_provider.clone()),
llm_model: Some(state.config.behavior_model.model.clone()),
llm_version: None,
prompt_hash: None,
output_hash: None,
original_prompt: None,
}),
};
let frozen_artifact = match state.artifact_freezer.freeze(&freeze_request).await {
Ok(a) => a,
Err(e) => {
error!("Failed to freeze critique artifact: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
info!(
"API critique completed. Artifact path: {}, Tokens used: {:?}, Cost: ${:.4}",
frozen_artifact.path,
critique.tokens_used,
critique.cost_usd.unwrap_or(0.0)
);
Ok(Json(ApiCritiqueResponse {
critique,
artifact_id: Some(frozen_artifact.path),
frozen: true,
}))
}
#[derive(Debug, Deserialize)]
pub struct SystemGenerationHttpRequest {
pub description: String,
#[serde(default)]
pub output_formats: Vec<String>,
pub workspace_id: Option<String>,
pub system_id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct SystemGenerationResponse {
pub system: mockforge_core::ai_studio::system_generator::GeneratedSystem,
}
#[derive(Debug, Deserialize)]
pub struct ApplySystemRequest {
#[serde(default)]
pub artifact_ids: Option<Vec<String>>,
}
#[derive(Debug, Serialize)]
pub struct ApplySystemResponse {
pub applied: mockforge_core::ai_studio::system_generator::AppliedSystem,
}
#[derive(Debug, Deserialize)]
pub struct FreezeArtifactsRequest {
pub artifact_ids: Vec<String>,
}
#[derive(Debug, Serialize)]
pub struct FreezeArtifactsResponse {
pub frozen_paths: Vec<String>,
}
pub async fn generate_system_handler(
State(state): State<AiStudioState>,
Json(request): Json<SystemGenerationHttpRequest>,
) -> Result<Json<SystemGenerationResponse>, StatusCode> {
info!("System generation request received");
let generation_request = SystemGenerationRequest {
description: request.description,
output_formats: request.output_formats,
workspace_id: request.workspace_id.or_else(|| state.workspace_id.clone()),
system_id: request.system_id,
};
let system = match state
.system_generator
.generate(&generation_request, state.deterministic_config.as_ref())
.await
{
Ok(s) => s,
Err(e) => {
error!("Failed to generate system: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
{
let mut storage = state.system_storage.write().await;
storage.insert(system.system_id.clone(), system.clone());
}
info!(
"System generation completed. System ID: {}, Version: {}, Status: {}, Tokens: {:?}, Cost: ${:.4}",
system.system_id,
system.version,
system.status,
system.tokens_used,
system.cost_usd.unwrap_or(0.0)
);
Ok(Json(SystemGenerationResponse { system }))
}
pub async fn apply_system_handler(
State(state): State<AiStudioState>,
Path(system_id): Path<String>,
Json(request): Json<ApplySystemRequest>,
) -> Result<Json<ApplySystemResponse>, StatusCode> {
info!("Apply system request received for system: {}", system_id);
let system = {
let storage = state.system_storage.read().await;
storage.get(&system_id).cloned()
};
let system = match system {
Some(s) => s,
None => {
warn!("System not found: {}", system_id);
return Err(StatusCode::NOT_FOUND);
}
};
let applied = match state
.system_generator
.apply_system_design(
&system,
state.deterministic_config.as_ref(),
request.artifact_ids.clone(),
)
.await
{
Ok(a) => a,
Err(e) => {
error!("Failed to apply system design: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
if applied.frozen {
let mut storage = state.system_storage.write().await;
if let Some(stored_system) = storage.get_mut(&system_id) {
stored_system.status = "frozen".to_string();
}
}
info!(
"System design applied. System ID: {}, Applied artifacts: {}, Frozen: {}",
applied.system_id,
applied.applied_artifacts.len(),
applied.frozen
);
Ok(Json(ApplySystemResponse { applied }))
}
pub async fn freeze_artifacts_handler(
State(state): State<AiStudioState>,
Path(system_id): Path<String>,
Json(request): Json<FreezeArtifactsRequest>,
) -> Result<Json<FreezeArtifactsResponse>, StatusCode> {
info!(
"Freeze artifacts request received for system: {}, artifacts: {:?}",
system_id, request.artifact_ids
);
let system = {
let storage = state.system_storage.read().await;
storage.get(&system_id).cloned()
};
let system = match system {
Some(s) => s,
None => {
warn!("System not found: {}", system_id);
return Err(StatusCode::NOT_FOUND);
}
};
let frozen_paths = match state
.system_generator
.freeze_artifacts(&system, request.artifact_ids.clone())
.await
{
Ok(paths) => paths,
Err(e) => {
error!("Failed to freeze artifacts: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
if !frozen_paths.is_empty() {
let mut storage = state.system_storage.write().await;
if let Some(stored_system) = storage.get_mut(&system_id) {
let all_frozen = stored_system
.artifacts
.values()
.all(|artifact| request.artifact_ids.contains(&artifact.artifact_id));
if all_frozen {
stored_system.status = "frozen".to_string();
}
}
}
info!(
"Artifacts frozen. System ID: {}, Frozen paths: {}",
system_id,
frozen_paths.len()
);
Ok(Json(FreezeArtifactsResponse { frozen_paths }))
}
#[derive(Debug, Deserialize)]
pub struct CreateAgentHttpRequest {
pub persona_id: Option<String>,
pub behavior_policy: Option<String>,
pub generate_persona: bool,
pub workspace_id: Option<String>,
}
#[derive(Debug, Serialize)]
pub struct CreateAgentResponse {
pub agent: mockforge_core::ai_studio::behavioral_simulator::NarrativeAgent,
}
#[derive(Debug, Deserialize)]
pub struct SimulateBehaviorHttpRequest {
pub agent_id: Option<String>,
pub persona_id: Option<String>,
pub current_state: mockforge_core::ai_studio::behavioral_simulator::AppState,
pub trigger_event: Option<String>,
pub workspace_id: Option<String>,
}
#[axum::debug_handler]
pub async fn create_agent_handler(
State(_state): State<AiStudioState>,
Json(_request): Json<CreateAgentHttpRequest>,
) -> Result<Json<CreateAgentResponse>, StatusCode> {
info!("Create agent request received");
let state = _state;
let request = _request;
let create_request = CreateAgentRequest {
persona_id: request.persona_id,
behavior_policy: request.behavior_policy,
generate_persona: request.generate_persona,
workspace_id: request.workspace_id.or_else(|| state.workspace_id.clone()),
};
let mut simulator = state.behavioral_simulator.lock().await;
let agent = match simulator.create_agent(&create_request).await {
Ok(a) => a,
Err(e) => {
error!("Failed to create agent: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
info!("Agent created: {}", agent.agent_id);
Ok(Json(CreateAgentResponse { agent }))
}
#[axum::debug_handler]
pub async fn simulate_behavior_handler(
State(_state): State<AiStudioState>,
Json(_request): Json<SimulateBehaviorHttpRequest>,
) -> Result<
Json<mockforge_core::ai_studio::behavioral_simulator::SimulateBehaviorResponse>,
StatusCode,
> {
info!("Simulate behavior request received");
let state = _state;
let request = _request;
let simulate_request = SimulateBehaviorRequest {
agent_id: request.agent_id,
persona_id: request.persona_id,
current_state: request.current_state,
trigger_event: request.trigger_event,
workspace_id: request.workspace_id.or_else(|| state.workspace_id.clone()),
};
let mut simulator = state.behavioral_simulator.lock().await;
let response = match simulator.simulate_behavior(&simulate_request).await {
Ok(r) => r,
Err(e) => {
error!("Failed to simulate behavior: {}", e);
return Err(StatusCode::INTERNAL_SERVER_ERROR);
}
};
info!(
"Behavior simulation completed. Intention: {:?}, Tokens: {:?}, Cost: ${:.4}",
response.intention,
response.tokens_used,
response.cost_usd.unwrap_or(0.0)
);
Ok(Json(response))
}
pub fn ai_studio_router(state: AiStudioState) -> Router {
let router = Router::new()
.route("/api-critique", post(api_critique_handler))
.route("/generate-system", post(generate_system_handler))
.route("/system/{system_id}/apply", post(apply_system_handler))
.route("/system/{system_id}/freeze", post(freeze_artifacts_handler))
.route("/simulate-behavior/create-agent", post(create_agent_handler))
.route("/simulate-behavior", post(simulate_behavior_handler));
router.with_state(state)
}
#[cfg(test)]
mod tests {
use super::*;
use mockforge_core::intelligent_behavior::config::BehaviorModelConfig;
fn create_test_state() -> AiStudioState {
let config = IntelligentBehaviorConfig {
behavior_model: BehaviorModelConfig {
llm_provider: "ollama".to_string(),
model: "llama2".to_string(),
api_endpoint: Some("http://localhost:11434/api/chat".to_string()),
api_key: None,
temperature: 0.7,
max_tokens: 2000,
rules: mockforge_core::intelligent_behavior::types::BehaviorRules::default(),
},
..Default::default()
};
AiStudioState::new(config)
}
#[test]
fn test_ai_studio_state_creation() {
let state = create_test_state();
assert!(state.workspace_id.is_none());
assert!(state.deterministic_config.is_none());
}
#[test]
fn test_api_critique_request_serialization() {
let request = ApiCritiqueRequest {
schema: serde_json::json!({"openapi": "3.0.0"}),
schema_type: "openapi".to_string(),
focus_areas: vec!["anti-patterns".to_string()],
workspace_id: None,
};
let json = serde_json::to_string(&request).unwrap();
assert!(json.contains("openapi"));
assert!(json.contains("anti-patterns"));
}
}