From 3a3439edc44df8c669922d4392942b68c30ccd6e Mon Sep 17 00:00:00 2001 From: Bart Veenstra <(bart.veenstra@gmail.com)> Date: Mon, 10 Mar 2025 14:14:39 +0100 Subject: [PATCH] feat: enhance AzureOpenAiResponseFormat to support JSON schema and builder pattern Signed-off-by: Bart Veenstra --- .../ai/azure/openai/AzureOpenAiChatModel.java | 14 +- .../openai/AzureOpenAiResponseFormat.java | 246 +++++++++++++++++- .../AzureChatCompletionsOptionsTests.java | 5 +- .../openai/AzureOpenAiChatOptionsTests.java | 12 +- 4 files changed, 267 insertions(+), 10 deletions(-) diff --git a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java index 78b482227..a94115a9b 100644 --- a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java +++ b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java @@ -16,6 +16,8 @@ package org.springframework.ai.azure.openai; +import com.azure.ai.openai.models.ChatCompletionsJsonSchemaResponseFormat; +import com.azure.ai.openai.models.ChatCompletionsJsonSchemaResponseFormatJsonSchema; import java.util.ArrayList; import java.util.Base64; import java.util.Collections; @@ -59,6 +61,8 @@ import io.micrometer.observation.ObservationRegistry; import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import org.springframework.ai.azure.openai.AzureOpenAiResponseFormat.JsonSchema; +import org.springframework.ai.azure.openai.AzureOpenAiResponseFormat.Type; import reactor.core.publisher.Flux; import reactor.core.scheduler.Schedulers; @@ -115,6 +119,7 @@ import org.springframework.util.CollectionUtils; * @author Alexandros Pappas * @author Berjan Jonker * @author Andres da Silva Santos + * @author Bart Veenstra * @see ChatModel * @see com.azure.ai.openai.OpenAIClient * @since 1.0.0 @@ -918,9 +923,16 @@ public class AzureOpenAiChatModel implements ChatModel { * @return Azure response format */ private ChatCompletionsResponseFormat toAzureResponseFormat(AzureOpenAiResponseFormat responseFormat) { - if (responseFormat == AzureOpenAiResponseFormat.JSON) { + if (responseFormat.getType() == Type.JSON_OBJECT) { return new ChatCompletionsJsonResponseFormat(); } + if (responseFormat.getType() == Type.JSON_SCHEMA) { + JsonSchema jsonSchema = responseFormat.getJsonSchema(); + var responseFormatJsonSchema = new ChatCompletionsJsonSchemaResponseFormatJsonSchema(jsonSchema.getName()); + String jsonString = ModelOptionsUtils.toJsonString(jsonSchema.getSchema()); + responseFormatJsonSchema.setSchema(BinaryData.fromString(jsonString)); + return new ChatCompletionsJsonSchemaResponseFormat(responseFormatJsonSchema); + } return new ChatCompletionsTextResponseFormat(); } diff --git a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiResponseFormat.java b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiResponseFormat.java index b6f71dc41..b3dc0483c 100644 --- a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiResponseFormat.java +++ b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiResponseFormat.java @@ -16,16 +16,23 @@ package org.springframework.ai.azure.openai; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonInclude.Include; +import com.fasterxml.jackson.annotation.JsonProperty; +import java.util.Map; +import java.util.Objects; +import org.springframework.ai.model.ModelOptionsUtils; +import org.springframework.util.StringUtils; + /** * Utility enumeration for representing the response format that may be requested from the * Azure OpenAI model. Please check OpenAI * API documentation for more details. */ -public enum AzureOpenAiResponseFormat { +@JsonInclude(Include.NON_NULL) +public class AzureOpenAiResponseFormat { - // default value used by OpenAI - TEXT, /* * From the OpenAI API documentation: Compatibility: Compatible with GPT-4 Turbo and * all GPT-3.5 Turbo models newer than gpt-3.5-turbo-1106. Caveats: This enables JSON @@ -36,7 +43,238 @@ public enum AzureOpenAiResponseFormat { * long-running and seemingly "stuck" request. Also note that the message content may * be partially cut off if finish_reason="length", which indicates the generation * exceeded max_tokens or the conversation exceeded the max context length. + * + * Type Must be one of 'text', 'json_object' or 'json_schema'. */ - JSON + @JsonProperty("type") + private Type type; + + /** + * JSON schema object that describes the format of the JSON object. Only applicable + * when type is 'json_schema'. + */ + @JsonProperty("json_schema") + private JsonSchema jsonSchema = null; + + private String schema; + + public AzureOpenAiResponseFormat() { + + } + + public Type getType() { + return this.type; + } + + public void setType(Type type) { + this.type = type; + } + + public JsonSchema getJsonSchema() { + return this.jsonSchema; + } + + public void setJsonSchema(JsonSchema jsonSchema) { + this.jsonSchema = jsonSchema; + } + + public String getSchema() { + return this.schema; + } + + public void setSchema(String schema) { + this.schema = schema; + if (schema != null) { + this.jsonSchema = JsonSchema.builder().schema(schema).strict(true).build(); + } + } + + private AzureOpenAiResponseFormat(Type type, JsonSchema jsonSchema) { + this.type = type; + this.jsonSchema = jsonSchema; + } + + public AzureOpenAiResponseFormat(Type type, String schema) { + this(type, StringUtils.hasText(schema) ? JsonSchema.builder().schema(schema).strict(true).build() : null); + } + + public static Builder builder() { + return new Builder(); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + AzureOpenAiResponseFormat that = (AzureOpenAiResponseFormat) o; + return this.type == that.type && Objects.equals(this.jsonSchema, that.jsonSchema); + } + + @Override + public int hashCode() { + return Objects.hash(this.type, this.jsonSchema); + } + + @Override + public String toString() { + return "ResponseFormat{" + "type=" + this.type + ", jsonSchema=" + this.jsonSchema + '}'; + } + + public static final class Builder { + + private Type type; + + private JsonSchema jsonSchema; + + private Builder() { + } + + public Builder type(Type type) { + this.type = type; + return this; + } + + public Builder jsonSchema(JsonSchema jsonSchema) { + this.jsonSchema = jsonSchema; + return this; + } + + public Builder jsonSchema(String jsonSchema) { + this.jsonSchema = JsonSchema.builder().schema(jsonSchema).build(); + return this; + } + + public AzureOpenAiResponseFormat build() { + return new AzureOpenAiResponseFormat(this.type, this.jsonSchema); + } + + } + + public enum Type { + + /** + * Generates a text response. (default) + */ + @JsonProperty("text") + TEXT, + + /** + * Enables JSON mode, which guarantees the message the model generates is valid + * JSON. + */ + @JsonProperty("json_object") + JSON_OBJECT, + + /** + * Enables Structured Outputs which guarantees the model will match your supplied + * JSON schema. + */ + @JsonProperty("json_schema") + JSON_SCHEMA + + } + + /** + * JSON schema object that describes the format of the JSON object. Applicable for the + * 'json_schema' type only. + */ + @JsonInclude(Include.NON_NULL) + public static class JsonSchema { + + @JsonProperty("name") + private String name; + + @JsonProperty("schema") + private Map schema; + + @JsonProperty("strict") + private Boolean strict; + + public JsonSchema() { + + } + + public String getName() { + return this.name; + } + + public Map getSchema() { + return this.schema; + } + + public Boolean getStrict() { + return this.strict; + } + + private JsonSchema(String name, Map schema, Boolean strict) { + this.name = name; + this.schema = schema; + this.strict = strict; + } + + public static Builder builder() { + return new Builder(); + } + + @Override + public int hashCode() { + return Objects.hash(this.name, this.schema, this.strict); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + JsonSchema that = (JsonSchema) o; + return Objects.equals(this.name, that.name) && Objects.equals(this.schema, that.schema) + && Objects.equals(this.strict, that.strict); + } + + public static final class Builder { + + private String name = "custom_schema"; + + private Map schema; + + private Boolean strict = true; + + private Builder() { + } + + public Builder name(String name) { + this.name = name; + return this; + } + + public Builder schema(Map schema) { + this.schema = schema; + return this; + } + + public Builder schema(String schema) { + this.schema = ModelOptionsUtils.jsonToMap(schema); + return this; + } + + public Builder strict(Boolean strict) { + this.strict = strict; + return this; + } + + public JsonSchema build() { + return new JsonSchema(this.name, this.schema, this.strict); + } + + } + + } } diff --git a/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureChatCompletionsOptionsTests.java b/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureChatCompletionsOptionsTests.java index 46dcf5547..2c13ced56 100644 --- a/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureChatCompletionsOptionsTests.java +++ b/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureChatCompletionsOptionsTests.java @@ -30,6 +30,7 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.mockito.Mockito; +import org.springframework.ai.azure.openai.AzureOpenAiResponseFormat.Type; import org.springframework.ai.chat.prompt.Prompt; import static org.assertj.core.api.Assertions.assertThat; @@ -68,7 +69,7 @@ public class AzureChatCompletionsOptionsTests { .logprobs(true) .topLogprobs(5) .enhancements(mockAzureChatEnhancementConfiguration) - .responseFormat(AzureOpenAiResponseFormat.TEXT) + .responseFormat(AzureOpenAiResponseFormat.builder().type(Type.TEXT).build()) .build(); var client = AzureOpenAiChatModel.builder() @@ -114,7 +115,7 @@ public class AzureChatCompletionsOptionsTests { .logprobs(true) .topLogprobs(4) .enhancements(anotherMockAzureChatEnhancementConfiguration) - .responseFormat(AzureOpenAiResponseFormat.JSON) + .responseFormat(AzureOpenAiResponseFormat.builder().type(Type.JSON_OBJECT).build()) .build(); requestOptions = client.toAzureChatCompletionsOptions(new Prompt("Test message content", runtimeOptions)); diff --git a/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureOpenAiChatOptionsTests.java b/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureOpenAiChatOptionsTests.java index 524686432..60568f540 100644 --- a/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureOpenAiChatOptionsTests.java +++ b/models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/AzureOpenAiChatOptionsTests.java @@ -36,7 +36,9 @@ class AzureOpenAiChatOptionsTests { @Test void testBuilderWithAllFields() { - AzureOpenAiResponseFormat responseFormat = AzureOpenAiResponseFormat.TEXT; + AzureOpenAiResponseFormat responseFormat = AzureOpenAiResponseFormat.builder() + .type(AzureOpenAiResponseFormat.Type.TEXT) + .build(); ChatCompletionStreamOptions streamOptions = new ChatCompletionStreamOptions(); streamOptions.setIncludeUsage(true); @@ -75,7 +77,9 @@ class AzureOpenAiChatOptionsTests { @Test void testCopy() { - AzureOpenAiResponseFormat responseFormat = AzureOpenAiResponseFormat.TEXT; + AzureOpenAiResponseFormat responseFormat = AzureOpenAiResponseFormat.builder() + .type(AzureOpenAiResponseFormat.Type.TEXT) + .build(); ChatCompletionStreamOptions streamOptions = new ChatCompletionStreamOptions(); streamOptions.setIncludeUsage(true); @@ -113,7 +117,9 @@ class AzureOpenAiChatOptionsTests { @Test void testSetters() { - AzureOpenAiResponseFormat responseFormat = AzureOpenAiResponseFormat.TEXT; + AzureOpenAiResponseFormat responseFormat = AzureOpenAiResponseFormat.builder() + .type(AzureOpenAiResponseFormat.Type.TEXT) + .build(); ChatCompletionStreamOptions streamOptions = new ChatCompletionStreamOptions(); streamOptions.setIncludeUsage(true); AzureChatEnhancementConfiguration enhancements = new AzureChatEnhancementConfiguration();