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        // TODO: Gemini docs mention using a subset of OpenAPI 3, so this may benefit from using
29        // `SchemaSettings::openapi3()`.
30        LanguageModelToolSchemaFormat::JsonSchemaSubset => SchemaSettings::draft07()
31            .with(|settings| {
32                settings.meta_schema = None;
33                settings.inline_subschemas = true;
34            })
35            .with_transform(ToJsonSchemaSubsetTransform)
36            .into_generator(),
37    };
38    generator.root_schema_for::<T>()
39}
40
41#[derive(Debug, Clone)]
42struct ToJsonSchemaSubsetTransform;
43
44impl Transform for ToJsonSchemaSubsetTransform {
45    fn transform(&mut self, schema: &mut Schema) {
46        // Ensure that the type field is not an array, this happens when we use
47        // Option<T>, the type will be [T, "null"].
48        if let Some(type_field) = schema.get_mut("type") {
49            if let Some(types) = type_field.as_array() {
50                if let Some(first_type) = types.first() {
51                    *type_field = first_type.clone();
52                }
53            }
54        }
55
56        // oneOf is not supported, use anyOf instead
57        if let Some(one_of) = schema.remove("oneOf") {
58            schema.insert("anyOf".to_string(), one_of);
59        }
60
61        transform_subschemas(self, schema);
62    }
63}