Skip to content

Commit a1a57b3

Browse files
committed
GH-3403: Add caching for JSON schema generation in function calling
Fixes GH-3403 (#3403) Add ConcurrentHashMap-based caching to eliminate repeated schema generation for identical method signatures in tool invocations. * Cache tool definitions by Method instance in ToolDefinitions * Cache JSON schemas with method signature + options as key * Include comprehensive test coverage for caching behavior Performance improvement: Eliminates CPU overhead from repeated schema generation in function calling applications. Signed-off-by: Seol-JY <[email protected]>
1 parent daf1274 commit a1a57b3

File tree

4 files changed

+418
-14
lines changed

4 files changed

+418
-14
lines changed

spring-ai-model/src/main/java/org/springframework/ai/tool/support/ToolDefinitions.java

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
package org.springframework.ai.tool.support;
1818

1919
import java.lang.reflect.Method;
20+
import java.util.Map;
21+
import java.util.concurrent.ConcurrentHashMap;
2022

2123
import org.springframework.ai.tool.definition.DefaultToolDefinition;
2224
import org.springframework.ai.tool.definition.ToolDefinition;
@@ -33,10 +35,16 @@
3335
* </p>
3436
*
3537
* @author Mark Pollack
38+
* @author Seol-JY
3639
* @since 1.0.0
3740
*/
3841
public final class ToolDefinitions {
3942

43+
/**
44+
* Cache for tool definitions. Key is the Method instance.
45+
*/
46+
private static final Map<Method, ToolDefinition> toolDefinitionCache = new ConcurrentHashMap<>(256);
47+
4048
private ToolDefinitions() {
4149
// prevents instantiation.
4250
}
@@ -56,7 +64,16 @@ public static DefaultToolDefinition.Builder builder(Method method) {
5664
* Create a default {@link ToolDefinition} instance from a {@link Method}.
5765
*/
5866
public static ToolDefinition from(Method method) {
59-
return builder(method).build();
67+
Assert.notNull(method, "method cannot be null");
68+
return toolDefinitionCache.computeIfAbsent(method, ToolDefinitions::createToolDefinition);
69+
}
70+
71+
private static ToolDefinition createToolDefinition(Method method) {
72+
return DefaultToolDefinition.builder()
73+
.name(ToolUtils.getToolName(method))
74+
.description(ToolUtils.getToolDescription(method))
75+
.inputSchema(JsonSchemaGenerator.generateForMethodInput(method))
76+
.build();
6077
}
6178

6279
}

spring-ai-model/src/main/java/org/springframework/ai/util/json/schema/JsonSchemaGenerator.java

Lines changed: 82 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@
2020
import java.lang.reflect.Parameter;
2121
import java.lang.reflect.Type;
2222
import java.util.ArrayList;
23+
import java.util.Arrays;
2324
import java.util.List;
25+
import java.util.Map;
26+
import java.util.concurrent.ConcurrentHashMap;
2427
import java.util.stream.Stream;
2528

2629
import com.fasterxml.jackson.annotation.JsonProperty;
@@ -67,8 +70,11 @@
6770
* If none of these annotations are present, the default behavior is to consider the
6871
* property as required and not to include a description.
6972
* <p>
73+
* This class provides caching for method input schema generation to improve performance
74+
* when the same method signatures are processed multiple times.
7075
*
7176
* @author Thomas Vitale
77+
* @author Seol-JY
7278
* @since 1.0.0
7379
*/
7480
public final class JsonSchemaGenerator {
@@ -81,6 +87,8 @@ public final class JsonSchemaGenerator {
8187
*/
8288
private static final boolean PROPERTY_REQUIRED_BY_DEFAULT = true;
8389

90+
private static final Map<String, String> methodSchemaCache = new ConcurrentHashMap<>(256);
91+
8492
private static final SchemaGenerator TYPE_SCHEMA_GENERATOR;
8593

8694
private static final SchemaGenerator SUBTYPE_SCHEMA_GENERATOR;
@@ -116,8 +124,82 @@ private JsonSchemaGenerator() {
116124

117125
/**
118126
* Generate a JSON Schema for a method's input parameters.
127+
*
128+
* <p>
129+
* This method uses caching to improve performance when the same method signature is
130+
* processed multiple times. The cache key includes method signature and schema
131+
* options to ensure correct cache hits.
132+
* @param method the method to generate schema for
133+
* @param schemaOptions options for schema generation
134+
* @return JSON Schema as a string
135+
* @throws IllegalArgumentException if method is null
119136
*/
120137
public static String generateForMethodInput(Method method, SchemaOption... schemaOptions) {
138+
Assert.notNull(method, "method cannot be null");
139+
140+
String cacheKey = buildMethodCacheKey(method, schemaOptions);
141+
return methodSchemaCache.computeIfAbsent(cacheKey, key -> generateMethodSchemaInternal(method, schemaOptions));
142+
}
143+
144+
/**
145+
* Generate a JSON Schema for a class type.
146+
*/
147+
public static String generateForType(Type type, SchemaOption... schemaOptions) {
148+
Assert.notNull(type, "type cannot be null");
149+
ObjectNode schema = TYPE_SCHEMA_GENERATOR.generateSchema(type);
150+
if ((type == Void.class) && !schema.has("properties")) {
151+
schema.putObject("properties");
152+
}
153+
processSchemaOptions(schemaOptions, schema);
154+
return schema.toPrettyString();
155+
}
156+
157+
/**
158+
* Build cache key for method input schema generation.
159+
*
160+
* <p>
161+
* The cache key includes:
162+
* <ul>
163+
* <li>Declaring class name</li>
164+
* <li>Method name</li>
165+
* <li>Parameter types (including generics)</li>
166+
* <li>Schema options</li>
167+
* </ul>
168+
* @param method the method
169+
* @param schemaOptions schema generation options
170+
* @return unique cache key
171+
*/
172+
private static String buildMethodCacheKey(Method method, SchemaOption... schemaOptions) {
173+
StringBuilder keyBuilder = new StringBuilder(256);
174+
175+
// Class name
176+
keyBuilder.append(method.getDeclaringClass().getName());
177+
keyBuilder.append('#');
178+
179+
// Method name
180+
keyBuilder.append(method.getName());
181+
keyBuilder.append('(');
182+
183+
// Parameter types (including generic information)
184+
Type[] parameterTypes = method.getGenericParameterTypes();
185+
for (int i = 0; i < parameterTypes.length; i++) {
186+
if (i > 0) {
187+
keyBuilder.append(',');
188+
}
189+
keyBuilder.append(parameterTypes[i].getTypeName());
190+
}
191+
keyBuilder.append(')');
192+
193+
// Schema options
194+
if (schemaOptions.length > 0) {
195+
keyBuilder.append(':');
196+
keyBuilder.append(Arrays.toString(schemaOptions));
197+
}
198+
199+
return keyBuilder.toString();
200+
}
201+
202+
private static String generateMethodSchemaInternal(Method method, SchemaOption... schemaOptions) {
121203
ObjectNode schema = JsonParser.getObjectMapper().createObjectNode();
122204
schema.put("$schema", SchemaVersion.DRAFT_2020_12.getIdentifier());
123205
schema.put("type", "object");
@@ -155,19 +237,6 @@ public static String generateForMethodInput(Method method, SchemaOption... schem
155237
return schema.toPrettyString();
156238
}
157239

158-
/**
159-
* Generate a JSON Schema for a class type.
160-
*/
161-
public static String generateForType(Type type, SchemaOption... schemaOptions) {
162-
Assert.notNull(type, "type cannot be null");
163-
ObjectNode schema = TYPE_SCHEMA_GENERATOR.generateSchema(type);
164-
if ((type == Void.class) && !schema.has("properties")) {
165-
schema.putObject("properties");
166-
}
167-
processSchemaOptions(schemaOptions, schema);
168-
return schema.toPrettyString();
169-
}
170-
171240
private static void processSchemaOptions(SchemaOption[] schemaOptions, ObjectNode schema) {
172241
if (Stream.of(schemaOptions)
173242
.noneMatch(option -> option == SchemaOption.ALLOW_ADDITIONAL_PROPERTIES_BY_DEFAULT)) {

0 commit comments

Comments
 (0)