diff --git a/crates/assistant_tool/src/tool_schema.rs b/crates/assistant_tool/src/tool_schema.rs index 57fc2a4d49..39478499d9 100644 --- a/crates/assistant_tool/src/tool_schema.rs +++ b/crates/assistant_tool/src/tool_schema.rs @@ -16,11 +16,24 @@ pub fn adapt_schema_to_format( } match format { - LanguageModelToolSchemaFormat::JsonSchema => Ok(()), + LanguageModelToolSchemaFormat::JsonSchema => preprocess_json_schema(json), LanguageModelToolSchemaFormat::JsonSchemaSubset => adapt_to_json_schema_subset(json), } } +fn preprocess_json_schema(json: &mut Value) -> Result<()> { + // `additionalProperties` defaults to `false` unless explicitly specified. + // This prevents models from hallucinating tool parameters. + if let Value::Object(obj) = json { + if let Some(Value::String(type_str)) = obj.get("type") { + if type_str == "object" && !obj.contains_key("additionalProperties") { + obj.insert("additionalProperties".to_string(), Value::Bool(false)); + } + } + } + Ok(()) +} + /// Tries to adapt the json schema so that it is compatible with https://ai.google.dev/api/caching#Schema fn adapt_to_json_schema_subset(json: &mut Value) -> Result<()> { if let Value::Object(obj) = json { @@ -237,4 +250,59 @@ mod tests { assert!(adapt_to_json_schema_subset(&mut json).is_err()); } + + #[test] + fn test_preprocess_json_schema_adds_additional_properties() { + let mut json = json!({ + "type": "object", + "properties": { + "name": { + "type": "string" + } + } + }); + + preprocess_json_schema(&mut json).unwrap(); + + assert_eq!( + json, + json!({ + "type": "object", + "properties": { + "name": { + "type": "string" + } + }, + "additionalProperties": false + }) + ); + } + + #[test] + fn test_preprocess_json_schema_preserves_additional_properties() { + let mut json = json!({ + "type": "object", + "properties": { + "name": { + "type": "string" + } + }, + "additionalProperties": true + }); + + preprocess_json_schema(&mut json).unwrap(); + + assert_eq!( + json, + json!({ + "type": "object", + "properties": { + "name": { + "type": "string" + } + }, + "additionalProperties": true + }) + ); + } } diff --git a/crates/assistant_tools/src/assistant_tools.rs b/crates/assistant_tools/src/assistant_tools.rs index 6bed4c216b..de57e8ebb3 100644 --- a/crates/assistant_tools/src/assistant_tools.rs +++ b/crates/assistant_tools/src/assistant_tools.rs @@ -126,6 +126,7 @@ mod tests { } }, "required": ["location"], + "additionalProperties": false }) ); }