From bc3005154863ea0db050747e60af6623389159cc Mon Sep 17 00:00:00 2001
From: Mark Pollack <mpollack@vmware.com>
Date: Wed, 23 Aug 2023 12:47:56 -0700
Subject: [PATCH] add optional ObjectMapper to BeanObjectMapper and allow for
 subclassing

---
 .../ai/parser/BeanOutputParser.java           | 21 ++++++++++++++--
 .../ai/autoconfigure/NativeHints.java         | 24 +++++++++++++++++++
 .../openai/OpenAiAutoConfiguration.java       |  3 +++
 3 files changed, 46 insertions(+), 2 deletions(-)
 create mode 100644 spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java

diff --git a/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java b/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java
index 1fbb176c8..ca8f41807 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/parser/BeanOutputParser.java
@@ -1,6 +1,7 @@
 package org.springframework.ai.parser;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.github.victools.jsonschema.generator.OptionPreset;
@@ -19,9 +20,20 @@ public class BeanOutputParser<T> implements OutputParser<T> {
 
 	private Class clazz;
 
+	private ObjectMapper objectMapper;
+
 	public BeanOutputParser(Class<T> clazz) {
 		Objects.requireNonNull(clazz, "Java Class can not be null;");
 		this.clazz = clazz;
+		this.objectMapper = getObjectMapper();
+		generateSchema();
+	}
+
+	public BeanOutputParser(Class<T> clazz, ObjectMapper objectMapper) {
+		Objects.requireNonNull(clazz, "Java Class can not be null;");
+		Objects.requireNonNull(objectMapper, "ObjectMapper can not be null;");
+		this.clazz = clazz;
+		this.objectMapper = objectMapper;
 		generateSchema();
 	}
 
@@ -36,15 +48,20 @@ public class BeanOutputParser<T> implements OutputParser<T> {
 
 	@Override
 	public T parse(String text) {
-		ObjectMapper objectMapper = new ObjectMapper();
 		try {
-			return (T) objectMapper.readValue(text, this.clazz);
+			return (T) this.objectMapper.readValue(text, this.clazz);
 		}
 		catch (JsonProcessingException e) {
 			throw new RuntimeException(e);
 		}
 	}
 
+	protected ObjectMapper getObjectMapper() {
+		ObjectMapper objectMapper = new ObjectMapper();
+		objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+		return objectMapper;
+	}
+
 	@Override
 	public String getFormat() {
 
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
new file mode 100644
index 000000000..171181b33
--- /dev/null
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/NativeHints.java
@@ -0,0 +1,24 @@
+package org.springframework.ai.autoconfigure;
+
+import org.springframework.aot.hint.MemberCategory;
+import org.springframework.aot.hint.RuntimeHints;
+import org.springframework.aot.hint.TypeReference;
+
+import java.util.Set;
+
+import org.springframework.aot.hint.RuntimeHintsRegistrar;
+
+public class NativeHints implements RuntimeHintsRegistrar {
+
+	@Override
+	public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
+		hints.proxies().registerJdkProxy(TypeReference.of("com.theokanning.openai.OpenAiApi"));
+		for (var className : Set.of("com.theokanning.openai.Usage",
+				"com.theokanning.openai.completion.chat.ChatCompletionChoice",
+				"com.theokanning.openai.completion.chat.ChatCompletionRequest",
+				"com.theokanning.openai.completion.chat.ChatCompletionResult",
+				"com.theokanning.openai.completion.chat.ChatMessage"))
+			hints.reflection().registerType(TypeReference.of(className), MemberCategory.values());
+	}
+
+}
\ No newline at end of file
diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java
index 55d751b23..9acfa19db 100644
--- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java
+++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java
@@ -18,6 +18,7 @@ package org.springframework.ai.autoconfigure.openai;
 
 import com.theokanning.openai.service.OpenAiService;
 
+import org.springframework.ai.autoconfigure.NativeHints;
 import org.springframework.ai.embedding.EmbeddingClient;
 import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
 import org.springframework.ai.openai.client.OpenAiClient;
@@ -25,6 +26,7 @@ import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.ImportRuntimeHints;
 import org.springframework.util.StringUtils;
 
 import static org.springframework.ai.autoconfigure.openai.OpenAiProperties.CONFIG_PREFIX;
@@ -32,6 +34,7 @@ import static org.springframework.ai.autoconfigure.openai.OpenAiProperties.CONFI
 @AutoConfiguration
 @ConditionalOnClass(OpenAiService.class)
 @EnableConfigurationProperties(OpenAiProperties.class)
+@ImportRuntimeHints(NativeHints.class)
 public class OpenAiAutoConfiguration {
 
 	private final OpenAiProperties openAiProperties;