diff --git a/Cargo.lock b/Cargo.lock index f976aefdf3..c9f821f7f1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9651,6 +9651,7 @@ dependencies = [ "futures 0.3.32", "gpui_shared_string", "http_client", + "log", "partial-json-fixer", "schemars 1.0.4", "serde", diff --git a/crates/language_model_core/Cargo.toml b/crates/language_model_core/Cargo.toml index e9aa06400b..c254989b4d 100644 --- a/crates/language_model_core/Cargo.toml +++ b/crates/language_model_core/Cargo.toml @@ -19,6 +19,7 @@ cloud_llm_client.workspace = true futures.workspace = true gpui_shared_string.workspace = true http_client.workspace = true +log.workspace = true partial-json-fixer.workspace = true schemars.workspace = true serde.workspace = true diff --git a/crates/language_model_core/src/tool_schema.rs b/crates/language_model_core/src/tool_schema.rs index 86e6d6d137..729d625939 100644 --- a/crates/language_model_core/src/tool_schema.rs +++ b/crates/language_model_core/src/tool_schema.rs @@ -4,7 +4,7 @@ use schemars::{ generate::SchemaSettings, transform::{Transform, transform_subschemas}, }; -use serde_json::Value; +use serde_json::{Map, Value, json}; /// Indicates the format used to define the input schema for a language model tool. #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)] @@ -64,6 +64,8 @@ pub fn adapt_schema_to_format( json: &mut Value, format: LanguageModelToolSchemaFormat, ) -> Result<()> { + log::trace!("Adapting schema to format {:?}: {}", format, json); + if let Value::Object(obj) = json { obj.remove("$schema"); obj.remove("title"); @@ -73,7 +75,10 @@ pub fn adapt_schema_to_format( match format { LanguageModelToolSchemaFormat::JsonSchema => preprocess_json_schema(json), LanguageModelToolSchemaFormat::JsonSchemaSubset => adapt_to_json_schema_subset(json), - } + }?; + + log::trace!("Adapted schema: {}", json); + Ok(()) } fn preprocess_json_schema(json: &mut Value) -> Result<()> { @@ -118,6 +123,9 @@ fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { } } + convert_null_in_types_to_nullable(obj); + convert_types_to_any_of_defs(obj); + // Ensure that the type field is not an array. This can happen with MCP tool // schemas that use multiple types (e.g. `["string", "number"]` or `["string", "null"]`). if let Some(type_value) = obj.get_mut("type") @@ -126,7 +134,6 @@ fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { { *type_value = first_type; } - if matches!(obj.get("description"), Some(Value::String(_))) && !obj.contains_key("type") && !(obj.contains_key("anyOf") @@ -141,7 +148,7 @@ fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { { let subschemas_clone = subschemas.clone(); obj.remove("oneOf"); - obj.insert("anyOf".to_string(), subschemas_clone); + push_any_of_constraint(obj, subschemas_clone); } for (_, value) in obj.iter_mut() { @@ -157,11 +164,278 @@ fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { Ok(()) } +fn convert_null_in_types_to_nullable(obj: &mut Map) { + let mut nullable_found_in_type = false; + + if let Some(type_entry) = obj.get_mut("type") { + if let Some(types) = type_entry.as_array_mut() { + let mut had_null_type = false; + types.retain(|t| { + if t.as_str() == Some("null") { + had_null_type = true; + false + } else { + true + } + }); + + if had_null_type { + nullable_found_in_type = true; + if types.len() == 1 { + *type_entry = types.remove(0); + } else if types.is_empty() { + obj.remove("type"); + } + } + } else if let Some(type_str) = type_entry.as_str() { + if type_str == "null" { + nullable_found_in_type = true; + obj.remove("type"); + } + } + } + if nullable_found_in_type { + obj.insert("nullable".to_string(), Value::Bool(true)); + } +} + +fn convert_types_to_any_of_defs(obj: &mut Map) { + if let Some(type_entry) = obj.get_mut("type") { + if let Some(types) = type_entry.as_array_mut() { + if types.len() > 1 { + let remaining_types = std::mem::take(types); + let mut any_of_schemas = Vec::new(); + for t in remaining_types { + any_of_schemas.push(json!({"type": t})); + } + obj.remove("type"); + push_any_of_constraint(obj, Value::Array(any_of_schemas)); + } + } + } +} + +fn push_any_of_constraint(obj: &mut Map, any_of_schemas: Value) { + if let Some(existing_any_of) = obj.remove("anyOf") { + let mut all_of = obj + .remove("allOf") + .and_then(|v| v.as_array().cloned()) + .unwrap_or_default(); + if all_of.is_empty() { + all_of.push(json!({"anyOf": existing_any_of})); + } + all_of.push(json!({"anyOf": any_of_schemas})); + obj.insert("allOf".to_string(), Value::Array(all_of)); + } else if let Some(all_of) = obj.get_mut("allOf").and_then(|v| v.as_array_mut()) { + all_of.push(json!({"anyOf": any_of_schemas})); + } else { + obj.insert("anyOf".to_string(), any_of_schemas); + } +} + #[cfg(test)] mod tests { use super::*; use serde_json::json; + #[test] + fn test_convert_null_in_types_to_nullable() { + // ["string", "null"] -> "string", nullable: true + let mut obj = json!({"type": ["string", "null"]}) + .as_object_mut() + .unwrap() + .to_owned(); + convert_null_in_types_to_nullable(&mut obj); + assert_eq!( + obj, + json!({"type": "string", "nullable": true}) + .as_object() + .unwrap() + .to_owned() + ); + + // "null" -> nullable: true + let mut obj = json!({"type": "null"}).as_object_mut().unwrap().to_owned(); + convert_null_in_types_to_nullable(&mut obj); + assert_eq!( + obj, + json!({"nullable": true}).as_object().unwrap().to_owned() + ); + + // ["string", "number", "null"] -> ["string", "number"], nullable: true (anyOf handled elsewhere) + let mut obj = json!({"type": ["string", "number", "null"]}) + .as_object_mut() + .unwrap() + .to_owned(); + convert_null_in_types_to_nullable(&mut obj); + assert_eq!( + obj, + json!({"type": ["string", "number"], "nullable": true}) + .as_object() + .unwrap() + .to_owned() + ); + + // "string" (no change, not nullable) + let mut obj = json!({"type": "string"}) + .as_object_mut() + .unwrap() + .to_owned(); + convert_null_in_types_to_nullable(&mut obj); + assert_eq!( + obj, + json!({"type": "string"}).as_object().unwrap().to_owned() + ); + + // ["string", "number"] (no change, not nullable) + let mut obj = json!({"type": ["string", "number"]}) + .as_object_mut() + .unwrap() + .to_owned(); + convert_null_in_types_to_nullable(&mut obj); + assert_eq!( + obj, + json!({"type": ["string", "number"]}) + .as_object() + .unwrap() + .to_owned() + ); + + // object with other properties, ["boolean", "null"] + let mut obj = json!({ + "description": "A test field", + "type": ["boolean", "null"] + }) + .as_object_mut() + .unwrap() + .to_owned(); + convert_null_in_types_to_nullable(&mut obj); + assert_eq!( + obj, + json!({ + "description": "A test field", + "type": "boolean", + "nullable": true + }) + .as_object() + .unwrap() + .to_owned() + ); + } + + #[test] + fn test_convert_types_to_any_of_defs() { + // ["string", "number"] -> anyOf with string and number + let mut obj = json!({"type": ["string", "number"]}) + .as_object_mut() + .unwrap() + .to_owned(); + convert_types_to_any_of_defs(&mut obj); + assert_eq!( + obj, + json!({ + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ] + }) + .as_object() + .unwrap() + .to_owned() + ); + + // "string" (no change) + let mut obj = json!({"type": "string"}) + .as_object_mut() + .unwrap() + .to_owned(); + convert_types_to_any_of_defs(&mut obj); + assert_eq!( + obj, + json!({"type": "string"}).as_object().unwrap().to_owned() + ); + + // object with other properties, ["string", "number"] + let mut obj = json!({ + "description": "A test field", + "type": ["string", "number"] + }) + .as_object_mut() + .unwrap() + .to_owned(); + convert_types_to_any_of_defs(&mut obj); + assert_eq!( + obj, + json!({ + "description": "A test field", + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ] + }) + .as_object() + .unwrap() + .to_owned() + ); + + // anyOf already present (no change) + let mut obj = json!({ + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ] + }) + .as_object_mut() + .unwrap() + .to_owned(); + convert_types_to_any_of_defs(&mut obj); + assert_eq!( + obj, + json!({ + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ] + }) + .as_object() + .unwrap() + .to_owned() + ); + + // both type array and anyOf present + let mut obj = json!({ + "type": ["string", "number"], + "anyOf": [ + {"format": "email"} + ] + }) + .as_object_mut() + .unwrap() + .to_owned(); + convert_types_to_any_of_defs(&mut obj); + assert_eq!( + obj, + json!({ + "allOf": [ + { + "anyOf": [ + {"format": "email"} + ] + }, + { + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ] + } + ] + }) + .as_object() + .unwrap() + .to_owned() + ); + } + #[test] fn test_transform_adds_type_when_missing() { let mut json = json!({ @@ -259,6 +533,69 @@ mod tests { ); } + #[test] + fn test_transform_null_in_any_of() { + let mut json = json!({ + "anyOf": [ + { "type": "string" }, + { "type": "null" } + ] + }); + + adapt_to_json_schema_subset(&mut json).unwrap(); + + assert_eq!( + json, + json!({ + "anyOf": [ + { "type": "string" }, + { "nullable": true } + ] + }) + ); + } + + #[test] + fn test_transform_conflicting_any_of_sources() { + let mut json = json!({ + "type": ["string", "number"], + "anyOf": [ + { "minLength": 5 } + ], + "oneOf": [ + { "pattern": "^a" }, + { "pattern": "^b" } + ] + }); + + adapt_to_json_schema_subset(&mut json).unwrap(); + + assert_eq!( + json, + json!({ + "allOf": [ + { + "anyOf": [ + { "minLength": 5 }, + ] + }, + { + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ] + }, + { + "anyOf": [ + { "pattern": "^a" }, + { "pattern": "^b" } + ] + } + ] + }) + ); + } + #[test] fn test_transform_one_of_to_any_of() { let mut json = json!({ @@ -308,8 +645,8 @@ mod tests { "nested": { "anyOf": [ { "type": "string" }, - { "type": "null" } - ] + { "nullable": true } + ], } } }) @@ -340,12 +677,16 @@ mod tests { "type": "object", "properties": { "projectSlugOrId": { - "type": "string", + "anyOf": [ + {"type": "string"}, + {"type": "number"} + ], "description": "Project slug or numeric ID" }, "optionalName": { "type": "string", - "description": "An optional name" + "description": "An optional name", + "nullable": true } } })