use crate::ai_studio::artifact_freezer::FreezeMetadata;
use crate::ai_studio::config::DeterministicModeConfig;
use crate::ai_studio::frozen_generation::{self, FrozenGeneration};
use crate::intelligent_behavior::llm_client::LlmClient;
use crate::intelligent_behavior::types::LlmGenerationRequest;
use crate::intelligent_behavior::IntelligentBehaviorConfig;
use mockforge_foundation::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
pub struct PersonaGenerator {
llm_client: LlmClient,
config: IntelligentBehaviorConfig,
}
impl PersonaGenerator {
pub fn new() -> Self {
let config = IntelligentBehaviorConfig::default();
Self {
llm_client: LlmClient::new(config.behavior_model.clone()),
config,
}
}
pub fn with_config(config: IntelligentBehaviorConfig) -> Self {
Self {
llm_client: LlmClient::new(config.behavior_model.clone()),
config,
}
}
pub async fn generate(
&self,
request: &PersonaGenerationRequest,
ai_mode: Option<crate::ai_studio::config::AiMode>,
deterministic_config: Option<&DeterministicModeConfig>,
) -> Result<PersonaGenerationResponse> {
let snapshot = FrozenGeneration::new(
serde_json::to_value(request)?,
&self.config,
ai_mode,
deterministic_config,
)?;
if let Some(frozen) = snapshot.load("persona").await? {
let persona = frozen_generation::payload(&frozen);
return Ok(PersonaGenerationResponse {
replayed: true,
persona: Some(persona),
message: format!(
"Loaded frozen persona artifact from {} (deterministic mode)",
frozen.path
),
frozen_artifact: Some(frozen),
});
}
let system_prompt = r#"You are an expert at creating realistic user personas for API testing.
Generate a complete persona profile from a natural language description.
For the persona, provide:
1. A unique ID (e.g., "user:premium-001", "customer:churned-002")
2. A descriptive name
3. A business domain (e.g., "ecommerce", "saas", "banking", "healthcare")
4. Realistic traits as key-value pairs (e.g., "subscription_tier": "premium", "spending_level": "high")
5. A narrative backstory explaining the persona's characteristics
6. Optional lifecycle state (e.g., "active", "trial", "churned", "premium")
Return your response as a JSON object with this structure:
{
"id": "string (unique persona ID)",
"name": "string (descriptive name)",
"domain": "string (business domain)",
"traits": {
"trait_name": "trait_value",
...
},
"backstory": "string (narrative description)",
"lifecycle_state": "string (optional, e.g., active, trial, churned)",
"metadata": {
"additional": "metadata fields"
}
}
Make the persona realistic and consistent. Traits should align with the description."#;
let user_prompt =
format!("Generate a persona from this description:\n\n{}", request.description);
let llm_request = LlmGenerationRequest {
system_prompt: system_prompt.to_string(),
user_prompt,
temperature: 0.7, max_tokens: 1500,
schema: None,
seed: None,
};
let response = self.llm_client.generate(&llm_request).await?;
let persona_json = if let Some(_id) = response.get("id") {
response.clone()
} else {
let uuid_str = uuid::Uuid::new_v4().to_string();
let short_id = uuid_str.split('-').next().unwrap_or("generated");
serde_json::json!({
"id": format!("user:generated-{}", short_id),
"name": response.get("name").and_then(|v| v.as_str()).unwrap_or("Generated Persona"),
"domain": response.get("domain").and_then(|v| v.as_str()).unwrap_or("general"),
"traits": response.get("traits").cloned().unwrap_or_else(|| serde_json::json!({})),
"backstory": response.get("backstory").and_then(|v| v.as_str()).unwrap_or("AI-generated persona"),
"lifecycle_state": response.get("lifecycle_state").and_then(|v| v.as_str()).unwrap_or("active"),
})
};
let persona_name = persona_json
.get("name")
.and_then(|v| v.as_str())
.unwrap_or("Generated Persona")
.to_string();
let traits: HashMap<String, String> = persona_json
.get("traits")
.and_then(|v| v.as_object())
.map(|obj| {
obj.iter()
.filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
.collect()
})
.unwrap_or_default();
let persona_value = serde_json::json!({
"name": persona_name,
"traits": traits,
"id": persona_json.get("id"),
"domain": persona_json.get("domain"),
"backstory": persona_json.get("backstory"),
"lifecycle_state": persona_json.get("lifecycle_state"),
});
let frozen_artifact = snapshot
.freeze(
"persona",
persona_value.clone(),
FreezeMetadata {
llm_provider: Some(self.config.behavior_model.llm_provider.clone()),
llm_model: Some(self.config.behavior_model.model.clone()),
llm_version: None,
prompt_hash: None,
output_hash: None,
original_prompt: Some(request.description.clone()),
},
)
.await?;
let persona_value = frozen_artifact
.as_ref()
.map(frozen_generation::payload)
.unwrap_or(persona_value);
Ok(PersonaGenerationResponse {
replayed: false,
persona: Some(persona_value),
message: format!(
"Successfully generated persona '{}' with {} traits{}",
persona_name,
traits.len(),
if frozen_artifact.is_some() {
" (auto-frozen)"
} else {
""
}
),
frozen_artifact,
})
}
pub async fn tweak(
&self,
base_persona: &serde_json::Value,
description: &str,
) -> Result<PersonaGenerationResponse> {
let system_prompt = r#"You are an expert at modifying user personas for API testing.
Given an existing persona and a description of desired changes, update the persona accordingly.
You can:
- Modify existing traits
- Add new traits
- Update the backstory
- Change lifecycle state
- Adjust domain if needed
Return the updated persona in the same JSON structure as the input."#;
let user_prompt = format!(
"Base persona:\n{}\n\nDesired changes: {}\n\nProvide the updated persona.",
serde_json::to_string_pretty(base_persona)?,
description
);
let llm_request = LlmGenerationRequest {
system_prompt: system_prompt.to_string(),
user_prompt,
temperature: 0.5,
max_tokens: 1500,
schema: None,
seed: None,
};
let response = self.llm_client.generate(&llm_request).await?;
Ok(PersonaGenerationResponse {
replayed: false,
persona: Some(response),
message: "Successfully updated persona".to_string(),
frozen_artifact: None,
})
}
}
impl Default for PersonaGenerator {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersonaGenerationRequest {
pub description: String,
pub base_persona_id: Option<String>,
pub workspace_id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PersonaGenerationResponse {
#[serde(default)]
pub replayed: bool,
pub persona: Option<serde_json::Value>,
pub message: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub frozen_artifact: Option<crate::ai_studio::artifact_freezer::FrozenArtifact>,
}