From c1fc687730eaeeed2e4727bbab8b22f27a96dbb6 Mon Sep 17 00:00:00 2001
From: Ilayaperumal Gopinathan <ilayaperumal.gopinathan@broadcom.com>
Date: Fri, 8 Nov 2024 13:34:07 +0000
Subject: [PATCH] Update OpenAiChatModelResponseFormatIT with JSON order
 usecase

 - Verify that the BeanOutputConverter converts with the right order as specified in the JSON schema

Update the BeanOutputConverterTest's test case to tweak the order for validation
---
 .../chat/OpenAiChatModelResponseFormatIT.java | 55 +++++++++++++++++--
 .../ai/converter/BeanOutputConverterTest.java |  4 +-
 2 files changed, 53 insertions(+), 6 deletions(-)

diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelResponseFormatIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelResponseFormatIT.java
index 3f24a24b3..39da40ea3 100644
--- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelResponseFormatIT.java
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelResponseFormatIT.java
@@ -17,11 +17,13 @@
 package org.springframework.ai.openai.chat;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.annotation.JsonPropertyOrder;
 import com.fasterxml.jackson.core.JacksonException;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.JsonMappingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.Assert;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
 import org.slf4j.Logger;
@@ -40,10 +42,12 @@ import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.context.annotation.Bean;
 
+import static org.assertj.core.api.Assertions.as;
 import static org.assertj.core.api.Assertions.assertThat;
 
 /**
  * @author Christian Tzolov
+ * @author Ilayaperumal Gopinathan
  */
 @SpringBootTest(classes = OpenAiChatModelResponseFormatIT.Config.class)
 @EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
@@ -140,11 +144,13 @@ public class OpenAiChatModelResponseFormatIT {
 	@Test
 	void jsonSchemaBeanConverter() throws JsonMappingException, JsonProcessingException {
 
+		@JsonPropertyOrder({ "steps", "final_answer" })
 		record MathReasoning(@JsonProperty(required = true, value = "steps") Steps steps,
 				@JsonProperty(required = true, value = "final_answer") String finalAnswer) {
 
 			record Steps(@JsonProperty(required = true, value = "items") Items[] items) {
 
+				@JsonPropertyOrder({ "output", "explanation" })
 				record Items(@JsonProperty(required = true, value = "explanation") String explanation,
 						@JsonProperty(required = true, value = "output") String output) {
 
@@ -156,9 +162,45 @@ public class OpenAiChatModelResponseFormatIT {
 
 		var outputConverter = new BeanOutputConverter<>(MathReasoning.class);
 
+		var expectedJsonSchema = """
+				{
+				  "$schema" : "https://json-schema.org/draft/2020-12/schema",
+				  "type" : "object",
+				  "properties" : {
+				    "steps" : {
+				      "type" : "object",
+				      "properties" : {
+				        "items" : {
+				          "type" : "array",
+				          "items" : {
+				            "type" : "object",
+				            "properties" : {
+				              "output" : {
+				                "type" : "string"
+				              },
+				              "explanation" : {
+				                "type" : "string"
+				              }
+				            },
+				            "required" : [ "output", "explanation" ],
+				            "additionalProperties" : false
+				          }
+				        }
+				      },
+				      "required" : [ "items" ],
+				      "additionalProperties" : false
+				    },
+				    "final_answer" : {
+				      "type" : "string"
+				    }
+				  },
+				  "required" : [ "steps", "final_answer" ],
+				  "additionalProperties" : false
+				}""";
 		var jsonSchema1 = outputConverter.getJsonSchema();
 
-		System.out.println(jsonSchema1);
+		assertThat(jsonSchema1).isNotNull();
+		assertThat(jsonSchema1).isEqualTo(expectedJsonSchema);
 
 		Prompt prompt = new Prompt("how can I solve 8x + 7 = -23",
 				OpenAiChatOptions.builder()
@@ -174,11 +216,16 @@ public class OpenAiChatModelResponseFormatIT {
 
 		logger.info("Response content: {}", content);
 
+		assertThat(isValidJson(content)).isTrue();
+
+		// Check if the order is correct as specified in the schema. Steps should come
+		// first before final answer.
+		assertThat(content.startsWith("{\"steps\":{\"items\":["));
+
 		MathReasoning mathReasoning = outputConverter.convert(content);
 
-		System.out.println(mathReasoning);
-
-		assertThat(isValidJson(content)).isTrue();
+		assertThat(mathReasoning).isNotNull();
+		logger.info(mathReasoning.toString());
 	}
 
 	@SpringBootConfiguration
diff --git a/spring-ai-core/src/test/java/org/springframework/ai/converter/BeanOutputConverterTest.java b/spring-ai-core/src/test/java/org/springframework/ai/converter/BeanOutputConverterTest.java
index 3653d06cf..ec3f5c5f9 100644
--- a/spring-ai-core/src/test/java/org/springframework/ai/converter/BeanOutputConverterTest.java
+++ b/spring-ai-core/src/test/java/org/springframework/ai/converter/BeanOutputConverterTest.java
@@ -120,9 +120,9 @@ class BeanOutputConverterTest {
 	record TestClassWithJsonPropertyOrder(
 			@JsonProperty("string_property") @JsonPropertyDescription("string_property_description") String someString,
 
-			@JsonProperty(required = true, value = "foo_property") String foo,
+			@JsonProperty(required = true, value = "bar_property") String bar,
 
-			@JsonProperty(required = true, value = "bar_property") String bar) {
+			@JsonProperty(required = true, value = "foo_property") String foo) {
 	}
 
 	@Nested