From 95cbe54bae211da8e2aee0640ae6ff3e96bada1e Mon Sep 17 00:00:00 2001 From: Josh Long Date: Sun, 17 Dec 2023 17:28:26 -0800 Subject: [PATCH] vertex ai aot hints, plus a fix to localize the dependency on spring webflux --- spring-ai-core/pom.xml | 6 +- spring-ai-openai/pom.xml | 8 +-- .../ai/autoconfigure/NativeHints.java | 59 ++++++++++++------- .../vertexai/VertexAiAutoConfiguration.java | 3 + 4 files changed, 49 insertions(+), 27 deletions(-) diff --git a/spring-ai-core/pom.xml b/spring-ai-core/pom.xml index 6cd1c4617..bf7da5411 100644 --- a/spring-ai-core/pom.xml +++ b/spring-ai-core/pom.xml @@ -40,11 +40,11 @@ - org.springframework - spring-webflux - ${spring-framework.version} + io.projectreactor + reactor-core + org.springframework spring-messaging diff --git a/spring-ai-openai/pom.xml b/spring-ai-openai/pom.xml index 21bfe14a7..7aeaa0c52 100644 --- a/spring-ai-openai/pom.xml +++ b/spring-ai-openai/pom.xml @@ -60,17 +60,17 @@ org.springframework spring-context-support - org.springframework.boot spring-boot-starter-logging - - io.projectreactor - reactor-core + org.springframework + spring-webflux + ${spring-framework.version} + org.springframework.ai diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java index bd420cca6..e4cff74d3 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java @@ -1,23 +1,24 @@ package org.springframework.ai.autoconfigure; +import com.fasterxml.jackson.annotation.JsonInclude; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.springframework.ai.openai.api.OpenAiApi; +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 org.springframework.aot.hint.TypeReference; +import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider; +import org.springframework.core.io.ClassPathResource; +import org.springframework.core.io.support.PathMatchingResourcePatternResolver; +import org.springframework.core.type.filter.AnnotationTypeFilter; import java.io.IOException; import java.util.Objects; import java.util.Set; import java.util.stream.Collectors; -import org.springframework.core.type.filter.AnnotationTypeFilter; -import org.springframework.aot.hint.RuntimeHintsRegistrar; -import org.springframework.context.annotation.ClassPathScanningCandidateComponentProvider; -import org.springframework.core.io.ClassPathResource; -import org.springframework.core.io.support.PathMatchingResourcePatternResolver; - -import com.fasterxml.jackson.annotation.JsonInclude; - /*** * Native hints * @@ -25,26 +26,44 @@ import com.fasterxml.jackson.annotation.JsonInclude; */ public class NativeHints implements RuntimeHintsRegistrar { + static final Logger log = LoggerFactory.getLogger(NativeHints.class); + @Override public void registerHints(RuntimeHints hints, ClassLoader classLoader) { - new KnuddelsHints().registerHints(hints, classLoader); - new PdfReaderHints().registerHints(hints, classLoader); - new OpenAiHints().registerHints(hints, classLoader); + + for (var h : Set.of(new VertexAiHints(), new OpenAiHints(), new PdfReaderHints(), new KnuddelsHints())) + h.registerHints(hints, classLoader); + hints.resources().registerResource(new ClassPathResource("embedding/embedding-model-dimensions.properties")); } - static class OpenAiHints implements RuntimeHintsRegistrar { + private static Set findJsonAnnotatedClasses(Class packageClass) { + var packageName = packageClass.getPackageName(); + var classPathScanningCandidateComponentProvider = new ClassPathScanningCandidateComponentProvider(false); + classPathScanningCandidateComponentProvider.addIncludeFilter(new AnnotationTypeFilter(JsonInclude.class)); + return classPathScanningCandidateComponentProvider.findCandidateComponents(packageName) + .stream() + .map(bd -> TypeReference.of(Objects.requireNonNull(bd.getBeanClassName()))) + .peek(tr -> { + if (log.isDebugEnabled()) + log.debug("registering [" + tr.getName() + ']'); + }) + .collect(Collectors.toUnmodifiableSet()); + } - private static Set findJsonAnnotatedClasses(Class packageClass) { - var packageName = packageClass.getPackageName(); - var classPathScanningCandidateComponentProvider = new ClassPathScanningCandidateComponentProvider(false); - classPathScanningCandidateComponentProvider.addIncludeFilter(new AnnotationTypeFilter(JsonInclude.class)); - return classPathScanningCandidateComponentProvider.findCandidateComponents(packageName) - .stream() - .map(bd -> TypeReference.of(Objects.requireNonNull(bd.getBeanClassName()))) - .collect(Collectors.toUnmodifiableSet()); + static class VertexAiHints implements RuntimeHintsRegistrar { + + @Override + public void registerHints(RuntimeHints hints, ClassLoader classLoader) { + var mcs = MemberCategory.values(); + for (var tr : findJsonAnnotatedClasses(VertexAiApi.class)) + hints.reflection().registerType(tr, mcs); } + } + + static class OpenAiHints implements RuntimeHintsRegistrar { + @Override public void registerHints(RuntimeHints hints, ClassLoader classLoader) { var mcs = MemberCategory.values(); diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java index 4a08c3221..8cab8410d 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/vertexai/VertexAiAutoConfiguration.java @@ -16,6 +16,7 @@ package org.springframework.ai.autoconfigure.vertexai; +import org.springframework.ai.autoconfigure.NativeHints; import org.springframework.ai.vertex.api.VertexAiApi; import org.springframework.ai.vertex.embedding.VertexAiEmbeddingClient; import org.springframework.ai.vertex.generation.VertexAiChatClient; @@ -24,10 +25,12 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; import org.springframework.boot.context.properties.EnableConfigurationProperties; import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.ImportRuntimeHints; import org.springframework.web.client.RestClient; @AutoConfiguration @ConditionalOnClass(VertexAiApi.class) +@ImportRuntimeHints(NativeHints.class) @EnableConfigurationProperties({ VertexAiConnectionProperties.class, VertexAiChatProperties.class, VertexAiEmbeddingProperties.class }) public class VertexAiAutoConfiguration {