use std::ffi::CString;
use std::os::raw::c_char;
use std::ptr;
use std::path::PathBuf;
fn get_model_dir() -> PathBuf {
let mut path = std::env::current_dir().unwrap();
path.push("models");
if !path.exists() {
std::fs::create_dir_all(&path).unwrap();
}
path
}
fn train_and_save_edm() -> PathBuf {
use aip::edm::core::{EmotionDataModelTrainer, TrainingDataset, TrainingSample, EmotionState};
use aip::edm::roguelite::training::RogueliteEdmTrainer;
use candle_core::Device;
use std::collections::HashMap;
let device = Device::Cpu;
let mut trainer = RogueliteEdmTrainer::new(device).unwrap();
let samples: Vec<TrainingSample> = (0..100)
.map(|i| {
let mut features = HashMap::new();
for j in 0..15 {
features.insert(j as u32, (i as f32 + j as f32) / 115.0);
}
TrainingSample {
features,
emotion: EmotionState::new(0.5, 0.5, 0.5),
}
})
.collect();
let dataset = TrainingDataset::new(samples);
for _ in 0..5 {
trainer.train_epoch(&dataset).unwrap();
}
let model = trainer.to_model();
let model_dir = get_model_dir();
let model_path = model_dir.join("ffi_test_edm.safetensors");
model.save(&model_path).unwrap();
model_path
}
fn train_and_save_director() -> PathBuf {
use aip::director::core::{InteractionStrategyTrainer, Trajectory, TrajectoryStep, InteractionState, EmotionStats, InteractionParams};
use aip::director::roguelite::training::RogueliteDirectorTrainer;
use aip::edm::core::EmotionState;
use candle_core::Device;
use std::collections::HashMap;
let device = Device::Cpu;
let mut trainer = RogueliteDirectorTrainer::new(device).unwrap();
let trajectories: Vec<Trajectory> = (0..10)
.map(|t| {
let steps: Vec<TrajectoryStep> = (0..10)
.map(|s| {
let mut user_traits = HashMap::new();
for i in 0..8 {
user_traits.insert(i as u32, 0.5);
}
let mut env_state = HashMap::new();
for i in 0..6 {
env_state.insert(i as u32, 0.5);
}
TrajectoryStep {
state: InteractionState {
user_traits,
env_state,
emotion: EmotionState::new(0.5, 0.5, 0.5),
emotion_stats: EmotionStats::default(),
},
action: InteractionParams::default(),
reward: 0.1,
}
})
.collect();
Trajectory { steps }
})
.collect();
for _ in 0..3 {
trainer.train_epoch(&trajectories).unwrap();
}
let model = trainer.to_model();
let model_dir = get_model_dir();
let model_path = model_dir.join("ffi_test_director.safetensors");
model.save(&model_path).unwrap();
model_path
}
#[test]
fn test_ffi_version() {
unsafe {
let version_ptr = aip::ffi::aip_get_version();
assert!(!version_ptr.is_null());
let version = std::ffi::CStr::from_ptr(version_ptr);
assert!(!version.to_bytes().is_empty());
aip::ffi::aip_free_string(version_ptr as *mut c_char);
}
}
#[test]
fn test_ffi_error_messages() {
unsafe {
let msg_ptr = aip::ffi::aip_error_message(aip::ffi::AIP_SUCCESS);
let msg = std::ffi::CStr::from_ptr(msg_ptr);
assert_eq!(msg.to_str().unwrap(), "Success");
aip::ffi::aip_free_string(msg_ptr as *mut c_char);
let msg_ptr = aip::ffi::aip_error_message(aip::ffi::AIP_ERROR_NULL_POINTER);
let msg = std::ffi::CStr::from_ptr(msg_ptr);
assert_eq!(msg.to_str().unwrap(), "Null pointer");
aip::ffi::aip_free_string(msg_ptr as *mut c_char);
}
}
#[test]
fn test_ffi_edm_create_destroy() {
unsafe {
let model = aip::ffi::aip_edm_create();
assert!(!model.is_null());
aip::ffi::aip_edm_destroy(model);
}
}
#[test]
fn test_ffi_director_create_destroy() {
unsafe {
let model = aip::ffi::aip_director_create();
assert!(!model.is_null());
aip::ffi::aip_director_destroy(model);
}
}
#[test]
fn test_ffi_edm_infer() {
unsafe {
let model = aip::ffi::aip_edm_create();
assert!(!model.is_null());
let features: [f32; 15] = [0.5; 15];
let feature_array = aip::ffi::AipFeatureArray {
data: features.as_ptr(),
len: 15,
};
let result = aip::ffi::aip_edm_infer(model, feature_array);
assert_eq!(result.error_code, aip::ffi::AIP_SUCCESS);
assert!(result.valence >= 0.0 && result.valence <= 1.0);
assert!(result.arousal >= 0.0 && result.arousal <= 1.0);
assert!(result.dominance >= 0.0 && result.dominance <= 1.0);
aip::ffi::aip_edm_destroy(model);
}
}
#[test]
fn test_ffi_edm_infer_null_pointer() {
unsafe {
let features: [f32; 15] = [0.5; 15];
let feature_array = aip::ffi::AipFeatureArray {
data: features.as_ptr(),
len: 15,
};
let result = aip::ffi::aip_edm_infer(ptr::null(), feature_array);
assert_eq!(result.error_code, aip::ffi::AIP_ERROR_NULL_POINTER);
}
}
#[test]
fn test_ffi_edm_infer_invalid_data() {
unsafe {
let model = aip::ffi::aip_edm_create();
assert!(!model.is_null());
let features: [f32; 10] = [0.5; 10];
let feature_array = aip::ffi::AipFeatureArray {
data: features.as_ptr(),
len: 10,
};
let result = aip::ffi::aip_edm_infer(model, feature_array);
assert_eq!(result.error_code, aip::ffi::AIP_ERROR_INVALID_DATA);
aip::ffi::aip_edm_destroy(model);
}
}
#[test]
fn test_ffi_director_decide() {
unsafe {
let model = aip::ffi::aip_director_create();
assert!(!model.is_null());
let user_traits: [f32; 8] = [0.5; 8];
let env_state: [f32; 6] = [0.5; 6];
let result = aip::ffi::aip_director_decide(
model,
aip::ffi::AipUserTraits {
data: user_traits.as_ptr(),
len: 8,
},
aip::ffi::AipEnvState {
data: env_state.as_ptr(),
len: 6,
},
aip::ffi::AipEmotionInput {
valence: 0.5,
arousal: 0.5,
dominance: 0.5,
},
);
assert_eq!(result.error_code, aip::ffi::AIP_SUCCESS);
assert!(result.intensity_factor >= 0.5 && result.intensity_factor <= 2.0);
assert!(result.pace_speed >= 0.6 && result.pace_speed <= 1.8);
aip::ffi::aip_director_destroy(model);
}
}
#[test]
fn test_ffi_director_decide_null_pointer() {
unsafe {
let user_traits: [f32; 8] = [0.5; 8];
let env_state: [f32; 6] = [0.5; 6];
let result = aip::ffi::aip_director_decide(
ptr::null(),
aip::ffi::AipUserTraits {
data: user_traits.as_ptr(),
len: 8,
},
aip::ffi::AipEnvState {
data: env_state.as_ptr(),
len: 6,
},
aip::ffi::AipEmotionInput {
valence: 0.5,
arousal: 0.5,
dominance: 0.5,
},
);
assert_eq!(result.error_code, aip::ffi::AIP_ERROR_NULL_POINTER);
}
}
#[test]
fn test_ffi_director_decide_invalid_data() {
unsafe {
let model = aip::ffi::aip_director_create();
assert!(!model.is_null());
let user_traits: [f32; 5] = [0.5; 5];
let env_state: [f32; 6] = [0.5; 6];
let result = aip::ffi::aip_director_decide(
model,
aip::ffi::AipUserTraits {
data: user_traits.as_ptr(),
len: 5,
},
aip::ffi::AipEnvState {
data: env_state.as_ptr(),
len: 6,
},
aip::ffi::AipEmotionInput {
valence: 0.5,
arousal: 0.5,
dominance: 0.5,
},
);
assert_eq!(result.error_code, aip::ffi::AIP_ERROR_INVALID_DATA);
aip::ffi::aip_director_destroy(model);
}
}
#[test]
fn test_ffi_edm_load() {
let model_path = train_and_save_edm();
let path_cstr = CString::new(model_path.to_str().unwrap()).unwrap();
unsafe {
let model = aip::ffi::aip_edm_create();
assert!(!model.is_null());
let result = aip::ffi::aip_edm_load(model, path_cstr.as_ptr());
assert_eq!(result, aip::ffi::AIP_SUCCESS);
let features: [f32; 15] = [0.5; 15];
let feature_array = aip::ffi::AipFeatureArray {
data: features.as_ptr(),
len: 15,
};
let infer_result = aip::ffi::aip_edm_infer(model, feature_array);
assert_eq!(infer_result.error_code, aip::ffi::AIP_SUCCESS);
aip::ffi::aip_edm_destroy(model);
}
}
#[test]
fn test_ffi_edm_load_file_not_found() {
let path_cstr = CString::new("/nonexistent/path/model.safetensors").unwrap();
unsafe {
let model = aip::ffi::aip_edm_create();
assert!(!model.is_null());
let result = aip::ffi::aip_edm_load(model, path_cstr.as_ptr());
assert_eq!(result, aip::ffi::AIP_ERROR_FILE_NOT_FOUND);
aip::ffi::aip_edm_destroy(model);
}
}
#[test]
fn test_ffi_director_load() {
let model_path = train_and_save_director();
let path_cstr = CString::new(model_path.to_str().unwrap()).unwrap();
unsafe {
let model = aip::ffi::aip_director_create();
assert!(!model.is_null());
let result = aip::ffi::aip_director_load(model, path_cstr.as_ptr());
assert_eq!(result, aip::ffi::AIP_SUCCESS);
let user_traits: [f32; 8] = [0.5; 8];
let env_state: [f32; 6] = [0.5; 6];
let decide_result = aip::ffi::aip_director_decide(
model,
aip::ffi::AipUserTraits {
data: user_traits.as_ptr(),
len: 8,
},
aip::ffi::AipEnvState {
data: env_state.as_ptr(),
len: 6,
},
aip::ffi::AipEmotionInput {
valence: 0.5,
arousal: 0.5,
dominance: 0.5,
},
);
assert_eq!(decide_result.error_code, aip::ffi::AIP_SUCCESS);
aip::ffi::aip_director_destroy(model);
}
}
#[test]
fn test_ffi_full_pipeline() {
let edm_path = train_and_save_edm();
let director_path = train_and_save_director();
let edm_path_cstr = CString::new(edm_path.to_str().unwrap()).unwrap();
let director_path_cstr = CString::new(director_path.to_str().unwrap()).unwrap();
unsafe {
let edm = aip::ffi::aip_edm_create();
assert!(!edm.is_null());
let edm_load_result = aip::ffi::aip_edm_load(edm, edm_path_cstr.as_ptr());
assert_eq!(edm_load_result, aip::ffi::AIP_SUCCESS);
let director = aip::ffi::aip_director_create();
assert!(!director.is_null());
let director_load_result = aip::ffi::aip_director_load(director, director_path_cstr.as_ptr());
assert_eq!(director_load_result, aip::ffi::AIP_SUCCESS);
for i in 0..5 {
let mut features: [f32; 15] = [0.5; 15];
features[0] = 0.3 + i as f32 * 0.1;
features[1] = 0.4 + i as f32 * 0.1;
let emotion_result = aip::ffi::aip_edm_infer(
edm,
aip::ffi::AipFeatureArray {
data: features.as_ptr(),
len: 15,
},
);
assert_eq!(emotion_result.error_code, aip::ffi::AIP_SUCCESS);
let user_traits: [f32; 8] = [0.5; 8];
let env_state: [f32; 6] = [0.5; 6];
let decision_result = aip::ffi::aip_director_decide(
director,
aip::ffi::AipUserTraits {
data: user_traits.as_ptr(),
len: 8,
},
aip::ffi::AipEnvState {
data: env_state.as_ptr(),
len: 6,
},
aip::ffi::AipEmotionInput {
valence: emotion_result.valence,
arousal: emotion_result.arousal,
dominance: emotion_result.dominance,
},
);
assert_eq!(decision_result.error_code, aip::ffi::AIP_SUCCESS);
println!(
"Iteration {}: emotion=({:.2}, {:.2}, {:.2}), intensity={:.2}, pace={:.2}",
i,
emotion_result.valence,
emotion_result.arousal,
emotion_result.dominance,
decision_result.intensity_factor,
decision_result.pace_speed
);
}
aip::ffi::aip_edm_destroy(edm);
aip::ffi::aip_director_destroy(director);
}
}