Skip to content

Commit 923c5cc

Browse files
sobychackosobychacko
committed
Refactor RuntimeHints to scan entire packages for JsonInclude annotations
- Replace class-by-class scanning with comprehensive package-level scanning to capture all JSON-annotated classes within each model's package hierarchy. - Update tests to verify registered types and add specific type assertions. Signed-off-by: Soby Chacko <[email protected]>
1 parent b313645 commit 923c5cc

File tree

20 files changed

+153
-88
lines changed

20 files changed

+153
-88
lines changed

‎models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/aot/AnthropicRuntimeHints.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.ai.anthropic.aot;
1818

19+
import org.springframework.ai.anthropic.AnthropicChatOptions;
1920
import org.springframework.ai.anthropic.api.AnthropicApi;
2021
import org.springframework.aot.hint.MemberCategory;
2122
import org.springframework.aot.hint.RuntimeHints;
@@ -37,7 +38,8 @@ public class AnthropicRuntimeHints implements RuntimeHintsRegistrar {
3738
@Override
3839
public void registerHints(@NonNull RuntimeHints hints, @Nullable ClassLoader classLoader) {
3940
var mcs = MemberCategory.values();
40-
for (var tr : findJsonAnnotatedClassesInPackage(AnthropicApi.class)) {
41+
42+
for (var tr : findJsonAnnotatedClassesInPackage("org.springframework.ai.anthropic")) {
4143
hints.reflection().registerType(tr, mcs);
4244
}
4345
}

‎models/spring-ai-anthropic/src/test/java/org/springframework/ai/anthropic/aot/AnthropicRuntimeHintsTests.java‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.ai.anthropic.aot;
1818

19+
import java.util.HashSet;
1920
import java.util.Set;
2021

2122
import org.junit.jupiter.api.Test;
@@ -26,7 +27,6 @@
2627

2728
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
2829
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
29-
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
3030

3131
class AnthropicRuntimeHintsTests {
3232

@@ -36,10 +36,23 @@ void registerHints() {
3636
AnthropicRuntimeHints anthropicRuntimeHints = new AnthropicRuntimeHints();
3737
anthropicRuntimeHints.registerHints(runtimeHints, null);
3838

39-
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(AnthropicApi.class);
39+
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage("org.springframework.ai.anthropic");
40+
41+
Set<TypeReference> registeredTypes = new HashSet<>();
42+
runtimeHints.reflection().typeHints().forEach(typeHint -> registeredTypes.add(typeHint.getType()));
43+
4044
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
41-
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
45+
assertThat(registeredTypes.contains(jsonAnnotatedClass)).isTrue();
4246
}
47+
48+
// Check a few more specific ones
49+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.class))).isTrue();
50+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.Role.class))).isTrue();
51+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.ThinkingType.class))).isTrue();
52+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.EventType.class))).isTrue();
53+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.ContentBlock.class))).isTrue();
54+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.ChatCompletionRequest.class))).isTrue();
55+
assertThat(registeredTypes.contains(TypeReference.of(AnthropicApi.AnthropicMessage.class))).isTrue();
4356
}
4457

4558
}

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java‎

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -41,21 +41,8 @@ public class BedrockRuntimeHints implements RuntimeHintsRegistrar {
4141
@Override
4242
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
4343
var mcs = MemberCategory.values();
44-
for (var tr : findJsonAnnotatedClassesInPackage(AbstractBedrockApi.class)) {
45-
hints.reflection().registerType(tr, mcs);
46-
}
4744

48-
for (var tr : findJsonAnnotatedClassesInPackage(CohereEmbeddingBedrockApi.class)) {
49-
hints.reflection().registerType(tr, mcs);
50-
}
51-
for (var tr : findJsonAnnotatedClassesInPackage(BedrockCohereEmbeddingOptions.class)) {
52-
hints.reflection().registerType(tr, mcs);
53-
}
54-
55-
for (var tr : findJsonAnnotatedClassesInPackage(BedrockTitanEmbeddingOptions.class)) {
56-
hints.reflection().registerType(tr, mcs);
57-
}
58-
for (var tr : findJsonAnnotatedClassesInPackage(TitanEmbeddingBedrockApi.class)) {
45+
for (var tr : findJsonAnnotatedClassesInPackage("org.springframework.ai.bedrock")) {
5946
hints.reflection().registerType(tr, mcs);
6047
}
6148
}

‎models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHintsTests.java‎

Lines changed: 17 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,20 +16,21 @@
1616

1717
package org.springframework.ai.bedrock.aot;
1818

19-
import java.util.Arrays;
20-
import java.util.List;
19+
import java.util.HashSet;
2120
import java.util.Set;
2221

2322
import org.junit.jupiter.api.Test;
2423

24+
import org.springframework.ai.bedrock.api.AbstractBedrockApi;
25+
import org.springframework.ai.bedrock.cohere.BedrockCohereEmbeddingOptions;
2526
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
27+
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingOptions;
2628
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
2729
import org.springframework.aot.hint.RuntimeHints;
2830
import org.springframework.aot.hint.TypeReference;
2931

3032
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
3133
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
32-
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
3334

3435
class BedrockRuntimeHintsTests {
3536

@@ -39,15 +40,22 @@ void registerHints() {
3940
BedrockRuntimeHints bedrockRuntimeHints = new BedrockRuntimeHints();
4041
bedrockRuntimeHints.registerHints(runtimeHints, null);
4142

42-
List<Class> classList = Arrays.asList(CohereEmbeddingBedrockApi.class, TitanEmbeddingBedrockApi.class);
43+
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage("org.springframework.ai.bedrock");
4344

44-
for (Class aClass : classList) {
45-
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(aClass);
46-
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
47-
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
48-
}
45+
Set<TypeReference> registeredTypes = new HashSet<>();
46+
runtimeHints.reflection().typeHints().forEach(typeHint -> registeredTypes.add(typeHint.getType()));
47+
48+
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
49+
assertThat(registeredTypes.contains(jsonAnnotatedClass)).isTrue();
4950
}
5051

52+
// Check a few more specific ones
53+
assertThat(registeredTypes.contains(TypeReference.of(AbstractBedrockApi.AmazonBedrockInvocationMetrics.class)))
54+
.isTrue();
55+
assertThat(registeredTypes.contains(TypeReference.of(CohereEmbeddingBedrockApi.class))).isTrue();
56+
assertThat(registeredTypes.contains(TypeReference.of(BedrockCohereEmbeddingOptions.class))).isTrue();
57+
assertThat(registeredTypes.contains(TypeReference.of(BedrockTitanEmbeddingOptions.class))).isTrue();
58+
assertThat(registeredTypes.contains(TypeReference.of(TitanEmbeddingBedrockApi.class))).isTrue();
5159
}
5260

5361
}

‎models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/aot/MiniMaxRuntimeHints.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,8 @@ public class MiniMaxRuntimeHints implements RuntimeHintsRegistrar {
3737
@Override
3838
public void registerHints(@NonNull RuntimeHints hints, @Nullable ClassLoader classLoader) {
3939
var mcs = MemberCategory.values();
40-
for (var tr : findJsonAnnotatedClassesInPackage(MiniMaxApi.class)) {
40+
41+
for (var tr : findJsonAnnotatedClassesInPackage("org.springframework.ai.minimax")) {
4142
hints.reflection().registerType(tr, mcs);
4243
}
4344
}

‎models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/aot/MistralAiRuntimeHints.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,14 @@
1616

1717
package org.springframework.ai.mistralai.aot;
1818

19-
import org.springframework.ai.mistralai.api.MistralAiApi;
2019
import org.springframework.aot.hint.MemberCategory;
2120
import org.springframework.aot.hint.RuntimeHints;
2221
import org.springframework.aot.hint.RuntimeHintsRegistrar;
2322
import org.springframework.lang.NonNull;
2423
import org.springframework.lang.Nullable;
2524

25+
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
26+
2627
/**
2728
* The MistralAiRuntimeHints class is responsible for registering runtime hints for
2829
* Mistral AI API classes.
@@ -35,7 +36,8 @@ public class MistralAiRuntimeHints implements RuntimeHintsRegistrar {
3536
@Override
3637
public void registerHints(@NonNull RuntimeHints hints, @Nullable ClassLoader classLoader) {
3738
var mcs = MemberCategory.values();
38-
for (var tr : org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage(MistralAiApi.class)) {
39+
40+
for (var tr : findJsonAnnotatedClassesInPackage("org.springframework.ai.mistralai")) {
3941
hints.reflection().registerType(tr, mcs);
4042
}
4143
}

‎models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/aot/MistralAiRuntimeHintsTests.java‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,19 @@
1616

1717
package org.springframework.ai.mistralai.aot;
1818

19+
import java.util.HashSet;
1920
import java.util.Set;
2021

2122
import org.junit.jupiter.api.Test;
2223

24+
import org.springframework.ai.mistralai.MistralAiChatOptions;
25+
import org.springframework.ai.mistralai.MistralAiEmbeddingOptions;
2326
import org.springframework.ai.mistralai.api.MistralAiApi;
2427
import org.springframework.aot.hint.RuntimeHints;
2528
import org.springframework.aot.hint.TypeReference;
2629

2730
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
31+
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
2832
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
2933

3034
class MistralAiRuntimeHintsTests {
@@ -35,11 +39,22 @@ void registerHints() {
3539
MistralAiRuntimeHints mistralAiRuntimeHints = new MistralAiRuntimeHints();
3640
mistralAiRuntimeHints.registerHints(runtimeHints, null);
3741

38-
Set<TypeReference> jsonAnnotatedClasses = org.springframework.ai.aot.AiRuntimeHints
39-
.findJsonAnnotatedClassesInPackage(MistralAiApi.class);
42+
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage("org.springframework.ai.mistralai");
43+
44+
Set<TypeReference> registeredTypes = new HashSet<>();
45+
runtimeHints.reflection().typeHints().forEach(typeHint -> registeredTypes.add(typeHint.getType()));
46+
4047
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
41-
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
48+
assertThat(registeredTypes.contains(jsonAnnotatedClass)).isTrue();
4249
}
50+
51+
// Check a few more specific ones
52+
assertThat(registeredTypes.contains(TypeReference.of(MistralAiApi.ChatCompletion.class))).isTrue();
53+
assertThat(registeredTypes.contains(TypeReference.of(MistralAiApi.ChatCompletionChunk.class))).isTrue();
54+
assertThat(registeredTypes.contains(TypeReference.of(MistralAiApi.LogProbs.class))).isTrue();
55+
assertThat(registeredTypes.contains(TypeReference.of(MistralAiApi.ChatCompletionFinishReason.class))).isTrue();
56+
assertThat(registeredTypes.contains(TypeReference.of(MistralAiChatOptions.class))).isTrue();
57+
assertThat(registeredTypes.contains(TypeReference.of(MistralAiEmbeddingOptions.class))).isTrue();
4358
}
4459

4560
}

‎models/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/aot/MoonshotRuntimeHints.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ public class MoonshotRuntimeHints implements RuntimeHintsRegistrar {
3636
@Override
3737
public void registerHints(@NonNull RuntimeHints hints, @Nullable ClassLoader classLoader) {
3838
var mcs = MemberCategory.values();
39-
for (var tr : findJsonAnnotatedClassesInPackage(MoonshotApi.class)) {
39+
for (var tr : findJsonAnnotatedClassesInPackage("org.springframework.ai.moonshot")) {
4040
hints.reflection().registerType(tr, mcs);
4141
}
4242
}

‎models/spring-ai-moonshot/src/test/java/org/springframework/ai/moonshot/aot/MoonshotRuntimeHintsTests.java‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,12 @@
1616

1717
package org.springframework.ai.moonshot.aot;
1818

19+
import java.util.HashSet;
1920
import java.util.Set;
2021

2122
import org.junit.jupiter.api.Test;
2223

24+
import org.springframework.ai.moonshot.MoonshotChatOptions;
2325
import org.springframework.ai.moonshot.api.MoonshotApi;
2426
import org.springframework.aot.hint.RuntimeHints;
2527
import org.springframework.aot.hint.TypeReference;
@@ -39,10 +41,21 @@ void registerHints() {
3941
MoonshotRuntimeHints moonshotRuntimeHints = new MoonshotRuntimeHints();
4042
moonshotRuntimeHints.registerHints(runtimeHints, null);
4143

42-
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(MoonshotApi.class);
44+
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage("org.springframework.ai.moonshot");
45+
46+
Set<TypeReference> registeredTypes = new HashSet<>();
47+
runtimeHints.reflection().typeHints().forEach(typeHint -> registeredTypes.add(typeHint.getType()));
48+
4349
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
44-
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
50+
assertThat(registeredTypes.contains(jsonAnnotatedClass)).isTrue();
4551
}
52+
53+
// Check a few more specific ones
54+
assertThat(registeredTypes.contains(TypeReference.of(MoonshotApi.ChatCompletion.class))).isTrue();
55+
assertThat(registeredTypes.contains(TypeReference.of(MoonshotApi.ChatCompletionRequest.class))).isTrue();
56+
assertThat(registeredTypes.contains(TypeReference.of(MoonshotApi.ChatCompletionChunk.class))).isTrue();
57+
assertThat(registeredTypes.contains(TypeReference.of(MoonshotApi.Usage.class))).isTrue();
58+
assertThat(registeredTypes.contains(TypeReference.of(MoonshotChatOptions.class))).isTrue();
4659
}
4760

4861
}

‎models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/aot/OllamaRuntimeHints.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,7 @@ public class OllamaRuntimeHints implements RuntimeHintsRegistrar {
3737
@Override
3838
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
3939
var mcs = MemberCategory.values();
40-
for (var tr : findJsonAnnotatedClassesInPackage(OllamaApi.class)) {
41-
hints.reflection().registerType(tr, mcs);
42-
}
43-
for (var tr : findJsonAnnotatedClassesInPackage(OllamaOptions.class)) {
40+
for (var tr : findJsonAnnotatedClassesInPackage("org.springframework.ai.ollama")) {
4441
hints.reflection().registerType(tr, mcs);
4542
}
4643
}

0 commit comments

Comments
 (0)