1use anyhow::Result;
 2use language_model::LanguageModelToolSchemaFormat;
 3use schemars::{
 4    JsonSchema,
 5    schema::{RootSchema, Schema, SchemaObject},
 6};
 7
 8pub fn json_schema_for<T: JsonSchema>(
 9    format: LanguageModelToolSchemaFormat,
10) -> Result<serde_json::Value> {
11    let schema = root_schema_for::<T>(format);
12    schema_to_json(&schema, format)
13}
14
15fn schema_to_json(
16    schema: &RootSchema,
17    format: LanguageModelToolSchemaFormat,
18) -> Result<serde_json::Value> {
19    let mut value = serde_json::to_value(schema)?;
20    assistant_tool::adapt_schema_to_format(&mut value, format)?;
21    Ok(value)
22}
23
24fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> RootSchema {
25    let mut generator = match format {
26        LanguageModelToolSchemaFormat::JsonSchema => schemars::SchemaGenerator::default(),
27        LanguageModelToolSchemaFormat::JsonSchemaSubset => {
28            schemars::r#gen::SchemaSettings::default()
29                .with(|settings| {
30                    settings.meta_schema = None;
31                    settings.inline_subschemas = true;
32                    settings
33                        .visitors
34                        .push(Box::new(TransformToJsonSchemaSubsetVisitor));
35                })
36                .into_generator()
37        }
38    };
39    generator.root_schema_for::<T>()
40}
41
42#[derive(Debug, Clone)]
43struct TransformToJsonSchemaSubsetVisitor;
44
45impl schemars::visit::Visitor for TransformToJsonSchemaSubsetVisitor {
46    fn visit_root_schema(&mut self, root: &mut RootSchema) {
47        schemars::visit::visit_root_schema(self, root)
48    }
49
50    fn visit_schema(&mut self, schema: &mut Schema) {
51        schemars::visit::visit_schema(self, schema)
52    }
53
54    fn visit_schema_object(&mut self, schema: &mut SchemaObject) {
55        // Ensure that the type field is not an array, this happens when we use
56        // Option<T>, the type will be [T, "null"].
57        if let Some(instance_type) = schema.instance_type.take() {
58            schema.instance_type = match instance_type {
59                schemars::schema::SingleOrVec::Single(t) => {
60                    Some(schemars::schema::SingleOrVec::Single(t))
61                }
62                schemars::schema::SingleOrVec::Vec(items) => items
63                    .into_iter()
64                    .next()
65                    .map(schemars::schema::SingleOrVec::from),
66            };
67        }
68
69        // One of is not supported, use anyOf instead.
70        if let Some(subschema) = schema.subschemas.as_mut() {
71            if let Some(one_of) = subschema.one_of.take() {
72                subschema.any_of = Some(one_of);
73            }
74        }
75
76        schemars::visit::visit_schema_object(self, schema)
77    }
78}