Allow custom StructuredOutputConverter(s) to participate in Native Structured Output

spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/DefaultChatClient.java

@@ -432,10 +432,10 @@ protected <T> ResponseEntity<ChatResponse, T> doResponseEntity(StructuredOutputC
 
 			this.request.context().put(ChatClientAttributes.OUTPUT_FORMAT.getKey(), outputConverter.getFormat());
 
-			if (Boolean.TRUE.equals(this.request.context().get(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey()))
-					&& outputConverter instanceof BeanOutputConverter beanOutputConverter) {
+			if (Boolean.TRUE
+				.equals(this.request.context().get(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey()))) {
 				this.request.context()
-					.put(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey(), beanOutputConverter.getJsonSchema());
+					.put(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey(), outputConverter.getJsonSchema());
 			}
 
 			var chatResponse = doGetObservableChatClientResponse(this.request).chatResponse();
@@ -474,12 +474,12 @@ protected <T> ResponseEntity<ChatResponse, T> doResponseEntity(StructuredOutputC
 				this.request.context().put(ChatClientAttributes.OUTPUT_FORMAT.getKey(), outputConverter.getFormat());
 			}
 
-			if (Boolean.TRUE.equals(this.request.context().get(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey()))
-					&& outputConverter instanceof BeanOutputConverter beanOutputConverter) {
+			if (Boolean.TRUE
+				.equals(this.request.context().get(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey()))) {
 				// Used for native structured output support, e.g. AI model API should
 				// provide structured output support.
 				this.request.context()
-					.put(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey(), beanOutputConverter.getJsonSchema());
+					.put(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey(), outputConverter.getJsonSchema());
 
 			}
 

spring-ai-client-chat/src/main/java/org/springframework/ai/chat/client/advisor/ChatModelCallAdvisor.java

@@ -16,7 +16,7 @@
 
 package org.springframework.ai.chat.client.advisor;
 
-import java.util.Map;
+import java.util.HashMap;
 
 import org.jspecify.annotations.Nullable;
 
@@ -59,7 +59,7 @@ public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAd
 
 		return ChatClientResponse.builder()
 			.chatResponse(chatResponse)
-			.context(Map.copyOf(formattedChatClientRequest.context()))
+			.context(new HashMap<>(formattedChatClientRequest.context()))
 			.build();
 	}
 
@@ -88,10 +88,7 @@ private static ChatClientRequest augmentWithFormatInstructions(ChatClientRequest
 				.text(userMessage.getText() + System.lineSeparator() + outputFormat)
 				.build());
 
-		return ChatClientRequest.builder()
-			.prompt(augmentedPrompt)
-			.context(Map.copyOf(chatClientRequest.context()))
-			.build();
+		return chatClientRequest.mutate().prompt(augmentedPrompt).build();
 	}
 
 	@Override

spring-ai-client-chat/src/test/java/org/springframework/ai/chat/client/ChatClientNativeStructuredResponseTests.java

@@ -16,18 +16,20 @@
 
 package org.springframework.ai.chat.client;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
 
 import net.javacrumbs.jsonunit.assertj.JsonAssertions;
 import net.javacrumbs.jsonunit.core.Option;
+import org.jspecify.annotations.NonNull;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Captor;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import tools.jackson.databind.JsonNode;
 
 import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
 import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
@@ -39,7 +41,9 @@
 import org.springframework.ai.chat.model.ChatResponse;
 import org.springframework.ai.chat.model.Generation;
 import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.converter.StructuredOutputConverter;
 import org.springframework.ai.model.tool.StructuredOutputChatOptions;
+import org.springframework.ai.util.json.JsonParser;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -325,12 +329,199 @@ public void dynamicDisableNativeEntityTest() {
 		verify(this.structuredOutputChatOptions, never()).setOutputSchema(anyString());
 	}
 
+	@Test
+	public void nativeWithCustomStructuredOutputConverterResponseEntityTest(
+			@Captor ArgumentCaptor<String> outputSchemaCaptor) {
+
+		ChatResponseMetadata metadata = ChatResponseMetadata.builder().keyValue("key1", "value1").build();
+
+		var chatResponse = new ChatResponse(List.of(new Generation(new AssistantMessage("""
+				{"name":"John", "age":30}
+				"""))), metadata);
+
+		given(this.chatModel.call(this.promptCaptor.capture())).willReturn(chatResponse);
+		willDoNothing().given(this.structuredOutputChatOptions).setOutputSchema(outputSchemaCaptor.capture());
+
+		var textCallAdvisor = new ContextCatcherCallAdvisor();
+
+		ResponseEntity<ChatResponse, JsonNode> responseEntity = ChatClient.builder(this.chatModel)
+			.build()
+			.prompt()
+			.options(this.structuredOutputChatOptions)
+			.advisors(AdvisorParams.ENABLE_NATIVE_STRUCTURED_OUTPUT)
+			.advisors(textCallAdvisor)
+			.user("Tell me about John")
+			.call()
+			.responseEntity(new CustomJsonSchemaOutputConverter(USER_JSON_SCHEMA));
+
+		var context = textCallAdvisor.getContext();
+
+		assertThat(context).containsKey(ChatClientAttributes.OUTPUT_FORMAT.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey());
+
+		assertThat(responseEntity.getResponse()).isEqualTo(chatResponse);
+		assertThat(responseEntity.getResponse().getMetadata().get("key1").toString()).isEqualTo("value1");
+
+		JsonAssertions.assertThatJson(responseEntity.getEntity()).isEqualTo("""
+				{
+					name: 'John',
+					age: 30
+				}
+				""");
+
+		Message userMessage = this.promptCaptor.getValue().getInstructions().get(0);
+		assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
+		assertThat(userMessage.getText()).isEqualTo("Tell me about John");
+
+		JsonAssertions.assertThatJson(outputSchemaCaptor.getValue())
+			.when(Option.IGNORING_ARRAY_ORDER)
+			.isEqualTo(USER_JSON_SCHEMA);
+	}
+
+	@Test
+	public void nativeWithCustomStructuredOutputConverterEntityTest(@Captor ArgumentCaptor<String> outputSchemaCaptor) {
+
+		ChatResponseMetadata metadata = ChatResponseMetadata.builder().keyValue("key1", "value1").build();
+
+		var chatResponse = new ChatResponse(List.of(new Generation(new AssistantMessage("""
+				{"name":"John", "age":30}
+				"""))), metadata);
+
+		given(this.chatModel.call(this.promptCaptor.capture())).willReturn(chatResponse);
+		willDoNothing().given(this.structuredOutputChatOptions).setOutputSchema(outputSchemaCaptor.capture());
+
+		var textCallAdvisor = new ContextCatcherCallAdvisor();
+
+		JsonNode entity = ChatClient.builder(this.chatModel)
+			.build()
+			.prompt()
+			.options(this.structuredOutputChatOptions)
+			.advisors(AdvisorParams.ENABLE_NATIVE_STRUCTURED_OUTPUT)
+			.advisors(textCallAdvisor)
+			.user("Tell me about John")
+			.call()
+			.entity(new CustomJsonSchemaOutputConverter(USER_JSON_SCHEMA));
+
+		var context = textCallAdvisor.getContext();
+
+		assertThat(context).containsKey(ChatClientAttributes.OUTPUT_FORMAT.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey());
+
+		JsonAssertions.assertThatJson(entity).isEqualTo("""
+				{
+					name: 'John',
+					age: 30
+				}
+				""");
+
+		Message userMessage = this.promptCaptor.getValue().getInstructions().get(0);
+		assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
+		assertThat(userMessage.getText()).isEqualTo("Tell me about John");
+
+		JsonAssertions.assertThatJson(outputSchemaCaptor.getValue())
+			.when(Option.IGNORING_ARRAY_ORDER)
+			.isEqualTo(USER_JSON_SCHEMA);
+	}
+
+	@Test
+	public void nativeWithCustomStructuredOutputConverterWithoutJsonSchemaResponseEntityTest() {
+
+		ChatResponseMetadata metadata = ChatResponseMetadata.builder().keyValue("key1", "value1").build();
+
+		var chatResponse = new ChatResponse(List.of(new Generation(new AssistantMessage("""
+				{"name":"John", "age":30}
+				"""))), metadata);
+
+		given(this.chatModel.call(this.promptCaptor.capture())).willReturn(chatResponse);
+		given(this.structuredOutputChatOptions.copy()).willReturn(this.structuredOutputChatOptions);
+
+		var textCallAdvisor = new ContextCatcherCallAdvisor();
+
+		ResponseEntity<ChatResponse, JsonNode> responseEntity = ChatClient.builder(this.chatModel)
+			.build()
+			.prompt()
+			.options(this.structuredOutputChatOptions)
+			.advisors(AdvisorParams.ENABLE_NATIVE_STRUCTURED_OUTPUT)
+			.advisors(textCallAdvisor)
+			.user("Tell me about John")
+			.call()
+			.responseEntity(new CustomJsonSchemaOutputConverter(null));
+
+		var context = textCallAdvisor.getContext();
+
+		assertThat(context).containsKey(ChatClientAttributes.OUTPUT_FORMAT.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey());
+
+		assertThat(responseEntity.getResponse()).isEqualTo(chatResponse);
+		assertThat(responseEntity.getResponse().getMetadata().get("key1").toString()).isEqualTo("value1");
+
+		JsonAssertions.assertThatJson(responseEntity.getEntity()).isEqualTo("""
+				{
+					name: 'John',
+					age: 30
+				}
+				""");
+
+		Message userMessage = this.promptCaptor.getValue().getInstructions().get(0);
+		assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
+		assertThat(userMessage.getText()).contains("Tell me about John", "Your response should be in JSON format");
+
+		verify(this.structuredOutputChatOptions, never()).setOutputSchema(anyString());
+	}
+
+	@Test
+	public void nativeWithCustomStructuredOutputConverterWithoutJsonSchemaEntityTest() {
+
+		ChatResponseMetadata metadata = ChatResponseMetadata.builder().keyValue("key1", "value1").build();
+
+		var chatResponse = new ChatResponse(List.of(new Generation(new AssistantMessage("""
+				{"name":"John", "age":30}
+				"""))), metadata);
+
+		given(this.chatModel.call(this.promptCaptor.capture())).willReturn(chatResponse);
+		given(this.structuredOutputChatOptions.copy()).willReturn(this.structuredOutputChatOptions);
+
+		var textCallAdvisor = new ContextCatcherCallAdvisor();
+
+		JsonNode entity = ChatClient.builder(this.chatModel)
+			.build()
+			.prompt()
+			.options(this.structuredOutputChatOptions)
+			.advisors(AdvisorParams.ENABLE_NATIVE_STRUCTURED_OUTPUT)
+			.advisors(textCallAdvisor)
+			.user("Tell me about John")
+			.call()
+			.entity(new CustomJsonSchemaOutputConverter(null));
+
+		var context = textCallAdvisor.getContext();
+
+		assertThat(context).containsKey(ChatClientAttributes.OUTPUT_FORMAT.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_SCHEMA.getKey());
+		assertThat(context).containsKey(ChatClientAttributes.STRUCTURED_OUTPUT_NATIVE.getKey());
+
+		JsonAssertions.assertThatJson(entity).isEqualTo("""
+				{
+					name: 'John',
+					age: 30
+				}
+				""");
+
+		Message userMessage = this.promptCaptor.getValue().getInstructions().get(0);
+		assertThat(userMessage.getMessageType()).isEqualTo(MessageType.USER);
+		assertThat(userMessage.getText()).contains("Tell me about John", "Your response should be in JSON format");
+
+		verify(this.structuredOutputChatOptions, never()).setOutputSchema(anyString());
+	}
+
 	record UserEntity(String name, int age) {
 	}
 
 	private static class ContextCatcherCallAdvisor implements CallAdvisor {
 
-		private Map<String, Object> context = new ConcurrentHashMap<>();
+		private Map<String, Object> context = new HashMap<>();
 
 		@Override
 		public String getName() {
@@ -355,4 +546,33 @@ public Map<String, Object> getContext() {
 
 	};
 
+	private static final class CustomJsonSchemaOutputConverter implements StructuredOutputConverter<JsonNode> {
+
+		private final String jsonSchema;
+
+		private CustomJsonSchemaOutputConverter(String jsonSchema) {
+			this.jsonSchema = jsonSchema;
+		}
+
+		@Override
+		public @NonNull String getFormat() {
+			return """
+					Your response should be in JSON format.
+					The JSON Schema is:
+					```%s```
+					""".formatted(this.jsonSchema);
+		}
+
+		@Override
+		public String getJsonSchema() {
+			return this.jsonSchema;
+		}
+
+		@Override
+		public JsonNode convert(String source) {
+			return JsonParser.getJsonMapper().readTree(source);
+		}
+
+	}
+
 }

spring-ai-model/src/main/java/org/springframework/ai/converter/BeanOutputConverter.java

@@ -279,6 +279,7 @@ public String getFormat() {
 	 * Provides the generated JSON schema for the target type.
 	 * @return The generated JSON schema.
 	 */
+	@Override
 	public String getJsonSchema() {
 		return this.jsonSchema;
 	}

spring-ai-model/src/main/java/org/springframework/ai/converter/StructuredOutputConverter.java

@@ -16,6 +16,8 @@
 
 package org.springframework.ai.converter;
 
+import org.jspecify.annotations.Nullable;
+
 import org.springframework.core.convert.converter.Converter;
 
 /**
@@ -26,7 +28,17 @@
  * @param <T> Specifies the desired response type.
  * @author Mark Pollack
  * @author Christian Tzolov
+ * @author Filip Hrisafov
  */
 public interface StructuredOutputConverter<T> extends Converter<String, T>, FormatProvider {
 
+	/**
+	 * Returns the JSON schema for the structured output of an LLM call.
+	 * @return the JSON schema or {@code null} if not available
+	 * @since 2.0.0
+	 */
+	@Nullable default String getJsonSchema() {
+		return null;
+	}
+
 }