1use language_model::LanguageModelToolSchemaFormat;
 2use schemars::{
 3    JsonSchema, Schema,
 4    generate::SchemaSettings,
 5    transform::{Transform, transform_subschemas},
 6};
 7
 8pub(crate) fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> Schema {
 9    let mut generator = match format {
10        LanguageModelToolSchemaFormat::JsonSchema => SchemaSettings::draft07().into_generator(),
11        LanguageModelToolSchemaFormat::JsonSchemaSubset => SchemaSettings::openapi3()
12            .with(|settings| {
13                settings.meta_schema = None;
14                settings.inline_subschemas = true;
15            })
16            .with_transform(ToJsonSchemaSubsetTransform)
17            .into_generator(),
18    };
19    generator.root_schema_for::<T>()
20}
21
22#[derive(Debug, Clone)]
23struct ToJsonSchemaSubsetTransform;
24
25impl Transform for ToJsonSchemaSubsetTransform {
26    fn transform(&mut self, schema: &mut Schema) {
27        // Ensure that the type field is not an array, this happens when we use
28        // Option<T>, the type will be [T, "null"].
29        if let Some(type_field) = schema.get_mut("type")
30            && let Some(types) = type_field.as_array()
31            && let Some(first_type) = types.first()
32        {
33            *type_field = first_type.clone();
34        }
35
36        // oneOf is not supported, use anyOf instead
37        if let Some(one_of) = schema.remove("oneOf") {
38            schema.insert("anyOf".to_string(), one_of);
39        }
40
41        transform_subschemas(self, schema);
42    }
43}