Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Move native hints into individual modules

Browse Source 
* Use aot.factories approach for registration
* Add tests
* final tweaks- Thanks Josh!
This commit is contained in:
Mark Pollack
2024-02-13 11:11:24 -05:00
parent ea7dce3833
commit 596f2b06c0
34 changed files with 609 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHints.java Normal file
View File

@@ -0,0 +1,45 @@
package org.springframework.ai.bedrock.aot;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
import org.springframework.aot.hint.MemberCategory;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.RuntimeHintsRegistrar;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
/**
* The BedrockRuntimeHints class is responsible for registering runtime hints for Bedrock
* AI API classes.
*
* @author Josh Long
* @author Christian Tzolov
* @author Mark Pollack
*/
public class BedrockRuntimeHints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
var mcs = MemberCategory.values();
for (var tr : findJsonAnnotatedClassesInPackage(Ai21Jurassic2ChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(CohereChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(CohereEmbeddingBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(Llama2ChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(TitanChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(TitanEmbeddingBedrockApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(AnthropicChatBedrockApi.class))
hints.reflection().registerType(tr, mcs);
}
}

2
models/spring-ai-bedrock/src/main/resources/META-INF/spring/aot.factories Normal file
View File

@@ -0,0 +1,2 @@
org.springframework.aot.hint.RuntimeHintsRegistrar=\
org.springframework.ai.bedrock.aot.BedrockRuntimeHints

43
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/aot/BedrockRuntimeHintsTests.java Normal file
View File

@@ -0,0 +1,43 @@
package org.springframework.ai.bedrock.aot;
import org.junit.jupiter.api.Test;
import org.springframework.ai.bedrock.anthropic.api.AnthropicChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereChatBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
import org.springframework.ai.bedrock.jurassic2.api.Ai21Jurassic2ChatBedrockApi;
import org.springframework.ai.bedrock.llama2.api.Llama2ChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanChatBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.TypeReference;
import java.util.List;
import java.util.Set;
import java.util.Arrays;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
class BedrockRuntimeHintsTests {
@Test
void registerHints() {
RuntimeHints runtimeHints = new RuntimeHints();
BedrockRuntimeHints bedrockRuntimeHints = new BedrockRuntimeHints();
bedrockRuntimeHints.registerHints(runtimeHints, null);
List<Class> classList = Arrays.asList(Ai21Jurassic2ChatBedrockApi.class, CohereChatBedrockApi.class,
CohereEmbeddingBedrockApi.class, Llama2ChatBedrockApi.class, TitanChatBedrockApi.class,
TitanEmbeddingBedrockApi.class, AnthropicChatBedrockApi.class);
for (Class aClass : classList) {
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(aClass);
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
}
}
}
}

30
models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/aot/OllamaRuntimeHints.java Normal file
View File

@@ -0,0 +1,30 @@
package org.springframework.ai.ollama.aot;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.aot.hint.MemberCategory;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.RuntimeHintsRegistrar;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
/**
* The OllamaRuntimeHints class is responsible for registering runtime hints for Ollama AI
* API classes.
*
* @author Josh Long
* @author Christian Tzolov
* @author Mark Pollack
*/
public class OllamaRuntimeHints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
var mcs = MemberCategory.values();
for (var tr : findJsonAnnotatedClassesInPackage(OllamaApi.class))
hints.reflection().registerType(tr, mcs);
for (var tr : findJsonAnnotatedClassesInPackage(OllamaOptions.class))
hints.reflection().registerType(tr, mcs);
}
}

2
models/spring-ai-ollama/src/main/resources/META-INF/spring/aot.factories Normal file
View File

@@ -0,0 +1,2 @@
org.springframework.aot.hint.RuntimeHintsRegistrar=\
org.springframework.ai.vertex.aot.OllamaRuntimeHints

34
models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/aot/OllamaRuntimeHintsTests.java Normal file
View File

@@ -0,0 +1,34 @@
package org.springframework.ai.ollama.aot;
import org.junit.jupiter.api.Test;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.TypeReference;
import java.util.Set;
import static org.assertj.core.api.AssertionsForClassTypes.*;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.*;
class OllamaRuntimeHintsTests {
@Test
void registerHints() {
RuntimeHints runtimeHints = new RuntimeHints();
OllamaRuntimeHints ollamaRuntimeHints = new OllamaRuntimeHints();
ollamaRuntimeHints.registerHints(runtimeHints, null);
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(OllamaApi.class);
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
}
jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(OllamaOptions.class);
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
}
}
}

27
models/spring-ai-openai/src/main/java/org/springframework/ai/openai/aot/OpenAiRuntimeHints.java Normal file
View File

@@ -0,0 +1,27 @@
package org.springframework.ai.openai.aot;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.aot.hint.MemberCategory;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.RuntimeHintsRegistrar;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
/**
* The OpenAiRuntimeHints class is responsible for registering runtime hints for OpenAI
* API classes.
*
* @author Josh Long
* @author Christian Tzolov
* @author Mark Pollack
*/
public class OpenAiRuntimeHints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
var mcs = MemberCategory.values();
for (var tr : findJsonAnnotatedClassesInPackage(OpenAiApi.class))
hints.reflection().registerType(tr, mcs);
}
}

2
models/spring-ai-openai/src/main/resources/META-INF/spring/aot.factories Normal file
View File

@@ -0,0 +1,2 @@
org.springframework.aot.hint.RuntimeHintsRegistrar=\
org.springframework.ai.openai.aot.OpenAiRuntimeHints

28
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/aot/OpenAiRuntimeHintsTests.java Normal file
View File

@@ -0,0 +1,28 @@
package org.springframework.ai.openai.aot;
import org.junit.jupiter.api.Test;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.TypeReference;
import java.util.Set;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
class OpenAiRuntimeHintsTests {
@Test
void registerHints() {
RuntimeHints runtimeHints = new RuntimeHints();
OpenAiRuntimeHints openAiRuntimeHints = new OpenAiRuntimeHints();
openAiRuntimeHints.registerHints(runtimeHints, null);
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(OpenAiApi.class);
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
}
}
}

27
models/spring-ai-vertex-ai/src/main/java/org/springframework/ai/vertex/aot/VertexRuntimeHints.java Normal file
View File

@@ -0,0 +1,27 @@
package org.springframework.ai.vertex.aot;
import org.springframework.ai.vertex.api.VertexAiApi;
import org.springframework.aot.hint.MemberCategory;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.RuntimeHintsRegistrar;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
/**
* The VertexRuntimeHints class is responsible for registering runtime hints for Vertex AI
* API classes.
*
* @author Josh Long
* @author Christian Tzolov
* @author Mark Pollack
*/
public class VertexRuntimeHints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
var mcs = MemberCategory.values();
for (var tr : findJsonAnnotatedClassesInPackage(VertexAiApi.class))
hints.reflection().registerType(tr, mcs);
}
}

2
models/spring-ai-vertex-ai/src/main/resources/META-INF/spring/aot.factories Normal file
View File

@@ -0,0 +1,2 @@
org.springframework.aot.hint.RuntimeHintsRegistrar=\
org.springframework.ai.vertex.aot.VertexRuntimeHints

27
models/spring-ai-vertex-ai/src/test/java/org/springframework/ai/vertex/aot/VertexRuntimeHintsTests.java Normal file
View File

@@ -0,0 +1,27 @@
package org.springframework.ai.vertex.aot;
import org.junit.jupiter.api.Test;
import org.springframework.ai.vertex.api.VertexAiApi;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.TypeReference;
import java.util.Set;
import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClassesInPackage;
import static org.springframework.aot.hint.predicate.RuntimeHintsPredicates.reflection;
class VertexRuntimeHintsTests {
@Test
void registerHints() {
RuntimeHints runtimeHints = new RuntimeHints();
VertexRuntimeHints vertexRuntimeHints = new VertexRuntimeHints();
vertexRuntimeHints.registerHints(runtimeHints, null);
Set<TypeReference> jsonAnnotatedClasses = findJsonAnnotatedClassesInPackage(VertexAiApi.class);
for (TypeReference jsonAnnotatedClass : jsonAnnotatedClasses) {
assertThat(runtimeHints).matches(reflection().onType(jsonAnnotatedClass));
}
}
}