Skip to content

Commit 5c69d08

Browse files
itxaiohangloveritxaiohanglover
committed
feat(tool): add additionalProperties support to @tool annotation
Add additionalProperties parameter to @tool annotation to control whether LLM can pass extra parameters not defined in the schema. When set to false, ToolValidator rejects hallucinated params. Closes #1379
1 parent 9e83aa7 commit 5c69d08

5 files changed

Lines changed: 313 additions & 1 deletion

File tree

‎agentscope-core/src/main/java/io/agentscope/core/tool/Tool.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,4 +130,18 @@
130130
* @see DefaultToolResultConverter
131131
*/
132132
Class<? extends ToolResultConverter> converter() default DefaultToolResultConverter.class;
133+
134+
/**
135+
* Whether to allow additional properties beyond those defined in the schema.
136+
*
137+
* <p>Corresponds to the {@code additionalProperties} keyword in JSON Schema. When set to
138+
* {@code false}, any extra parameters passed by the LLM that are not defined in the schema will
139+
* cause a validation error, preventing the LLM from hallucinating undefined parameters.
140+
*
141+
* <p>This setting is applied recursively to all nested objects within the generated schema.
142+
*
143+
* @return {@code true} to allow additional properties (default, backward compatible), {@code
144+
* false} to disallow
145+
*/
146+
boolean additionalProperties() default true;
133147
}

‎agentscope-core/src/main/java/io/agentscope/core/tool/ToolSchemaGenerator.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,77 @@
3333
*/
3434
class ToolSchemaGenerator {
3535

36+
/**
37+
* Generate parameter schema with additionalProperties control.
38+
*
39+
* @param method the method to generate schema for
40+
* @param excludeParams set of parameter names to exclude (may be null or empty)
41+
* @param allowAdditionalProperties whether to allow additional properties in the schema
42+
* @return JSON Schema map
43+
*/
44+
@SuppressWarnings("unchecked")
45+
Map<String, Object> generateParameterSchema(
46+
Method method, Set<String> excludeParams, boolean allowAdditionalProperties) {
47+
Map<String, Object> schema = generateParameterSchema(method, excludeParams);
48+
49+
if (!allowAdditionalProperties) {
50+
addAdditionalPropertiesFalseRecursively(schema);
51+
}
52+
53+
return schema;
54+
}
55+
56+
@SuppressWarnings("unchecked")
57+
private void addAdditionalPropertiesFalseRecursively(Map<String, Object> schema) {
58+
if (!"object".equals(schema.get("type"))) {
59+
return;
60+
}
61+
// Skip if user already set additionalProperties explicitly
62+
if (schema.containsKey("additionalProperties")) {
63+
return;
64+
}
65+
schema.put("additionalProperties", false);
66+
67+
Object propsObj = schema.get("properties");
68+
if (propsObj instanceof Map) {
69+
Map<String, Object> props = (Map<String, Object>) propsObj;
70+
for (Map.Entry<String, Object> entry : props.entrySet()) {
71+
if (entry.getValue() instanceof Map) {
72+
addAdditionalPropertiesFalseRecursively((Map<String, Object>) entry.getValue());
73+
}
74+
}
75+
}
76+
77+
Object itemsObj = schema.get("items");
78+
if (itemsObj instanceof Map) {
79+
addAdditionalPropertiesFalseRecursively((Map<String, Object>) itemsObj);
80+
}
81+
82+
for (String key : new String[] {"oneOf", "anyOf", "allOf"}) {
83+
Object listObj = schema.get(key);
84+
if (listObj instanceof Iterable) {
85+
for (Object item : (Iterable<?>) listObj) {
86+
if (item instanceof Map) {
87+
addAdditionalPropertiesFalseRecursively((Map<String, Object>) item);
88+
}
89+
}
90+
}
91+
}
92+
93+
for (String defsKey : new String[] {"$defs", "definitions"}) {
94+
Object defsObj = schema.get(defsKey);
95+
if (defsObj instanceof Map) {
96+
Map<String, Object> defs = (Map<String, Object>) defsObj;
97+
for (Map.Entry<String, Object> entry : defs.entrySet()) {
98+
if (entry.getValue() instanceof Map) {
99+
addAdditionalPropertiesFalseRecursively(
100+
(Map<String, Object>) entry.getValue());
101+
}
102+
}
103+
}
104+
}
105+
}
106+
36107
/**
37108
* Generate parameter schema for a method with excluded parameters.
38109
*

‎agentscope-core/src/main/java/io/agentscope/core/tool/Toolkit.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -373,7 +373,8 @@ public Map<String, Object> getParameters() {
373373
presetParameters != null
374374
? presetParameters.keySet()
375375
: Collections.emptySet();
376-
return schemaGenerator.generateParameterSchema(method, excludeParams);
376+
return schemaGenerator.generateParameterSchema(
377+
method, excludeParams, toolAnnotation.additionalProperties());
377378
}
378379

379380
@Override

‎agentscope-core/src/test/java/io/agentscope/core/tool/ToolAnnotationTest.java‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,10 @@
1515
*/
1616
package io.agentscope.core.tool;
1717

18+
import static org.junit.jupiter.api.Assertions.assertEquals;
19+
import static org.junit.jupiter.api.Assertions.assertFalse;
1820
import static org.junit.jupiter.api.Assertions.assertNotNull;
21+
import static org.junit.jupiter.api.Assertions.assertNull;
1922
import static org.junit.jupiter.api.Assertions.assertTrue;
2023

2124
import io.agentscope.core.tool.test.SampleTools;
@@ -142,4 +145,104 @@ void testInvalidAnnotations() {
142145
assertTrue(
143146
registeredTools <= methods.length, "Should not register more tools than methods");
144147
}
148+
149+
@SuppressWarnings("unchecked")
150+
@Test
151+
@DisplayName("Should include additionalProperties=false in schema when set")
152+
void testAdditionalPropertiesFalseInSchema() {
153+
toolkit.registerTool(
154+
new Object() {
155+
@Tool(
156+
name = "query_data",
157+
description = "Query data",
158+
additionalProperties = false)
159+
public String queryData(
160+
@ToolParam(name = "keyword", description = "keyword") String keyword) {
161+
return "result";
162+
}
163+
});
164+
165+
Map<String, Object> schema = toolkit.getToolSchemas().get(0).getParameters();
166+
167+
assertTrue(
168+
schema.containsKey("additionalProperties"),
169+
"Schema should contain additionalProperties key");
170+
assertEquals(
171+
false, schema.get("additionalProperties"), "additionalProperties should be false");
172+
173+
Map<String, Object> properties = (Map<String, Object>) schema.get("properties");
174+
assertFalse(
175+
properties.containsKey("additionalProperties"),
176+
"additionalProperties should NOT be inside properties map");
177+
}
178+
179+
@Test
180+
@DisplayName("Should not include additionalProperties when default (true)")
181+
void testAdditionalPropertiesDefaultNotInSchema() {
182+
toolkit.registerTool(
183+
new Object() {
184+
@Tool(name = "query_data", description = "Query data")
185+
public String queryData(
186+
@ToolParam(name = "keyword", description = "keyword") String keyword) {
187+
return "result";
188+
}
189+
});
190+
191+
Map<String, Object> schema = toolkit.getToolSchemas().get(0).getParameters();
192+
193+
assertFalse(
194+
schema.containsKey("additionalProperties"),
195+
"Default behavior should NOT add additionalProperties to schema");
196+
}
197+
198+
@Test
199+
@DisplayName("Should reject hallucinated params when additionalProperties=false")
200+
void testRejectHallucinatedParamsWhenStrict() {
201+
toolkit.registerTool(
202+
new Object() {
203+
@Tool(
204+
name = "query_data",
205+
description = "Query data",
206+
additionalProperties = false)
207+
public String queryData(
208+
@ToolParam(name = "keyword", description = "keyword") String keyword) {
209+
return "result";
210+
}
211+
});
212+
213+
Map<String, Object> schema = toolkit.getToolSchemas().get(0).getParameters();
214+
String hallucinatedInput = "{\"keyword\": \"test\", \"owners\": [\"admin\"]}";
215+
216+
String validationError = ToolValidator.validateInput(hallucinatedInput, schema);
217+
218+
assertNotNull(
219+
validationError,
220+
"ToolValidator should reject hallucinated params when "
221+
+ "additionalProperties=false");
222+
assertTrue(
223+
validationError.contains("owners"),
224+
"Error message should mention the hallucinated param");
225+
}
226+
227+
@Test
228+
@DisplayName("Should allow extra params by default (backward compatible)")
229+
void testAllowExtraParamsByDefault() {
230+
toolkit.registerTool(
231+
new Object() {
232+
@Tool(name = "query_data", description = "Query data")
233+
public String queryData(
234+
@ToolParam(name = "keyword", description = "keyword") String keyword) {
235+
return "result";
236+
}
237+
});
238+
239+
Map<String, Object> schema = toolkit.getToolSchemas().get(0).getParameters();
240+
String hallucinatedInput = "{\"keyword\": \"test\", \"owners\": [\"admin\"]}";
241+
242+
String validationError = ToolValidator.validateInput(hallucinatedInput, schema);
243+
244+
assertNull(
245+
validationError,
246+
"Default behavior should allow extra params (backward compatible)");
247+
}
145248
}

‎agentscope-core/src/test/java/io/agentscope/core/tool/ToolSchemaGeneratorTest.java‎

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,12 +19,14 @@
1919
import static org.junit.jupiter.api.Assertions.assertEquals;
2020
import static org.junit.jupiter.api.Assertions.assertFalse;
2121
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
22+
import static org.junit.jupiter.api.Assertions.assertNull;
2223
import static org.junit.jupiter.api.Assertions.assertThrows;
2324
import static org.junit.jupiter.api.Assertions.assertTrue;
2425

2526
import java.lang.reflect.InvocationTargetException;
2627
import java.lang.reflect.Method;
2728
import java.util.HashMap;
29+
import java.util.List;
2830
import java.util.Map;
2931
import org.junit.jupiter.api.DisplayName;
3032
import org.junit.jupiter.api.Tag;
@@ -83,4 +85,125 @@ void testHoistDefsEquivalent() throws Exception {
8385
assertEquals(definition, target.get("Material"));
8486
assertFalse(paramSchema.containsKey("$defs"));
8587
}
88+
89+
@SuppressWarnings("unchecked")
90+
@Test
91+
@DisplayName("Should inject additionalProperties=false recursively into nested structures")
92+
void testAddAdditionalPropertiesFalseRecursively() throws Exception {
93+
Method recMethod =
94+
ToolSchemaGenerator.class.getDeclaredMethod(
95+
"addAdditionalPropertiesFalseRecursively", Map.class);
96+
recMethod.setAccessible(true);
97+
98+
// Object with items
99+
Map<String, Object> schemaWithItems = new HashMap<>();
100+
schemaWithItems.put("type", "object");
101+
Map<String, Object> itemsObj = new HashMap<>();
102+
itemsObj.put("type", "object");
103+
schemaWithItems.put("items", itemsObj);
104+
105+
recMethod.invoke(generator, schemaWithItems);
106+
107+
assertEquals(false, schemaWithItems.get("additionalProperties"));
108+
assertEquals(false, itemsObj.get("additionalProperties"));
109+
110+
// Object with oneOf
111+
Map<String, Object> schemaWithOneOf = new HashMap<>();
112+
schemaWithOneOf.put("type", "object");
113+
Map<String, Object> oneOfItem = new HashMap<>();
114+
oneOfItem.put("type", "object");
115+
schemaWithOneOf.put("oneOf", List.of(oneOfItem));
116+
117+
recMethod.invoke(generator, schemaWithOneOf);
118+
119+
assertEquals(false, schemaWithOneOf.get("additionalProperties"));
120+
assertEquals(false, oneOfItem.get("additionalProperties"));
121+
122+
// Object with anyOf and allOf
123+
Map<String, Object> schemaWithAnyOf = new HashMap<>();
124+
schemaWithAnyOf.put("type", "object");
125+
Map<String, Object> anyOfItem = new HashMap<>();
126+
anyOfItem.put("type", "object");
127+
Map<String, Object> allOfItem = new HashMap<>();
128+
allOfItem.put("type", "object");
129+
schemaWithAnyOf.put("anyOf", List.of(anyOfItem));
130+
schemaWithAnyOf.put("allOf", List.of(allOfItem));
131+
132+
recMethod.invoke(generator, schemaWithAnyOf);
133+
134+
assertEquals(false, anyOfItem.get("additionalProperties"));
135+
assertEquals(false, allOfItem.get("additionalProperties"));
136+
137+
// Object with $defs
138+
Map<String, Object> schemaWithDefs = new HashMap<>();
139+
schemaWithDefs.put("type", "object");
140+
Map<String, Object> defEntry = new HashMap<>();
141+
defEntry.put("type", "object");
142+
schemaWithDefs.put("$defs", Map.of("MyType", defEntry));
143+
144+
recMethod.invoke(generator, schemaWithDefs);
145+
146+
assertEquals(false, defEntry.get("additionalProperties"));
147+
148+
// Object with definitions (legacy key)
149+
Map<String, Object> schemaWithDefinitions = new HashMap<>();
150+
schemaWithDefinitions.put("type", "object");
151+
Map<String, Object> defEntry2 = new HashMap<>();
152+
defEntry2.put("type", "object");
153+
schemaWithDefinitions.put("definitions", Map.of("LegacyType", defEntry2));
154+
155+
recMethod.invoke(generator, schemaWithDefinitions);
156+
157+
assertEquals(false, defEntry2.get("additionalProperties"));
158+
}
159+
160+
@Test
161+
@DisplayName("Should skip when additionalProperties already set")
162+
void testSkipWhenAdditionalPropertiesAlreadySet() throws Exception {
163+
Method recMethod =
164+
ToolSchemaGenerator.class.getDeclaredMethod(
165+
"addAdditionalPropertiesFalseRecursively", Map.class);
166+
recMethod.setAccessible(true);
167+
168+
Map<String, Object> schema = new HashMap<>();
169+
schema.put("type", "object");
170+
schema.put("additionalProperties", true);
171+
172+
recMethod.invoke(generator, schema);
173+
174+
assertEquals(true, schema.get("additionalProperties"));
175+
}
176+
177+
@Test
178+
@DisplayName("Should skip non-object schemas")
179+
void testSkipNonObjectSchemas() throws Exception {
180+
Method recMethod =
181+
ToolSchemaGenerator.class.getDeclaredMethod(
182+
"addAdditionalPropertiesFalseRecursively", Map.class);
183+
recMethod.setAccessible(true);
184+
185+
Map<String, Object> schema = new HashMap<>();
186+
schema.put("type", "string");
187+
188+
recMethod.invoke(generator, schema);
189+
190+
assertNull(schema.get("additionalProperties"));
191+
}
192+
193+
@Test
194+
@DisplayName("Should handle properties with non-Map values")
195+
void testPropertiesWithNonMapValues() throws Exception {
196+
Method recMethod =
197+
ToolSchemaGenerator.class.getDeclaredMethod(
198+
"addAdditionalPropertiesFalseRecursively", Map.class);
199+
recMethod.setAccessible(true);
200+
201+
Map<String, Object> schema = new HashMap<>();
202+
schema.put("type", "object");
203+
schema.put("properties", Map.of("simple", "not-a-map"));
204+
205+
recMethod.invoke(generator, schema);
206+
207+
assertEquals(false, schema.get("additionalProperties"));
208+
}
86209
}

0 commit comments

Comments
 (0)