ai_lib_rust/structured/
schema.rs

1//! Schema generation utilities.
2
3use serde_json::json;
4
5/// Generator for JSON schemas with customization options.
6#[derive(Debug, Clone, Default)]
7pub struct SchemaGenerator {
8    title: Option<String>,
9    description: Option<String>,
10    properties: Vec<(String, serde_json::Value)>,
11    required: Vec<String>,
12    additional_properties: bool,
13}
14
15impl SchemaGenerator {
16    pub fn new() -> Self {
17        Self::default()
18    }
19
20    pub fn title(mut self, title: impl Into<String>) -> Self {
21        self.title = Some(title.into());
22        self
23    }
24
25    pub fn description(mut self, description: impl Into<String>) -> Self {
26        self.description = Some(description.into());
27        self
28    }
29
30    pub fn add_property(mut self, name: impl Into<String>, schema: serde_json::Value) -> Self {
31        self.properties.push((name.into(), schema));
32        self
33    }
34
35    pub fn set_required(mut self, required: &[String]) -> Self {
36        self.required = required.to_vec();
37        self
38    }
39
40    pub fn set_additional_properties(mut self, additional: bool) -> Self {
41        self.additional_properties = additional;
42        self
43    }
44
45    pub fn build(self) -> serde_json::Value {
46        let mut map = serde_json::Map::new();
47        map.insert("type".into(), json!("object"));
48
49        let mut properties = serde_json::Map::new();
50        for (name, schema) in self.properties {
51            properties.insert(name, schema);
52        }
53        map.insert("properties".into(), properties.into());
54
55        if !self.required.is_empty() {
56            map.insert("required".into(), self.required.into());
57        }
58
59        if !self.additional_properties {
60            map.insert("additionalProperties".into(), json!(false));
61        }
62
63        if let Some(title) = self.title {
64            map.insert("title".into(), title.into());
65        }
66        if let Some(desc) = self.description {
67            map.insert("description".into(), desc.into());
68        }
69
70        map.into()
71    }
72}
73
74pub fn schema_from_type_name(type_name: &str) -> serde_json::Value {
75    match type_name {
76        "string" => json!({"type": "string"}),
77        "integer" => json!({"type": "integer"}),
78        "number" => json!({"type": "number"}),
79        "boolean" => json!({"type": "boolean"}),
80        "array" => json!({"type": "array"}),
81        "object" => json!({"type": "object"}),
82        "null" => json!({"type": "null"}),
83        _ => json!({"type": "object"}),
84    }
85}
86
87pub fn json_schema_from_type<T: schemars::JsonSchema>() -> serde_json::Value {
88    let schema = schemars::schema_for!(T);
89    serde_json::to_value(&schema).unwrap_or_else(|_| json!({}))
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn test_schema_generator_basic() {
98        let generator = SchemaGenerator::new()
99            .add_property("name", json!({"type": "string"}))
100            .add_property("age", json!({"type": "integer"}));
101
102        let schema = generator.build();
103        assert_eq!(schema["type"], "object");
104        assert_eq!(schema["properties"]["name"]["type"], "string");
105        assert_eq!(schema["properties"]["age"]["type"], "integer");
106    }
107
108    #[test]
109    fn test_schema_generator_with_required() {
110        let generator = SchemaGenerator::new()
111            .add_property("name", json!({"type": "string"}))
112            .set_required(&vec!["name".to_string()]);
113
114        let schema = generator.build();
115        assert!(schema["required"].is_array());
116        assert_eq!(schema["required"][0], "name");
117    }
118
119    #[test]
120    fn test_schema_generator_additional_properties_false() {
121        let generator = SchemaGenerator::new().set_additional_properties(false);
122        let schema = generator.build();
123        assert_eq!(schema["additionalProperties"], false);
124    }
125
126    #[test]
127    fn test_schema_from_type_name() {
128        let string_schema = schema_from_type_name("string");
129        assert_eq!(string_schema["type"], "string");
130
131        let integer_schema = schema_from_type_name("integer");
132        assert_eq!(integer_schema["type"], "integer");
133
134        let array_schema = schema_from_type_name("array");
135        assert_eq!(array_schema["type"], "array");
136    }
137}
ai_lib_rust/structured/schema.rs

ai_lib_rust/structured/
schema.rs
