3d7e012e09
Closes #33894 After #33635 Gemini Integration was broken because we now produce `const` fields for enums, which are not supported. Changing this to `openapi3` fixes the issue. Release Notes: - Fixed an issue where Gemini Models would not work because of incompatible tool schemas
61 lines
2 KiB
Rust
61 lines
2 KiB
Rust
use anyhow::Result;
|
|
use language_model::LanguageModelToolSchemaFormat;
|
|
use schemars::{
|
|
JsonSchema, Schema,
|
|
generate::SchemaSettings,
|
|
transform::{Transform, transform_subschemas},
|
|
};
|
|
|
|
pub fn json_schema_for<T: JsonSchema>(
|
|
format: LanguageModelToolSchemaFormat,
|
|
) -> Result<serde_json::Value> {
|
|
let schema = root_schema_for::<T>(format);
|
|
schema_to_json(&schema, format)
|
|
}
|
|
|
|
fn schema_to_json(
|
|
schema: &Schema,
|
|
format: LanguageModelToolSchemaFormat,
|
|
) -> Result<serde_json::Value> {
|
|
let mut value = serde_json::to_value(schema)?;
|
|
assistant_tool::adapt_schema_to_format(&mut value, format)?;
|
|
Ok(value)
|
|
}
|
|
|
|
fn root_schema_for<T: JsonSchema>(format: LanguageModelToolSchemaFormat) -> Schema {
|
|
let mut generator = match format {
|
|
LanguageModelToolSchemaFormat::JsonSchema => SchemaSettings::draft07().into_generator(),
|
|
LanguageModelToolSchemaFormat::JsonSchemaSubset => SchemaSettings::openapi3()
|
|
.with(|settings| {
|
|
settings.meta_schema = None;
|
|
settings.inline_subschemas = true;
|
|
})
|
|
.with_transform(ToJsonSchemaSubsetTransform)
|
|
.into_generator(),
|
|
};
|
|
generator.root_schema_for::<T>()
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct ToJsonSchemaSubsetTransform;
|
|
|
|
impl Transform for ToJsonSchemaSubsetTransform {
|
|
fn transform(&mut self, schema: &mut Schema) {
|
|
// Ensure that the type field is not an array, this happens when we use
|
|
// Option<T>, the type will be [T, "null"].
|
|
if let Some(type_field) = schema.get_mut("type") {
|
|
if let Some(types) = type_field.as_array() {
|
|
if let Some(first_type) = types.first() {
|
|
*type_field = first_type.clone();
|
|
}
|
|
}
|
|
}
|
|
|
|
// oneOf is not supported, use anyOf instead
|
|
if let Some(one_of) = schema.remove("oneOf") {
|
|
schema.insert("anyOf".to_string(), one_of);
|
|
}
|
|
|
|
transform_subschemas(self, schema);
|
|
}
|
|
}
|