Skip to main content

ferrin_schema/
transform.rs

1//! Provider-oriented JSON Schema rewrites.
2
3use serde_json::Map;
4use serde_json::Value;
5use serde_json::json;
6
7/// A rewrite applied to a JSON Schema before it is sent to a provider.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9#[non_exhaustive]
10pub enum SchemaTransform {
11    /// Set `additionalProperties: false` on every object schema.
12    ///
13    /// Applied by default to schemas derived from Rust types, so that models
14    /// cannot invent properties the type does not declare.
15    AdditionalPropertiesFalse,
16    /// Remove `propertyNames` keywords everywhere.
17    RemovePropertyNames,
18    /// OpenAI strict mode: `additionalProperties: false`, no `propertyNames`,
19    /// every property required, previously optional properties nullable.
20    OpenAiStrict,
21}
22
23impl SchemaTransform {
24    /// The OpenAI strict-mode transform.
25    #[must_use]
26    pub fn openai_strict() -> Self {
27        Self::OpenAiStrict
28    }
29
30    /// The `additionalProperties: false` transform.
31    #[must_use]
32    pub fn additional_properties_false() -> Self {
33        Self::AdditionalPropertiesFalse
34    }
35
36    /// Applies the transform in place.
37    pub fn apply(self, schema: &mut Value) {
38        match self {
39            Self::AdditionalPropertiesFalse => add_additional_properties_false(schema),
40            Self::RemovePropertyNames => remove_property_names(schema),
41            Self::OpenAiStrict => to_openai_strict(schema),
42        }
43    }
44
45    /// Applies the transform to a copy and returns it.
46    #[must_use]
47    pub fn applied(self, mut schema: Value) -> Value {
48        self.apply(&mut schema);
49        schema
50    }
51}
52
53/// Sets `additionalProperties: false` on every object schema that does not
54/// already declare a schema for additional properties.
55///
56/// Objects are recognized by `type: "object"` (or a type array containing
57/// `object`). Recurses into `properties`, `additionalProperties`, `items`,
58/// `anyOf`, `allOf`, `oneOf` and `definitions`/`$defs`.
59pub fn add_additional_properties_false(schema: &mut Value) {
60    let Value::Object(obj) = schema else { return };
61    if type_includes(obj, "object") {
62        match obj.get_mut("additionalProperties") {
63            Some(nested @ Value::Object(_)) => add_additional_properties_false(nested),
64            _ => {
65                obj.insert("additionalProperties".to_owned(), Value::Bool(false));
66            }
67        }
68        if let Some(Value::Object(properties)) = obj.get_mut("properties") {
69            for property in properties.values_mut() {
70                add_additional_properties_false(property);
71            }
72        }
73    }
74    for_each_subschema(obj, &mut add_additional_properties_false);
75}
76
77/// Removes `propertyNames` from every schema.
78pub fn remove_property_names(schema: &mut Value) {
79    let Value::Object(obj) = schema else { return };
80    obj.remove("propertyNames");
81    if let Some(Value::Object(properties)) = obj.get_mut("properties") {
82        for property in properties.values_mut() {
83            remove_property_names(property);
84        }
85    }
86    if let Some(nested @ Value::Object(_)) = obj.get_mut("additionalProperties") {
87        remove_property_names(nested);
88    }
89    for_each_subschema(obj, &mut remove_property_names);
90}
91
92/// Rewrites a schema for OpenAI strict mode.
93///
94/// Every object gets `additionalProperties: false`, `propertyNames` is
95/// removed, every property is listed in `required`, and properties that were
96/// optional become nullable (`type` gains `"null"`, or the schema is wrapped
97/// in `anyOf: [.., {type: null}]`). Objects are recognized by `type` or by
98/// the presence of `properties`.
99pub fn to_openai_strict(schema: &mut Value) {
100    let Value::Object(obj) = schema else { return };
101    obj.remove("propertyNames");
102    let is_object = match obj.get("type") {
103        Some(Value::String(kind)) => kind == "object",
104        Some(Value::Array(kinds)) => kinds.iter().any(|kind| kind == "object"),
105        _ => obj.contains_key("properties"),
106    };
107    if is_object {
108        let required: Vec<String> = obj
109            .get("required")
110            .and_then(Value::as_array)
111            .map(|items| {
112                items
113                    .iter()
114                    .filter_map(Value::as_str)
115                    .map(str::to_owned)
116                    .collect()
117            })
118            .unwrap_or_default();
119        if let Some(Value::Object(properties)) = obj.get_mut("properties") {
120            let names: Vec<String> = properties.keys().cloned().collect();
121            for (name, property) in properties.iter_mut() {
122                if !required.contains(name) {
123                    make_nullable(property);
124                }
125                to_openai_strict(property);
126            }
127            obj.insert("required".to_owned(), json!(names));
128        }
129        match obj.get_mut("additionalProperties") {
130            Some(nested @ Value::Object(_)) => to_openai_strict(nested),
131            _ => {
132                obj.insert("additionalProperties".to_owned(), Value::Bool(false));
133            }
134        }
135    }
136    for key in ["not", "if", "then", "else", "contains"] {
137        if let Some(child) = obj.get_mut(key) {
138            to_openai_strict(child);
139        }
140    }
141    if let Some(Value::Array(items)) = obj.get_mut("prefixItems") {
142        items.iter_mut().for_each(to_openai_strict);
143    }
144    if let Some(Value::Object(map)) = obj.get_mut("patternProperties") {
145        map.values_mut().for_each(to_openai_strict);
146    }
147    for_each_subschema(obj, &mut to_openai_strict);
148}
149
150/// Makes a property schema accept `null`.
151fn make_nullable(property: &mut Value) {
152    let Value::Object(obj) = property else { return };
153    match obj.get_mut("type") {
154        Some(Value::String(kind)) => {
155            let kind = kind.clone();
156            obj.insert("type".to_owned(), json!([kind, "null"]));
157        }
158        Some(Value::Array(kinds)) => {
159            if !kinds.iter().any(|kind| kind == "null") {
160                kinds.push(json!("null"));
161            }
162        }
163        _ => {
164            let inner = Value::Object(std::mem::take(obj));
165            let mut wrapper = Map::new();
166            wrapper.insert("anyOf".to_owned(), json!([inner, { "type": "null" }]));
167            *obj = wrapper;
168        }
169    }
170}
171
172fn type_includes(obj: &Map<String, Value>, kind: &str) -> bool {
173    match obj.get("type") {
174        Some(Value::String(actual)) => actual == kind,
175        Some(Value::Array(kinds)) => kinds.iter().any(|actual| actual == kind),
176        _ => false,
177    }
178}
179
180/// Visits `items`, `anyOf`, `allOf`, `oneOf`, `definitions` and `$defs`.
181fn for_each_subschema(obj: &mut Map<String, Value>, visit: &mut dyn FnMut(&mut Value)) {
182    if let Some(items) = obj.get_mut("items") {
183        match items {
184            Value::Array(list) => list.iter_mut().for_each(&mut *visit),
185            other => visit(other),
186        }
187    }
188    for key in ["anyOf", "allOf", "oneOf"] {
189        if let Some(Value::Array(list)) = obj.get_mut(key) {
190            list.iter_mut().for_each(&mut *visit);
191        }
192    }
193    for key in ["definitions", "$defs"] {
194        if let Some(Value::Object(map)) = obj.get_mut(key) {
195            map.values_mut().for_each(&mut *visit);
196        }
197    }
198}