Use WebClient.Builder and RestClient.Builder as ctor args in OpenAiAutoConfiguration

Fixes #609

Author: deepakn27

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

@@ -72,7 +72,7 @@ public OpenAiApi(String openAiToken) {
 	 * @param openAiToken OpenAI apiKey.
 	 */
 	public OpenAiApi(String baseUrl, String openAiToken) {
-		this(baseUrl, openAiToken, RestClient.builder());
+		this(baseUrl, openAiToken, RestClient.builder(), WebClient.builder());
 	}
 
 	/**
@@ -82,8 +82,8 @@ public OpenAiApi(String baseUrl, String openAiToken) {
 	 * @param openAiToken OpenAI apiKey.
 	 * @param restClientBuilder RestClient builder.
 	 */
-	public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder) {
-		this(baseUrl, openAiToken, restClientBuilder, RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
+	public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder) {
+		this(baseUrl, openAiToken, restClientBuilder, webClientBuilder, RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
 	}
 
 	/**
@@ -94,15 +94,15 @@ public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClie
 	 * @param restClientBuilder RestClient builder.
 	 * @param responseErrorHandler Response error handler.
 	 */
-	public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder, ResponseErrorHandler responseErrorHandler) {
+	public OpenAiApi(String baseUrl, String openAiToken, RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler) {
 
 		this.restClient = restClientBuilder
 				.baseUrl(baseUrl)
 				.defaultHeaders(ApiUtils.getJsonContentHeaders(openAiToken))
 				.defaultStatusHandler(responseErrorHandler)
 				.build();
 
-		this.webClient = WebClient.builder()
+		this.webClient = webClientBuilder
 				.baseUrl(baseUrl)
 				.defaultHeaders(ApiUtils.getJsonContentHeaders(openAiToken))
 				.build();

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatModelWithChatResponseMetadataTests.java

@@ -39,6 +39,7 @@
 import org.springframework.http.MediaType;
 import org.springframework.test.web.client.MockRestServiceServer;
 import org.springframework.web.client.RestClient;
+import org.springframework.web.reactive.function.client.WebClient;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.springframework.test.web.client.match.MockRestRequestMatchers.header;
@@ -166,8 +167,8 @@ private String getJson() {
 	static class Config {
 
 		@Bean
-		public OpenAiApi chatCompletionApi(RestClient.Builder builder) {
-			return new OpenAiApi("", TEST_API_KEY, builder);
+		public OpenAiApi chatCompletionApi(RestClient.Builder builder, WebClient.Builder webClientBuilder) {
+			return new OpenAiApi("", TEST_API_KEY, builder, webClientBuilder);
 		}
 
 		@Bean

spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/openai/OpenAiAutoConfiguration.java

@@ -39,6 +39,7 @@
 import org.springframework.util.StringUtils;
 import org.springframework.web.client.ResponseErrorHandler;
 import org.springframework.web.client.RestClient;
+import org.springframework.web.reactive.function.client.WebClient;
 
 /**
  * @author Christian Tzolov
@@ -56,11 +57,13 @@ public class OpenAiAutoConfiguration {
 			matchIfMissing = true)
 	public OpenAiChatModel openAiChatModel(OpenAiConnectionProperties commonProperties,
 			OpenAiChatProperties chatProperties, RestClient.Builder restClientBuilder,
-			List<FunctionCallback> toolFunctionCallbacks, FunctionCallbackContext functionCallbackContext,
-			RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) {
+			WebClient.Builder webClientBuilder, List<FunctionCallback> toolFunctionCallbacks,
+			FunctionCallbackContext functionCallbackContext, RetryTemplate retryTemplate,
+			ResponseErrorHandler responseErrorHandler) {
 
 		var openAiApi = openAiApi(chatProperties.getBaseUrl(), commonProperties.getBaseUrl(),
-				chatProperties.getApiKey(), commonProperties.getApiKey(), restClientBuilder, responseErrorHandler);
+				chatProperties.getApiKey(), commonProperties.getApiKey(), restClientBuilder, webClientBuilder,
+				responseErrorHandler);
 
 		if (!CollectionUtils.isEmpty(toolFunctionCallbacks)) {
 			chatProperties.getOptions().getFunctionCallbacks().addAll(toolFunctionCallbacks);
@@ -75,25 +78,29 @@ public OpenAiChatModel openAiChatModel(OpenAiConnectionProperties commonProperti
 			matchIfMissing = true)
 	public OpenAiEmbeddingModel openAiEmbeddingModel(OpenAiConnectionProperties commonProperties,
 			OpenAiEmbeddingProperties embeddingProperties, RestClient.Builder restClientBuilder,
-			RetryTemplate retryTemplate, ResponseErrorHandler responseErrorHandler) {
+			WebClient.Builder webClientBuilder, RetryTemplate retryTemplate,
+			ResponseErrorHandler responseErrorHandler) {
 
 		var openAiApi = openAiApi(embeddingProperties.getBaseUrl(), commonProperties.getBaseUrl(),
-				embeddingProperties.getApiKey(), commonProperties.getApiKey(), restClientBuilder, responseErrorHandler);
+				embeddingProperties.getApiKey(), commonProperties.getApiKey(), restClientBuilder, webClientBuilder,
+				responseErrorHandler);
 
 		return new OpenAiEmbeddingModel(openAiApi, embeddingProperties.getMetadataMode(),
 				embeddingProperties.getOptions(), retryTemplate);
 	}
 
 	private OpenAiApi openAiApi(String baseUrl, String commonBaseUrl, String apiKey, String commonApiKey,
-			RestClient.Builder restClientBuilder, ResponseErrorHandler responseErrorHandler) {
+			RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder,
+			ResponseErrorHandler responseErrorHandler) {
 
 		String resolvedBaseUrl = StringUtils.hasText(baseUrl) ? baseUrl : commonBaseUrl;
 		Assert.hasText(resolvedBaseUrl, "OpenAI base URL must be set");
 
 		String resolvedApiKey = StringUtils.hasText(apiKey) ? apiKey : commonApiKey;
 		Assert.hasText(resolvedApiKey, "OpenAI API key must be set");
 
-		return new OpenAiApi(resolvedBaseUrl, resolvedApiKey, restClientBuilder, responseErrorHandler);
+		return new OpenAiApi(resolvedBaseUrl, resolvedApiKey, restClientBuilder, webClientBuilder,
+				responseErrorHandler);
 	}
 
 	@Bean
@@ -122,6 +129,7 @@ public OpenAiImageModel openAiImageModel(OpenAiConnectionProperties commonProper
 	@ConditionalOnMissingBean
 	public OpenAiAudioTranscriptionModel openAiAudioTranscriptionModel(OpenAiConnectionProperties commonProperties,
 			OpenAiAudioTranscriptionProperties transcriptionProperties, RetryTemplate retryTemplate,
+			RestClient.Builder restClientBuilder, WebClient.Builder webClientBuilder,
 			ResponseErrorHandler responseErrorHandler) {
 
 		String apiKey = StringUtils.hasText(transcriptionProperties.getApiKey()) ? transcriptionProperties.getApiKey()
@@ -133,7 +141,8 @@ public OpenAiAudioTranscriptionModel openAiAudioTranscriptionModel(OpenAiConnect
 		Assert.hasText(apiKey, "OpenAI API key must be set");
 		Assert.hasText(baseUrl, "OpenAI base URL must be set");
 
-		var openAiAudioApi = new OpenAiAudioApi(baseUrl, apiKey, RestClient.builder(), responseErrorHandler);
+		var openAiAudioApi = new OpenAiAudioApi(baseUrl, apiKey, restClientBuilder, webClientBuilder,
+				responseErrorHandler);
 
 		OpenAiAudioTranscriptionModel openAiChatModel = new OpenAiAudioTranscriptionModel(openAiAudioApi,
 				transcriptionProperties.getOptions(), retryTemplate);
@@ -143,8 +152,9 @@ public OpenAiAudioTranscriptionModel openAiAudioTranscriptionModel(OpenAiConnect
 
 	@Bean
 	@ConditionalOnMissingBean
-	public OpenAiAudioSpeechModel openAiAudioSpeechModel(OpenAiConnectionProperties commonProperties,
-			OpenAiAudioSpeechProperties speechProperties, ResponseErrorHandler responseErrorHandler) {
+	public OpenAiAudioSpeechModel openAiAudioSpeechClient(OpenAiConnectionProperties commonProperties,
+			OpenAiAudioSpeechProperties speechProperties, RestClient.Builder restClientBuilder,
+			WebClient.Builder webClientBuilder, ResponseErrorHandler responseErrorHandler) {
 
 		String apiKey = StringUtils.hasText(speechProperties.getApiKey()) ? speechProperties.getApiKey()
 				: commonProperties.getApiKey();
@@ -155,7 +165,8 @@ public OpenAiAudioSpeechModel openAiAudioSpeechModel(OpenAiConnectionProperties
 		Assert.hasText(apiKey, "OpenAI API key must be set");
 		Assert.hasText(baseUrl, "OpenAI base URL must be set");
 
-		var openAiAudioApi = new OpenAiAudioApi(baseUrl, apiKey, RestClient.builder(), responseErrorHandler);
+		var openAiAudioApi = new OpenAiAudioApi(baseUrl, apiKey, restClientBuilder, webClientBuilder,
+				responseErrorHandler);
 
 		OpenAiAudioSpeechModel openAiSpeechModel = new OpenAiAudioSpeechModel(openAiAudioApi,
 				speechProperties.getOptions());