first cut of aot improvements

Signed-off-by: Josh Long <[email protected]>
Signed-off-by: Theo van Kraay <[email protected]>

Author: joshlong

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiEmbeddingModel.java

@@ -50,6 +50,8 @@
  *
  * @author Christian Tzolov
  * @author Thomas Vitale
+ * @author Josh Long
+ *
  */
 public class OpenAiEmbeddingModel extends AbstractEmbeddingModel {
 

models/spring-ai-openai/src/main/java/org/springframework/ai/openai/aot/OpenAiRuntimeHints.java

@@ -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 @@
  */
 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);
+		}
+
 	}
 
 }

spring-ai-model/src/main/java/org/springframework/ai/embedding/AbstractEmbeddingModel.java

@@ -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 static int dimensions(EmbeddingModel embeddingModel, String modelName, St
 
 	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())));