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