unistructgen_codegen/
json_schema.rs

1use serde_json::{json, Map, Value};
2use unistructgen_core::{
3    CodeGenerator, GeneratorMetadata, IRField, IRModule, IRStruct, IRType, IRTypeRef, PrimitiveKind,
4    IREnum,
5};
6use thiserror::Error;
7
8/// Errors that can occur during JSON Schema generation
9#[derive(Debug, Error)]
10pub enum JsonSchemaError {
11    #[error("Serialization error: {0}")]
12    Serialization(#[from] serde_json::Error),
13    #[error("Unsupported type for JSON Schema: {0}")]
14    UnsupportedType(String),
15}
16
17/// Generator for JSON Schema (Draft 2020-12)
18///
19/// This generator produces a JSON Schema compatible with OpenAI's `response_format`
20/// and other AI tools that require structured output validation.
21///
22/// # Example
23///
24/// ```ignore
25/// use unistructgen_codegen::JsonSchemaRenderer;
26/// use unistructgen_core::CodeGenerator;
27///
28/// let generator = JsonSchemaRenderer::default();
29/// let schema_json = generator.generate(&ir_module)?;
30/// ```
31#[derive(Debug, Default)]
32pub struct JsonSchemaRenderer {
33    /// If true, the generated schema will not include the $schema keyword,
34    /// making it suitable for embedding in other schemas.
35    pub fragment_mode: bool,
36}
37
38impl JsonSchemaRenderer {
39    pub fn new() -> Self {
40        Self::default()
41    }
42
43    /// Enable fragment mode (no $schema, no generic wrapper if possible)
44    pub fn fragment(mut self) -> Self {
45        self.fragment_mode = true;
46        self
47    }
48}
49
50impl CodeGenerator for JsonSchemaRenderer {
51    type Error = JsonSchemaError;
52
53    fn generate(&self, module: &IRModule) -> Result<String, Self::Error> {
54        let mut defs = Map::new();
55        let mut root_ref = None;
56
57        // Process all types in the module to build definitions
58        for ty in &module.types {
59            let (name, schema) = match ty {
60                IRType::Struct(s) => (s.name.clone(), render_struct(s)?),
61                IRType::Enum(e) => (e.name.clone(), render_enum(e)?),
62            };
63
64            defs.insert(name.clone(), schema);
65            
66            // Heuristic: The last defined type or a type matching the module name is often the root
67            // For now, we'll just track the last one, but ideally the IRModule would specify a root.
68            root_ref = Some(name);
69        }
70
71        // If the module name matches a type, that's definitely the root
72        if defs.contains_key(&module.name) {
73            root_ref = Some(module.name.clone());
74        }
75
76        let mut root_schema = Map::new();
77        
78        if !self.fragment_mode {
79            root_schema.insert("$schema".to_string(), json!("https://json-schema.org/draft/2020-12/schema"));
80        }
81
82        root_schema.insert("$defs".to_string(), Value::Object(defs));
83
84        // If we identified a root type, point the main schema to it
85        if let Some(root_name) = root_ref {
86            root_schema.insert("$ref".to_string(), json!(format!("#/$defs/{}", root_name)));
87        } else {
88            // Fallback if no types: empty object
89            root_schema.insert("type".to_string(), json!("object"));
90        }
91
92        Ok(serde_json::to_string_pretty(&root_schema)?)
93    }
94
95    fn language(&self) -> &'static str {
96        "JSON Schema"
97    }
98
99    fn file_extension(&self) -> &str {
100        "json"
101    }
102
103    fn metadata(&self) -> GeneratorMetadata {
104        GeneratorMetadata::new()
105            .with_version("0.1.0")
106            .with_description("Generates JSON Schema (Draft 2020-12) for AI Structured Outputs")
107            .with_feature("recursive-types")
108            .with_feature("validation")
109    }
110}
111
112fn render_struct(s: &IRStruct) -> Result<Value, JsonSchemaError> {
113    let mut schema = Map::new();
114    schema.insert("type".to_string(), json!("object"));
115    schema.insert("additionalProperties".to_string(), json!(false));
116    
117    if let Some(doc) = &s.doc {
118        schema.insert("description".to_string(), json!(doc));
119    }
120
121    let mut properties = Map::new();
122    let mut required = Vec::new();
123
124    for field in &s.fields {
125        let field_name = field.source_name.as_ref().unwrap_or(&field.name).clone();
126        let field_schema = render_field(field)?;
127        
128        properties.insert(field_name.clone(), field_schema);
129
130        if !field.optional {
131            required.push(json!(field_name));
132        }
133    }
134
135    schema.insert("properties".to_string(), Value::Object(properties));
136    
137    if !required.is_empty() {
138        schema.insert("required".to_string(), Value::Array(required));
139    }
140
141    Ok(Value::Object(schema))
142}
143
144fn render_enum(e: &IREnum) -> Result<Value, JsonSchemaError> {
145    // Check if it's a simple string enum (no fields in variants)
146    // IR currently doesn't strictly distinguish variant types in the struct def, 
147    // but typical JSON enums are just strings.
148    // If future IR supports complex enums (sum types), this needs 'oneOf'.
149    
150    // Assuming simple string enum for now as per current IR usage
151    let mut schema = Map::new();
152    schema.insert("type".to_string(), json!("string"));
153    
154    if let Some(doc) = &e.doc {
155        schema.insert("description".to_string(), json!(doc));
156    }
157
158    let mut enum_values = Vec::new();
159    for variant in &e.variants {
160        let val = variant.source_value.as_ref().unwrap_or(&variant.name).clone();
161        enum_values.push(json!(val));
162    }
163    
164    schema.insert("enum".to_string(), Value::Array(enum_values));
165
166    Ok(Value::Object(schema))
167}
168
169fn render_field(field: &IRField) -> Result<Value, JsonSchemaError> {
170    let mut schema = render_type_ref(&field.ty)?;
171    
172    // Add field-specific documentation and constraints if it's a simple type object
173    if let Some(obj) = schema.as_object_mut() {
174        if let Some(doc) = &field.doc {
175            obj.insert("description".to_string(), json!(doc));
176        }
177        
178        // Apply constraints
179        if let Some(min) = field.constraints.min_length {
180            obj.insert("minLength".to_string(), json!(min));
181        }
182        if let Some(max) = field.constraints.max_length {
183            obj.insert("maxLength".to_string(), json!(max));
184        }
185        if let Some(min) = field.constraints.min_value {
186            obj.insert("minimum".to_string(), json!(min));
187        }
188        if let Some(max) = field.constraints.max_value {
189            obj.insert("maximum".to_string(), json!(max));
190        }
191        if let Some(pattern) = &field.constraints.pattern {
192            obj.insert("pattern".to_string(), json!(pattern));
193        }
194        if let Some(format) = &field.constraints.format {
195            obj.insert("format".to_string(), json!(format));
196        }
197    }
198
199    Ok(schema)
200}
201
202fn render_type_ref(ty: &IRTypeRef) -> Result<Value, JsonSchemaError> {
203    match ty {
204        IRTypeRef::Primitive(p) => render_primitive(p),
205        IRTypeRef::Named(name) => {
206            Ok(json!({ "$ref": format!("#/$defs/{}", name) }))
207        }
208        IRTypeRef::Option(inner) => {
209            // For JSON schema, optionality is usually handled by 'required'.
210            // However, explicit nullability can be defined.
211            // We'll return the inner type, and let the struct renderer handle 'required'.
212            // If explicit null is needed:
213            // let mut inner_schema = render_type_ref(inner)?;
214            // ... logic to add "null" to type ...
215            render_type_ref(inner)
216        }
217        IRTypeRef::Vec(inner) => {
218            Ok(json!({
219                "type": "array",
220                "items": render_type_ref(inner)?
221            }))
222        }
223        IRTypeRef::Map(_key, value) => {
224            // JSON keys are always strings
225            Ok(json!({
226                "type": "object",
227                "additionalProperties": render_type_ref(value)?
228            }))
229        }
230    }
231}
232
233fn render_primitive(p: &PrimitiveKind) -> Result<Value, JsonSchemaError> {
234    let (ty, format) = match p {
235        PrimitiveKind::String => ("string", None),
236        PrimitiveKind::I32 | PrimitiveKind::I64 | PrimitiveKind::I128 |
237        PrimitiveKind::U32 | PrimitiveKind::U64 | PrimitiveKind::U128 => ("integer", None),
238        PrimitiveKind::I8 | PrimitiveKind::I16 |
239        PrimitiveKind::U8 | PrimitiveKind::U16 => ("integer", None),
240        PrimitiveKind::F32 | PrimitiveKind::F64 | PrimitiveKind::Decimal => ("number", None),
241        PrimitiveKind::Bool => ("boolean", None),
242        PrimitiveKind::Char => ("string", Some("char")), // Custom format
243        PrimitiveKind::DateTime => ("string", Some("date-time")),
244        PrimitiveKind::Uuid => ("string", Some("uuid")),
245        PrimitiveKind::Json => ("object", None), // Generic object
246    };
247
248    let mut map = Map::new();
249    map.insert("type".to_string(), json!(ty));
250    if let Some(fmt) = format {
251        map.insert("format".to_string(), json!(fmt));
252    }
253    
254    Ok(Value::Object(map))
255}
unistructgen_codegen/json_schema.rs

unistructgen_codegen/
json_schema.rs
