assistant_tools: Disallow extra tool parameters by default (#32081)

Oleksiy Syvokon created 6 months ago

This prevents models from hallucinating tool parameters.


Release Notes:

- Prevent models from hallucinating tool parameters

Change summary

crates/assistant_tool/src/tool_schema.rs      | 70 ++++++++++++++++++++
crates/assistant_tools/src/assistant_tools.rs |  1 
2 files changed, 70 insertions(+), 1 deletion(-)

Detailed changes

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
+            })
+        );
+    }
 }

crates/assistant_tools/src/assistant_tools.rs

@@ -126,6 +126,7 @@ mod tests {
                     }
                 },
                 "required": ["location"],
+                "additionalProperties": false
             })
         );
     }