first cut of aot improvements

Signed-off-by: Josh Long <54473+joshlong@users.noreply.github.com>
This commit is contained in:
Josh Long
2025-04-07 19:31:43 -07:00
committed by Ilayaperumal Gopinathan
parent a43cdc8f79
commit fbec267eca
4 changed files with 36 additions and 21 deletions

View File

@@ -50,6 +50,8 @@ import org.springframework.util.Assert;
*
* @author Christian Tzolov
* @author Thomas Vitale
* @author Josh Long
*
*/
public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {

View File

@@ -16,15 +16,13 @@
package org.springframework.ai.openai.aot;
import java.util.Set;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.api.OpenAiAudioApi;
import org.springframework.ai.openai.api.OpenAiImageApi;
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.lang.NonNull;
import org.springframework.lang.Nullable;
@@ -40,23 +38,22 @@ import static org.springframework.ai.aot.AiRuntimeHints.findJsonAnnotatedClasses
*/
public class OpenAiRuntimeHints implements RuntimeHintsRegistrar {
private static Set<TypeReference> eval(Set<TypeReference> referenceSet) {
referenceSet.forEach(tr -> System.out.println(tr.toString()));
return referenceSet;
}
@Override
public void registerHints(@NonNull RuntimeHints hints, @Nullable ClassLoader classLoader) {
var mcs = MemberCategory.values();
for (var tr : eval(findJsonAnnotatedClassesInPackage(OpenAiApi.class))) {
for (var tr : (findJsonAnnotatedClassesInPackage(OpenAiChatOptions.class))) {
hints.reflection().registerType(tr, mcs);
}
for (var tr : eval(findJsonAnnotatedClassesInPackage(OpenAiAudioApi.class))) {
for (var tr : (findJsonAnnotatedClassesInPackage(OpenAiApi.class))) {
hints.reflection().registerType(tr, mcs);
}
for (var tr : eval(findJsonAnnotatedClassesInPackage(OpenAiImageApi.class))) {
for (var tr : (findJsonAnnotatedClassesInPackage(OpenAiAudioApi.class))) {
hints.reflection().registerType(tr, mcs);
}
for (var tr : findJsonAnnotatedClassesInPackage(OpenAiImageApi.class)) {
hints.reflection().registerType(tr, mcs);
}
}
}

View File

@@ -16,28 +16,41 @@
package org.springframework.ai.embedding;
import org.springframework.aot.hint.RuntimeHints;
import org.springframework.aot.hint.RuntimeHintsRegistrar;
import org.springframework.context.annotation.ImportRuntimeHints;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;
import org.springframework.util.Assert;
import java.io.IOException;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import org.springframework.core.io.DefaultResourceLoader;
/**
* Abstract implementation of the {@link EmbeddingModel} interface that provides
* dimensions calculation caching.
*
* @author Christian Tzolov
* @author Josh Long
*/
@ImportRuntimeHints(AbstractEmbeddingModel.Hints.class)
public abstract class AbstractEmbeddingModel implements EmbeddingModel {
private static final Resource EMBEDDING_MODEL_DIMENSIONS_PROPERTIES = new ClassPathResource(
"/embedding/embedding-model-dimensions.properties");
private static final Map<String, Integer> KNOWN_EMBEDDING_DIMENSIONS = loadKnownModelDimensions();
/**
* Default constructor.
*/
public AbstractEmbeddingModel() {
static class Hints implements RuntimeHintsRegistrar {
@Override
public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
hints.resources().registerResource(EMBEDDING_MODEL_DIMENSIONS_PROPERTIES);
}
}
/**
@@ -69,10 +82,13 @@ public abstract class AbstractEmbeddingModel implements EmbeddingModel {
private static Map<String, Integer> loadKnownModelDimensions() {
try {
Properties properties = new Properties();
properties.load(new DefaultResourceLoader()
.getResource("classpath:/embedding/embedding-model-dimensions.properties")
.getInputStream());
var resource = EMBEDDING_MODEL_DIMENSIONS_PROPERTIES;
Assert.notNull(resource, "the embedding dimensions must be non-null");
Assert.state(resource.exists(), "the embedding dimensions properties file must exist");
var properties = new Properties();
try (var in = resource.getInputStream()) {
properties.load(in);
}
return properties.entrySet()
.stream()
.collect(Collectors.toMap(e -> e.getKey().toString(), e -> Integer.parseInt(e.getValue().toString())));