vertex ai aot hints, plus a fix to localize the dependency on spring webflux

This commit is contained in:
Josh Long
2023-12-17 17:28:26 -08:00
committed by Christian Tzolov
parent c62fee3635
commit 95cbe54bae
4 changed files with 49 additions and 27 deletions

View File

@@ -40,11 +40,11 @@
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
<version>${spring-framework.version}</version>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-core</artifactId>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-messaging</artifactId>

View File

@@ -60,17 +60,17 @@
<groupId>org.springframework</groupId>
<artifactId>spring-context-support</artifactId>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-logging</artifactId>
</dependency>
<dependency>
<groupId>io.projectreactor</groupId>
<artifactId>reactor-core</artifactId>
<groupId>org.springframework</groupId>
<artifactId>spring-webflux</artifactId>
<version>${spring-framework.version}</version>
</dependency>
<!-- test dependencies -->
<dependency>
<groupId>org.springframework.ai</groupId>

View File

@@ -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<TypeReference> 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<TypeReference> 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();

View File

@@ -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 {