pub mod distance;
pub mod error;
pub mod repair;
pub mod sanitize;
pub mod schema;
pub use distance::{Algorithm, Match};
pub use error::FuzzyError;
pub use repair::{
repair_enum_array, repair_fields_with_list, repair_object_fields, repair_tagged_enum,
repair_tagged_enum_array, repair_tagged_enum_json, Correction, FuzzyOptions, RepairResult,
};
pub use sanitize::sanitize_json;
pub use schema::{ObjectSchema, TaggedEnumSchema};
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_full_workflow() {
let schema =
TaggedEnumSchema::new("type", &["AddDerive", "RenameIdent"], |tag| match tag {
"AddDerive" => Some(&["target", "derives"][..]),
"RenameIdent" => Some(&["from", "to", "kind"][..]),
_ => None,
});
let llm_output = r#"{
"type": "AddDeriv",
"taget": "DatabaseConfig",
"derives": ["Debug", "Clone"]
}"#;
let result =
repair_tagged_enum_json(llm_output, &schema, &FuzzyOptions::default()).unwrap();
assert_eq!(result.repaired["type"], "AddDerive");
assert_eq!(result.repaired["target"], "DatabaseConfig");
assert_eq!(result.corrections.len(), 2);
}
#[test]
fn test_sanitize_then_repair_workflow() {
let schema =
TaggedEnumSchema::new("type", &["AddDerive"], |_| Some(&["target", "derives"][..]))
.with_enum_array("derives", &["Debug", "Clone", "Serialize"]);
let malformed_llm_output = r#"{
"type": "AddDeriv",
"taget": "User",
"derives": ["Debg", "Clne",],
}"#;
let sanitized = sanitize_json(malformed_llm_output);
assert!(!sanitized.contains(",]"));
assert!(!sanitized.contains(",}"));
let result =
repair_tagged_enum_json(&sanitized, &schema, &FuzzyOptions::default()).unwrap();
assert_eq!(result.repaired["type"], "AddDerive");
assert!(result.repaired.get("target").is_some());
assert_eq!(result.repaired["derives"][0], "Debug");
assert_eq!(result.repaired["derives"][1], "Clone");
}
#[test]
fn test_sanitize_then_repair_missing_brace() {
let schema = TaggedEnumSchema::new("type", &["Action"], |_| Some(&["name"][..]));
let truncated = r#"{"type": "Action", "name": "test"#;
let sanitized = sanitize_json(truncated);
let result =
repair_tagged_enum_json(&sanitized, &schema, &FuzzyOptions::default()).unwrap();
assert_eq!(result.repaired["type"], "Action");
assert_eq!(result.repaired["name"], "test");
}
}