perf(azure-openai): support HTTP client timeout configuration

Signed-off-by: yinh <[email protected]>

Author: Fottas

auto-configurations/models/spring-ai-autoconfigure-model-azure-openai/src/main/java/org/springframework/ai/model/azure/openai/autoconfigure/AzureOpenAiClientBuilderConfiguration.java

@@ -23,8 +23,8 @@
 import com.azure.ai.openai.OpenAIClientBuilder;
 import com.azure.core.credential.AzureKeyCredential;
 import com.azure.core.credential.KeyCredential;
-import com.azure.core.util.ClientOptions;
 import com.azure.core.util.Header;
+import com.azure.core.util.HttpClientOptions;
 import com.azure.identity.DefaultAzureCredentialBuilder;
 
 import org.springframework.beans.factory.ObjectProvider;
@@ -56,23 +56,17 @@ public OpenAIClientBuilder openAIClientBuilder(AzureOpenAiConnectionProperties c
 
 		final OpenAIClientBuilder clientBuilder;
 
+		HttpClientOptions clientOptions = createHttpClientOptions(connectionProperties);
+
 		// Connect to OpenAI (e.g. not the Azure OpenAI). The deploymentName property is
 		// used as OpenAI model name.
 		if (StringUtils.hasText(connectionProperties.getOpenAiApiKey())) {
 			clientBuilder = new OpenAIClientBuilder().endpoint("https://api.openai.com/v1")
 				.credential(new KeyCredential(connectionProperties.getOpenAiApiKey()))
-				.clientOptions(new ClientOptions().setApplicationId(APPLICATION_ID));
+				.clientOptions(clientOptions);
 			applyOpenAIClientBuilderCustomizers(clientBuilder, customizers);
 			return clientBuilder;
 		}
-
-		Map<String, String> customHeaders = connectionProperties.getCustomHeaders();
-		List<Header> headers = customHeaders.entrySet()
-			.stream()
-			.map(entry -> new Header(entry.getKey(), entry.getValue()))
-			.collect(Collectors.toList());
-		ClientOptions clientOptions = new ClientOptions().setApplicationId(APPLICATION_ID).setHeaders(headers);
-
 		Assert.hasText(connectionProperties.getEndpoint(), "Endpoint must not be empty");
 
 		if (!StringUtils.hasText(connectionProperties.getApiKey())) {
@@ -96,4 +90,44 @@ private void applyOpenAIClientBuilderCustomizers(OpenAIClientBuilder clientBuild
 		customizers.orderedStream().forEach(customizer -> customizer.customize(clientBuilder));
 	}
 
+	/**
+	 * Create HttpClientOptions
+	 */
+	private HttpClientOptions createHttpClientOptions(AzureOpenAiConnectionProperties connectionProperties) {
+		// Create HttpClientOptions and apply the configuration
+		HttpClientOptions options = new HttpClientOptions();
+
+		options.setApplicationId(APPLICATION_ID);
+
+		Map<String, String> customHeaders = connectionProperties.getCustomHeaders();
+		List<Header> headers = customHeaders.entrySet()
+			.stream()
+			.map(entry -> new Header(entry.getKey(), entry.getValue()))
+			.collect(Collectors.toList());
+
+		options.setHeaders(headers);
+
+		if (connectionProperties.getConnectTimeout() != null) {
+			options.setConnectTimeout(connectionProperties.getConnectTimeout());
+		}
+
+		if (connectionProperties.getReadTimeout() != null) {
+			options.setReadTimeout(connectionProperties.getReadTimeout());
+		}
+
+		if (connectionProperties.getWriteTimeout() != null) {
+			options.setWriteTimeout(connectionProperties.getWriteTimeout());
+		}
+
+		if (connectionProperties.getResponseTimeout() != null) {
+			options.setResponseTimeout(connectionProperties.getResponseTimeout());
+		}
+
+		if (connectionProperties.getMaximumConnectionPoolSize() != null) {
+			options.setMaximumConnectionPoolSize(connectionProperties.getMaximumConnectionPoolSize());
+		}
+
+		return options;
+	}
+
 }

auto-configurations/models/spring-ai-autoconfigure-model-azure-openai/src/main/java/org/springframework/ai/model/azure/openai/autoconfigure/AzureOpenAiConnectionProperties.java

@@ -16,6 +16,7 @@
 
 package org.springframework.ai.model.azure.openai.autoconfigure;
 
+import java.time.Duration;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -46,6 +47,31 @@ public class AzureOpenAiConnectionProperties {
 
 	private Map<String, String> customHeaders = new HashMap<>();
 
+	/**
+	 * HTTP connection timeout
+	 */
+	private Duration connectTimeout;
+
+	/**
+	 * HTTP read timeout
+	 */
+	private Duration readTimeout;
+
+	/**
+	 * HTTP write timeout
+	 */
+	private Duration writeTimeout;
+
+	/**
+	 * HTTP response timeout
+	 */
+	private Duration responseTimeout;
+
+	/**
+	 * The maximum number of connections in the HTTP connection pool
+	 */
+	private Integer maximumConnectionPoolSize;
+
 	public String getEndpoint() {
 		return this.endpoint;
 	}
@@ -78,4 +104,44 @@ public void setCustomHeaders(Map<String, String> customHeaders) {
 		this.customHeaders = customHeaders;
 	}
 
+	public Duration getConnectTimeout() {
+		return this.connectTimeout;
+	}
+
+	public void setConnectTimeout(Duration connectTimeout) {
+		this.connectTimeout = connectTimeout;
+	}
+
+	public Duration getReadTimeout() {
+		return this.readTimeout;
+	}
+
+	public void setReadTimeout(Duration readTimeout) {
+		this.readTimeout = readTimeout;
+	}
+
+	public Duration getWriteTimeout() {
+		return this.writeTimeout;
+	}
+
+	public void setWriteTimeout(Duration writeTimeout) {
+		this.writeTimeout = writeTimeout;
+	}
+
+	public Duration getResponseTimeout() {
+		return this.responseTimeout;
+	}
+
+	public void setResponseTimeout(Duration responseTimeout) {
+		this.responseTimeout = responseTimeout;
+	}
+
+	public Integer getMaximumConnectionPoolSize() {
+		return this.maximumConnectionPoolSize;
+	}
+
+	public void setMaximumConnectionPoolSize(Integer maximumConnectionPoolSize) {
+		this.maximumConnectionPoolSize = maximumConnectionPoolSize;
+	}
+
 }

auto-configurations/models/spring-ai-autoconfigure-model-azure-openai/src/test/java/org/springframework/ai/model/azure/openai/autoconfigure/AzureOpenAiAutoConfigurationIT.java

@@ -54,6 +54,7 @@
 import org.springframework.util.ReflectionUtils;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
  * @author Christian Tzolov
@@ -288,4 +289,18 @@ void openAIClientBuilderCustomizer() {
 			});
 	}
 
+	@Test
+	void connectTimeoutShouldTakeEffect() {
+		new ApplicationContextRunner().withPropertyValues(
+						"spring.ai.azure.openai.connect-timeout=1ms"
+				)
+				.withConfiguration(SpringAiTestAutoConfigurations.of(AzureOpenAiChatAutoConfiguration.class))
+				.run(context -> {
+					AzureOpenAiChatModel chatModel = context.getBean(AzureOpenAiChatModel.class);
+
+					assertThatThrownBy(() -> chatModel.call(new Prompt("Hello")))
+							.isInstanceOf(Exception.class);
+				});
+	}
+
 }

auto-configurations/models/spring-ai-autoconfigure-model-azure-openai/src/test/java/org/springframework/ai/model/azure/openai/autoconfigure/AzureOpenAiAutoConfigurationPropertyTests.java

@@ -16,6 +16,8 @@
 
 package org.springframework.ai.model.azure.openai.autoconfigure;
 
+import java.time.Duration;
+
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.utils.SpringAiTestAutoConfigurations;
@@ -66,7 +68,13 @@ public void chatPropertiesTest() {
 				"spring.ai.azure.openai.chat.options.stop=boza,koza",
 				"spring.ai.azure.openai.chat.options.temperature=0.55",
 				"spring.ai.azure.openai.chat.options.topP=0.56",
-				"spring.ai.azure.openai.chat.options.user=userXYZ"
+				"spring.ai.azure.openai.chat.options.user=userXYZ",
+
+				"spring.ai.azure.openai.connect-timeout=10s",
+				"spring.ai.azure.openai.read-timeout=30s",
+				"spring.ai.azure.openai.write-timeout=30s",
+				"spring.ai.azure.openai.response-timeout=60s",
+				"spring.ai.azure.openai.maximum-connection-pool-size=50"
 				)
 			// @formatter:on
 			.withConfiguration(SpringAiTestAutoConfigurations.of(AzureOpenAiChatAutoConfiguration.class,
@@ -91,6 +99,12 @@ public void chatPropertiesTest() {
 				assertThat(chatProperties.getOptions().getTemperature()).isEqualTo(0.55);
 				assertThat(chatProperties.getOptions().getTopP()).isEqualTo(0.56);
 
+				assertThat(connectionProperties.getConnectTimeout()).isEqualTo(Duration.ofSeconds(10));
+				assertThat(connectionProperties.getReadTimeout()).isEqualTo(Duration.ofSeconds(30));
+				assertThat(connectionProperties.getWriteTimeout()).isEqualTo(Duration.ofSeconds(30));
+				assertThat(connectionProperties.getResponseTimeout()).isEqualTo(Duration.ofSeconds(60));
+				assertThat(connectionProperties.getMaximumConnectionPoolSize()).isEqualTo(50);
+
 				assertThat(chatProperties.getOptions().getUser()).isEqualTo("userXYZ");
 			});
 	}