pub mod client;
pub mod diff;
pub mod prompts;
use crate::schema::{Schema, SchemaError};
#[derive(Debug)]
pub enum GenerateError {
MissingApiKey,
Transport(String),
Schema(SchemaError),
EmptyResult,
}
impl std::fmt::Display for GenerateError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingApiKey => f.write_str(
"ANTHROPIC_API_KEY is not set. Set it in your environment before running \
`rustio ai generate`.",
),
Self::Transport(msg) => write!(f, "anthropic API transport error: {msg}"),
Self::Schema(err) => write!(f, "anthropic API returned invalid schema: {err}"),
Self::EmptyResult => f.write_str(
"Refusing to apply update: schema would become empty",
),
}
}
}
impl std::error::Error for GenerateError {}
impl From<SchemaError> for GenerateError {
fn from(e: SchemaError) -> Self {
GenerateError::Schema(e)
}
}
pub async fn generate(prompt: &str) -> Result<Schema, GenerateError> {
let api_key = api_key()?;
let body = client::request(&api_key, prompt)
.await
.map_err(|e| GenerateError::Transport(e.to_string()))?;
parse_response(&body)
}
pub async fn update(existing: &Schema, instruction: &str) -> Result<Schema, GenerateError> {
let api_key = api_key()?;
let existing_json = existing
.to_pretty_json()
.map_err(|e| GenerateError::Transport(format!("serialise existing schema: {e}")))?;
let body = client::request_update(&api_key, &existing_json, instruction)
.await
.map_err(|e| GenerateError::Transport(e.to_string()))?;
let updated = parse_response(&body)?;
check_not_empty(existing, &updated)?;
Ok(updated)
}
pub(crate) fn check_not_empty(old: &Schema, new: &Schema) -> Result<(), GenerateError> {
if new.models.is_empty() && !old.models.is_empty() {
return Err(GenerateError::EmptyResult);
}
Ok(())
}
fn api_key() -> Result<String, GenerateError> {
std::env::var("ANTHROPIC_API_KEY")
.ok()
.filter(|s| !s.trim().is_empty())
.ok_or(GenerateError::MissingApiKey)
}
#[derive(Debug, Clone, PartialEq)]
pub struct AnalyzeReport {
pub issues: Vec<String>,
pub suggestions: Vec<String>,
pub score: f32,
}
#[derive(Debug)]
pub enum AnalyzeError {
MissingApiKey,
Transport(String),
Encode(String),
}
impl std::fmt::Display for AnalyzeError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingApiKey => f.write_str(
"ANTHROPIC_API_KEY is not set. Set it in your environment before running \
`rustio ai analyze`.",
),
Self::Transport(msg) => write!(f, "anthropic API transport error: {msg}"),
Self::Encode(msg) => write!(f, "could not serialise schema for analyze: {msg}"),
}
}
}
impl std::error::Error for AnalyzeError {}
pub async fn analyze(schema: &Schema) -> Result<AnalyzeReport, AnalyzeError> {
let api_key = std::env::var("ANTHROPIC_API_KEY")
.ok()
.filter(|s| !s.trim().is_empty())
.ok_or(AnalyzeError::MissingApiKey)?;
let existing_json = schema
.to_pretty_json()
.map_err(|e| AnalyzeError::Encode(e.to_string()))?;
let body = client::request_analyze(&api_key, &existing_json)
.await
.map_err(AnalyzeError::Transport)?;
Ok(parse_analyze_response(&body))
}
pub fn parse_analyze_response(body: &str) -> AnalyzeReport {
let body = body.trim();
if body.is_empty() {
return AnalyzeReport {
issues: Vec::new(),
suggestions: Vec::new(),
score: 0.0,
};
}
let lower = body.to_lowercase();
let has_section_header = lower.contains("issues:")
|| lower.contains("suggestions:")
|| lower.contains("score:");
if !has_section_header {
let suggestions = collect_bullets(body);
return AnalyzeReport {
issues: Vec::new(),
suggestions,
score: 0.0,
};
}
let mut section = Section::None;
let mut issues: Vec<String> = Vec::new();
let mut suggestions: Vec<String> = Vec::new();
let mut score: f32 = 0.0;
for raw_line in body.lines() {
let line = raw_line.trim();
let lower = line.to_lowercase();
if lower.starts_with("issues:") {
section = Section::Issues;
continue;
}
if lower.starts_with("suggestions:") {
section = Section::Suggestions;
continue;
}
if lower.starts_with("score:") {
section = Section::Score;
score = parse_score(line["score:".len()..].trim()).unwrap_or(0.0);
continue;
}
if line.is_empty() || line.eq_ignore_ascii_case("(none)") {
continue;
}
let item = line
.strip_prefix("- ")
.or_else(|| line.strip_prefix("* "))
.unwrap_or(line)
.to_string();
match section {
Section::Issues => issues.push(item),
Section::Suggestions => suggestions.push(item),
Section::Score | Section::None => {}
}
}
AnalyzeReport { issues, suggestions, score }
}
enum Section {
None,
Issues,
Suggestions,
Score,
}
fn parse_score(s: &str) -> Option<f32> {
let s = s.trim();
let end = s
.char_indices()
.find(|(_, c)| !(c.is_ascii_digit() || *c == '.' || *c == '-'))
.map(|(i, _)| i)
.unwrap_or(s.len());
s[..end].parse::<f32>().ok()
}
fn collect_bullets(body: &str) -> Vec<String> {
body.lines()
.map(str::trim)
.filter(|l| !l.is_empty())
.map(|l| {
l.strip_prefix("- ")
.or_else(|| l.strip_prefix("* "))
.unwrap_or(l)
.to_string()
})
.collect()
}
#[derive(Debug, Clone, PartialEq)]
pub struct ExplainReport {
pub why: Vec<String>,
pub impact: Vec<String>,
}
#[derive(Debug)]
pub enum ExplainError {
MissingApiKey,
Transport(String),
Encode(String),
}
impl std::fmt::Display for ExplainError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::MissingApiKey => f.write_str(
"ANTHROPIC_API_KEY is not set. Set it in your environment before requesting \
an explanation.",
),
Self::Transport(msg) => write!(f, "anthropic API transport error: {msg}"),
Self::Encode(msg) => write!(f, "could not serialise schema for explain: {msg}"),
}
}
}
impl std::error::Error for ExplainError {}
pub async fn explain_diff(
old: &Schema,
new: &Schema,
api_key: &str,
) -> Result<ExplainReport, ExplainError> {
if api_key.trim().is_empty() {
return Err(ExplainError::MissingApiKey);
}
let old_json = old
.to_pretty_json()
.map_err(|e| ExplainError::Encode(e.to_string()))?;
let new_json = new
.to_pretty_json()
.map_err(|e| ExplainError::Encode(e.to_string()))?;
let body = client::request_explain(api_key, &old_json, &new_json)
.await
.map_err(ExplainError::Transport)?;
Ok(parse_explain_response(&body))
}
pub fn parse_explain_response(body: &str) -> ExplainReport {
let body = body.trim();
if body.is_empty() {
return ExplainReport { why: Vec::new(), impact: Vec::new() };
}
let lower = body.to_lowercase();
let has_section_header = lower.contains("why:") || lower.contains("impact:");
if !has_section_header {
return ExplainReport { why: collect_bullets(body), impact: Vec::new() };
}
let mut section = ExplainSection::None;
let mut why: Vec<String> = Vec::new();
let mut impact: Vec<String> = Vec::new();
for raw_line in body.lines() {
let line = raw_line.trim();
let lower = line.to_lowercase();
if lower.starts_with("why:") {
section = ExplainSection::Why;
continue;
}
if lower.starts_with("impact:") {
section = ExplainSection::Impact;
continue;
}
if line.is_empty() || line.eq_ignore_ascii_case("(none)") {
continue;
}
let bullet = line
.strip_prefix("- ")
.or_else(|| line.strip_prefix("* "));
match (§ion, bullet) {
(ExplainSection::Why, Some(item)) => why.push(item.to_string()),
(ExplainSection::Impact, Some(item)) => impact.push(item.to_string()),
(ExplainSection::Why | ExplainSection::Impact, None) => {
section = ExplainSection::None;
}
(ExplainSection::None, _) => {}
}
}
ExplainReport { why, impact }
}
enum ExplainSection {
None,
Why,
Impact,
}
pub fn parse_response(body: &str) -> Result<Schema, GenerateError> {
let json = extract_schema_json(body);
Ok(Schema::parse(json)?)
}
pub(crate) fn extract_schema_json(body: &str) -> &str {
let trimmed = body.trim();
let stripped = trimmed
.strip_prefix("```json")
.or_else(|| trimmed.strip_prefix("```"))
.unwrap_or(trimmed);
let stripped = stripped.trim_start_matches('\n');
stripped.strip_suffix("```").map_or(stripped, str::trim_end)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn extract_schema_json_strips_fence() {
let fenced = "```json\n{\"version\":2}\n```";
assert_eq!(extract_schema_json(fenced), "{\"version\":2}");
let fenced_no_lang = "```\n{\"version\":2}\n```";
assert_eq!(extract_schema_json(fenced_no_lang), "{\"version\":2}");
let plain = " {\"version\":2} ";
assert_eq!(extract_schema_json(plain), "{\"version\":2}");
}
#[test]
fn parse_response_accepts_valid_schema() {
let body = r#"{
"version": 2,
"rustio_version": "1.0.0",
"models": [
{
"name": "Post",
"table": "posts",
"admin_name": "posts",
"display_name": "Posts",
"singular_name": "Post",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true },
{ "name": "title", "type": "String", "nullable": false, "editable": true }
],
"relations": []
}
]
}"#;
let schema = parse_response(body).expect("valid response parses");
assert_eq!(schema.models.len(), 1);
assert_eq!(schema.models[0].name, "Post");
}
#[test]
fn update_adds_new_model() {
let response = r#"{
"version": 2,
"rustio_version": "1.0.0",
"models": [
{
"name": "Post",
"table": "posts",
"admin_name": "posts",
"display_name": "Posts",
"singular_name": "Post",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true },
{ "name": "title", "type": "String", "nullable": false, "editable": true }
],
"relations": []
},
{
"name": "Tag",
"table": "tags",
"admin_name": "tags",
"display_name": "Tags",
"singular_name": "Tag",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true },
{ "name": "label", "type": "String", "nullable": false, "editable": true }
],
"relations": []
}
]
}"#;
let updated = parse_response(response).expect("valid update parses");
assert!(updated.models.iter().any(|m| m.name == "Tag"));
assert!(updated.models.iter().any(|m| m.name == "Post"));
}
#[test]
fn update_preserves_existing_fields() {
let original = r#"{
"version": 2,
"rustio_version": "1.0.0",
"models": [
{
"name": "Post",
"table": "posts",
"admin_name": "posts",
"display_name": "Posts",
"singular_name": "Post",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true },
{ "name": "title", "type": "String", "nullable": false, "editable": true },
{ "name": "body", "type": "String", "nullable": false, "editable": true }
],
"relations": []
}
]
}"#;
let response = r#"{
"version": 2,
"rustio_version": "1.0.0",
"models": [
{
"name": "Post",
"table": "posts",
"admin_name": "posts",
"display_name": "Posts",
"singular_name": "Post",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true },
{ "name": "title", "type": "String", "nullable": false, "editable": true },
{ "name": "body", "type": "String", "nullable": false, "editable": true },
{ "name": "status", "type": "String", "nullable": false, "editable": true }
],
"relations": []
}
]
}"#;
let old = parse_response(original).expect("original parses");
let new = parse_response(response).expect("response parses");
let changes = diff::diff(&old, &new);
for surviving in ["id", "title", "body"] {
assert!(
!changes.iter().any(|c| matches!(c,
diff::Change::FieldRemoved { field, .. } if field == surviving
)),
"preserved field {surviving} surfaced as removed: {changes:?}"
);
}
let adds: Vec<_> = changes
.iter()
.filter(|c| matches!(c, diff::Change::FieldAdded { .. }))
.collect();
assert_eq!(adds.len(), 1);
}
#[test]
fn update_invalid_json_rejected() {
let bad = r#"{
"version": 2,
"rustio_version": "1.0.0",
"models": [],
}"#;
let err = parse_response(bad).expect_err("malformed JSON must be rejected");
assert!(matches!(err, GenerateError::Schema(_)));
}
#[test]
fn no_live_api_calls() {
let _ = std::env::var("ANTHROPIC_API_KEY"); let dummy = r#"{
"version": 2, "rustio_version": "1.0.0",
"models": [
{ "name": "Post", "table": "posts", "admin_name": "posts",
"display_name": "Posts", "singular_name": "Post",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true }
],
"relations": []
}
]
}"#;
let parsed = parse_response(dummy).expect("offline parse path works");
let _ = diff::diff(&parsed, &parsed); }
#[test]
fn analyze_detects_missing_relation_model() {
let body = "ISSUES:\n\
- Post.author_id has relation but User model missing\n\
\n\
SUGGESTIONS:\n\
- Add created_at timestamp to all models\n\
\n\
SCORE: 6.0\n";
let report = parse_analyze_response(body);
assert_eq!(report.issues.len(), 1);
assert!(report.issues[0].contains("author_id"));
assert!(report.issues[0].contains("User"));
assert_eq!(report.suggestions.len(), 1);
assert!((report.score - 6.0).abs() < f32::EPSILON);
}
#[test]
fn analyze_suggests_best_practices() {
let body = "ISSUES:\n\
(none)\n\
\n\
SUGGESTIONS:\n\
- Add created_at and updated_at to every model\n\
- Index Comment.post_id\n\
- Consider an enum for Post.status\n\
\n\
SCORE: 8.5\n";
let report = parse_analyze_response(body);
assert!(report.issues.is_empty(), "issues bucket should be empty");
assert_eq!(report.suggestions.len(), 3);
assert!(report.suggestions.iter().any(|s| s.contains("created_at")));
assert!(report.suggestions.iter().any(|s| s.contains("Index")));
assert!(report.suggestions.iter().any(|s| s.contains("enum")));
assert!((report.score - 8.5).abs() < f32::EPSILON);
}
#[test]
fn analyze_parsing_valid_output() {
let body = "ISSUES:\n\
- Post.author_id has relation but User model missing\n\
- Comment.post_id not indexed\n\
\n\
SUGGESTIONS:\n\
- Add created_at timestamp to all models\n\
- Add index on foreign keys\n\
\n\
SCORE: 7.5 / 10\n";
let report = parse_analyze_response(body);
assert_eq!(report.issues.len(), 2);
assert_eq!(report.suggestions.len(), 2);
assert!((report.score - 7.5).abs() < f32::EPSILON);
}
#[test]
fn analyze_handles_unstructured_output() {
let body = "Looks fine overall. Maybe think about adding indexes \n\
on the foreign keys, and consider an enum for Post.status.\n\
- Add created_at on every model.";
let report = parse_analyze_response(body);
assert!(report.issues.is_empty(), "unstructured input → issues must be empty");
assert!(
!report.suggestions.is_empty(),
"unstructured input → fallback should populate suggestions"
);
assert_eq!(report.suggestions.len(), 3);
assert!(report.suggestions[2].starts_with("Add created_at"));
assert_eq!(report.score, 0.0, "no SCORE: header → default 0.0");
}
#[test]
fn analyze_no_live_api_calls() {
let _ = std::env::var("ANTHROPIC_API_KEY"); let report = parse_analyze_response(
"ISSUES:\n(none)\n\nSUGGESTIONS:\n(none)\n\nSCORE: 9\n",
);
assert_eq!(report.issues.len(), 0);
assert_eq!(report.suggestions.len(), 0);
assert_eq!(report.score, 9.0);
}
#[test]
fn explain_parses_valid_response() {
let body = "WHY:\n\
- Tags allow flexible categorization of posts\n\
- Decoupling from rigid categories\n\
\n\
IMPACT:\n\
- Adds new table (Tag)\n\
- Introduces many-to-many relationship\n";
let report = parse_explain_response(body);
assert_eq!(report.why.len(), 2);
assert!(report.why[0].starts_with("Tags allow"));
assert!(report.why[1].starts_with("Decoupling"));
assert_eq!(report.impact.len(), 2);
assert!(report.impact[0].starts_with("Adds new table"));
assert!(report.impact[1].starts_with("Introduces"));
}
#[test]
fn explain_handles_missing_sections() {
let body = "IMPACT:\n- Adds Tag table\n";
let report = parse_explain_response(body);
assert!(report.why.is_empty(), "WHY missing → empty bucket");
assert_eq!(report.impact.len(), 1);
let body = "WHY:\n- Tags help categorize\n";
let report = parse_explain_response(body);
assert_eq!(report.why.len(), 1);
assert!(report.impact.is_empty(), "IMPACT missing → empty bucket");
let body = "WHY:\n(none)\n\nIMPACT:\n(none)\n";
let report = parse_explain_response(body);
assert!(report.why.is_empty());
assert!(report.impact.is_empty());
}
#[test]
fn explain_ignores_extra_text() {
let body = "WHY:\n\
- Improves categorization\n\
\n\
IMPACT:\n\
- New table\n\
\n\
This concludes the explanation. Hope it helps!\n";
let report = parse_explain_response(body);
assert_eq!(report.why.len(), 1);
assert_eq!(report.impact.len(), 1);
assert!(report.impact[0].starts_with("New table"));
assert!(
!report.impact.iter().any(|l| l.contains("This concludes")),
"trailing commentary leaked into impact: {:?}",
report.impact
);
assert!(
!report.why.iter().any(|l| l.contains("This concludes")),
"trailing commentary leaked into why: {:?}",
report.why
);
}
#[test]
fn explain_fallback_treats_unstructured_as_why() {
let body = "Tags help categorize posts.\n\
- New table is added.\n\
- Many-to-many relationship is introduced.";
let report = parse_explain_response(body);
assert!(
report.impact.is_empty(),
"no headers → impact must be empty"
);
assert_eq!(report.why.len(), 3);
assert_eq!(report.why[0], "Tags help categorize posts.");
assert_eq!(report.why[1], "New table is added.");
}
#[test]
fn update_refuses_empty_result() {
let one_model = crate::schema::Schema {
version: crate::schema::SCHEMA_VERSION,
rustio_version: "1.0.0".into(),
models: vec![crate::schema::SchemaModel {
name: "Post".into(),
table: "posts".into(),
admin_name: "posts".into(),
display_name: "Posts".into(),
singular_name: "Post".into(),
fields: vec![],
relations: vec![],
core: false,
}],
};
let empty = crate::schema::Schema {
version: crate::schema::SCHEMA_VERSION,
rustio_version: "1.0.0".into(),
models: vec![],
};
let err = check_not_empty(&one_model, &empty)
.expect_err("non-empty → empty must reject");
assert!(
matches!(err, GenerateError::EmptyResult),
"expected EmptyResult, got {err:?}"
);
assert_eq!(
err.to_string(),
"Refusing to apply update: schema would become empty"
);
check_not_empty(&one_model, &one_model)
.expect("non-empty preservation must pass");
check_not_empty(&empty, &empty).expect("empty no-op must pass");
check_not_empty(&empty, &one_model)
.expect("first-time fill must pass");
}
#[test]
fn parse_response_rejects_invalid_schema() {
let body = r#"{
"version": 2,
"rustio_version": "1.0.0",
"models": [
{
"name": "Post",
"table": "posts",
"admin_name": "posts",
"display_name": "Posts",
"singular_name": "Post",
"fields": [
{ "name": "id", "type": "i64", "nullable": false, "editable": true },
{ "name": "title", "type": "FooBar", "nullable": false, "editable": true }
],
"relations": []
}
]
}"#;
let err = parse_response(body).expect_err("invalid type must reject");
match err {
GenerateError::Schema(SchemaError::InvalidType { ref ty, .. }) => {
assert_eq!(ty, "FooBar");
}
other => panic!("expected Schema(InvalidType), got {other:?}"),
}
}
}