mod error;
pub use error::{Error, Result};
pub mod consistency;
pub mod dataset;
pub mod domains;
pub mod drift;
pub mod drift_learning;
pub mod faker;
pub mod generator;
pub mod intelligent_mock;
pub mod mock_generator;
pub mod mock_server;
pub mod persona;
pub mod persona_backstory;
pub mod persona_graph;
pub mod persona_lifecycle;
pub mod persona_lifecycle_response;
pub mod persona_templates;
pub mod provider;
pub mod rag;
pub mod replay_augmentation;
pub mod schema;
pub mod token_resolver;
#[cfg(test)]
mod mock_data_tests;
#[cfg(test)]
mod integration_tests;
pub use consistency::{ConsistencyStore, EntityIdExtractor, EntityType};
pub use dataset::{Dataset, DatasetValidationResult};
pub use domains::{Domain, DomainGenerator, ParseDomainError};
pub use drift::{DataDriftConfig, DataDriftEngine, DriftStrategy};
pub use drift_learning::{
BehaviorEvent, BehaviorEventType, DriftLearningEngine, LearnedPattern, LearningConfig,
LearningMode, PatternType, PersonaBehaviorLearner, TrafficPatternLearner,
};
pub use fake::Faker;
pub use generator::DataGenerator;
pub use intelligent_mock::{IntelligentMockConfig, IntelligentMockGenerator, ResponseMode};
pub use mock_generator::{MockDataGenerator, MockDataResult, MockGeneratorConfig, MockResponse};
pub use mock_server::{
start_mock_server, start_mock_server_with_config, MockServer, MockServerBuilder,
MockServerConfig,
};
pub use persona::{PersonaGenerator, PersonaProfile, PersonaRegistry};
pub use persona_backstory::{BackstoryGenerator, BackstoryTemplate};
pub use persona_graph::{
Edge, GraphStats, GraphVisualization, PersonaGraph, PersonaNode, VisualizationEdge,
VisualizationNode,
};
pub use persona_lifecycle::{
LifecyclePreset, LifecycleScenarios, LifecycleState, PersonaLifecycle, TransitionRule,
};
pub use persona_lifecycle_response::{
apply_billing_lifecycle_effects, apply_lifecycle_effects, apply_loan_lifecycle_effects,
apply_order_fulfillment_lifecycle_effects, apply_support_lifecycle_effects,
apply_user_engagement_lifecycle_effects,
};
pub use persona_templates::{
EcommercePersonaTemplate, FinancePersonaTemplate, HealthcarePersonaTemplate, PersonaTemplate,
PersonaTemplateRegistry,
};
pub use rag::{EmbeddingProvider, LlmProvider, RagConfig, RagEngine, SearchResult};
pub use replay_augmentation::{
EventStrategy, GeneratedEvent, ReplayAugmentationConfig, ReplayAugmentationEngine, ReplayMode,
};
pub use schema::{FieldDefinition, SchemaDefinition};
pub use token_resolver::{resolve_tokens, resolve_tokens_with_rag, TokenResolver, TokenType};
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize, Default)]
pub struct DataConfig {
#[serde(default = "default_rows")]
pub rows: usize,
pub seed: Option<u64>,
pub rag_enabled: bool,
#[serde(default = "default_rag_context_length")]
pub rag_context_length: usize,
pub format: OutputFormat,
}
fn default_rows() -> usize {
100
}
fn default_rag_context_length() -> usize {
1000
}
#[derive(Debug, Clone, serde::Deserialize, serde::Serialize, Default)]
#[serde(rename_all = "lowercase")]
pub enum OutputFormat {
#[default]
Json,
JsonLines,
Yaml,
Csv,
}
#[derive(Debug)]
pub struct GenerationResult {
pub data: Vec<serde_json::Value>,
pub count: usize,
pub generation_time_ms: u128,
pub warnings: Vec<String>,
}
impl GenerationResult {
pub fn new(data: Vec<serde_json::Value>, generation_time_ms: u128) -> Self {
Self {
count: data.len(),
data,
generation_time_ms,
warnings: Vec::new(),
}
}
pub fn with_warning(mut self, warning: String) -> Self {
self.warnings.push(warning);
self
}
pub fn to_json_string(&self) -> Result<String> {
Ok(serde_json::to_string_pretty(&self.data)?)
}
pub fn to_jsonl_string(&self) -> Result<String> {
let lines: Vec<String> = self
.data
.iter()
.map(serde_json::to_string)
.collect::<std::result::Result<_, _>>()?;
Ok(lines.join("\n"))
}
}
pub async fn generate_data(
schema: SchemaDefinition,
config: DataConfig,
) -> Result<GenerationResult> {
let mut generator = DataGenerator::new(schema, config)?;
generator.generate().await
}
pub async fn generate_from_json_schema(
json_schema: &serde_json::Value,
rows: usize,
) -> Result<GenerationResult> {
let schema = SchemaDefinition::from_json_schema(json_schema)?;
let config = DataConfig {
rows,
..Default::default()
};
generate_data(schema, config).await
}
pub async fn generate_from_openapi(
openapi_spec: &serde_json::Value,
rows: usize,
) -> Result<GenerationResult> {
let schema = SchemaDefinition::from_openapi_spec(openapi_spec)?;
let config = DataConfig {
rows,
..Default::default()
};
generate_data(schema, config).await
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_data_config_default() {
let config = DataConfig::default();
assert_eq!(config.rows, 0); assert_eq!(config.seed, None);
assert!(!config.rag_enabled);
assert_eq!(config.rag_context_length, 0); assert!(matches!(config.format, OutputFormat::Json));
}
#[test]
fn test_data_config_custom() {
let config = DataConfig {
rows: 50,
seed: Some(42),
rag_enabled: true,
rag_context_length: 2000,
format: OutputFormat::Csv,
};
assert_eq!(config.rows, 50);
assert_eq!(config.seed, Some(42));
assert!(config.rag_enabled);
assert_eq!(config.rag_context_length, 2000);
assert!(matches!(config.format, OutputFormat::Csv));
}
#[test]
fn test_output_format_variants() {
let json = OutputFormat::Json;
let jsonlines = OutputFormat::JsonLines;
let yaml = OutputFormat::Yaml;
let csv = OutputFormat::Csv;
assert!(matches!(json, OutputFormat::Json));
assert!(matches!(jsonlines, OutputFormat::JsonLines));
assert!(matches!(yaml, OutputFormat::Yaml));
assert!(matches!(csv, OutputFormat::Csv));
}
#[test]
fn test_generation_result_new() {
let data = vec![json!({"id": 1, "name": "test"})];
let result = GenerationResult::new(data.clone(), 100);
assert_eq!(result.count, 1);
assert_eq!(result.data.len(), 1);
assert_eq!(result.generation_time_ms, 100);
assert_eq!(result.warnings.len(), 0);
}
#[test]
fn test_generation_result_with_warning() {
let data = vec![json!({"id": 1})];
let result = GenerationResult::new(data, 50).with_warning("Test warning".to_string());
assert_eq!(result.warnings.len(), 1);
assert_eq!(result.warnings[0], "Test warning");
}
#[test]
fn test_generation_result_to_json_string() {
let data = vec![json!({"id": 1, "name": "test"})];
let result = GenerationResult::new(data, 10);
let json_string = result.to_json_string();
assert!(json_string.is_ok());
let json_str = json_string.unwrap();
assert!(json_str.contains("\"id\""));
assert!(json_str.contains("\"name\""));
}
#[test]
fn test_generation_result_to_jsonl_string() {
let data = vec![json!({"id": 1}), json!({"id": 2})];
let result = GenerationResult::new(data, 10);
let jsonl_string = result.to_jsonl_string();
assert!(jsonl_string.is_ok());
let jsonl_str = jsonl_string.unwrap();
assert!(jsonl_str.contains("{\"id\":1}"));
assert!(jsonl_str.contains("{\"id\":2}"));
assert!(jsonl_str.contains("\n"));
}
#[test]
fn test_generation_result_multiple_warnings() {
let data = vec![json!({"id": 1})];
let result = GenerationResult::new(data, 10)
.with_warning("Warning 1".to_string())
.with_warning("Warning 2".to_string());
assert_eq!(result.warnings.len(), 2);
assert_eq!(result.warnings[0], "Warning 1");
assert_eq!(result.warnings[1], "Warning 2");
}
#[test]
fn test_default_rows() {
assert_eq!(default_rows(), 100);
}
#[test]
fn test_default_rag_context_length() {
assert_eq!(default_rag_context_length(), 1000);
}
}