perf(stability-ai): Support HTTP client timeout configuration

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

Author: Fottas

auto-configurations/models/spring-ai-autoconfigure-model-stability-ai/src/main/java/org/springframework/ai/model/stabilityai/autoconfigure/StabilityAiConnectionProperties.java

@@ -16,8 +16,15 @@
 
 package org.springframework.ai.model.stabilityai.autoconfigure;
 
+import java.time.Duration;
+
+import jakarta.annotation.Nullable;
+
 import org.springframework.ai.stabilityai.api.StabilityAiApi;
 import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.context.properties.NestedConfigurationProperty;
+import org.springframework.boot.http.client.HttpRedirects;
+import org.springframework.boot.http.client.autoconfigure.HttpClientSettingsProperties;
 
 @ConfigurationProperties(StabilityAiConnectionProperties.CONFIG_PREFIX)
 public class StabilityAiConnectionProperties extends StabilityAiParentProperties {
@@ -26,8 +33,47 @@ public class StabilityAiConnectionProperties extends StabilityAiParentProperties
 
 	public static final String DEFAULT_BASE_URL = StabilityAiApi.DEFAULT_BASE_URL;
 
+	@NestedConfigurationProperty
+	private final HttpClientSettingsProperties http = new HttpClientSettingsProperties() {
+	};
+
 	public StabilityAiConnectionProperties() {
 		super.setBaseUrl(DEFAULT_BASE_URL);
 	}
 
+	@Nullable
+	public HttpRedirects getRedirects() {
+		return this.http.getRedirects();
+	}
+
+	public void setRedirects(HttpRedirects redirects) {
+		this.http.setRedirects(redirects);
+	}
+
+	@Nullable
+	public Duration getConnectTimeout() {
+		return this.http.getConnectTimeout();
+	}
+
+	public void setConnectTimeout(Duration connectTimeout) {
+		this.http.setConnectTimeout(connectTimeout);
+	}
+
+	@Nullable
+	public Duration getReadTimeout() {
+		return this.http.getReadTimeout();
+	}
+
+	public void setReadTimeout(Duration readTimeout) {
+		this.http.setReadTimeout(readTimeout);
+	}
+
+	public HttpClientSettingsProperties.Ssl getSsl() {
+		return this.http.getSsl();
+	}
+
+	public HttpClientSettingsProperties getHttp() {
+		return this.http;
+	}
+
 }

auto-configurations/models/spring-ai-autoconfigure-model-stability-ai/src/main/java/org/springframework/ai/model/stabilityai/autoconfigure/StabilityAiImageAutoConfiguration.java

@@ -26,8 +26,13 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
+import org.springframework.boot.http.client.ClientHttpRequestFactoryBuilder;
+import org.springframework.boot.http.client.HttpClientSettings;
+import org.springframework.boot.http.client.autoconfigure.HttpClientSettingsPropertyMapper;
 import org.springframework.boot.restclient.autoconfigure.RestClientAutoConfiguration;
+import org.springframework.boot.ssl.SslBundles;
 import org.springframework.context.annotation.Bean;
+import org.springframework.http.client.ClientHttpRequestFactory;
 import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 import org.springframework.web.client.RestClient;
@@ -50,7 +55,9 @@ public class StabilityAiImageAutoConfiguration {
 	@Bean
 	@ConditionalOnMissingBean
 	public StabilityAiApi stabilityAiApi(StabilityAiConnectionProperties commonProperties,
-			StabilityAiImageProperties imageProperties, ObjectProvider<RestClient.Builder> restClientBuilderProvider) {
+			StabilityAiImageProperties imageProperties, ObjectProvider<RestClient.Builder> restClientBuilderProvider,
+			ObjectProvider<SslBundles> sslBundles, ObjectProvider<HttpClientSettings> globalHttpClientSettings,
+			ObjectProvider<ClientHttpRequestFactoryBuilder<?>> factoryBuilder) {
 
 		String apiKey = StringUtils.hasText(imageProperties.getApiKey()) ? imageProperties.getApiKey()
 				: commonProperties.getApiKey();
@@ -61,8 +68,15 @@ public StabilityAiApi stabilityAiApi(StabilityAiConnectionProperties commonPrope
 		Assert.hasText(apiKey, "StabilityAI API key must be set");
 		Assert.hasText(baseUrl, "StabilityAI base URL must be set");
 
-		return new StabilityAiApi(apiKey, imageProperties.getOptions().getModel(), baseUrl,
-				restClientBuilderProvider.getIfAvailable(RestClient::builder));
+		HttpClientSettingsPropertyMapper mapper = new HttpClientSettingsPropertyMapper(sslBundles.getIfAvailable(),
+				globalHttpClientSettings.getIfAvailable());
+		HttpClientSettings httpClientSettings = mapper.map(commonProperties.getHttp());
+
+		RestClient.Builder restClientBuilder = restClientBuilderProvider.getIfAvailable(RestClient::builder);
+		applyRestClientSettings(restClientBuilder, httpClientSettings,
+				factoryBuilder.getIfAvailable(ClientHttpRequestFactoryBuilder::detect));
+
+		return new StabilityAiApi(apiKey, imageProperties.getOptions().getModel(), baseUrl, restClientBuilder);
 	}
 
 	@Bean
@@ -72,4 +86,10 @@ public StabilityAiImageModel stabilityAiImageModel(StabilityAiApi stabilityAiApi
 		return new StabilityAiImageModel(stabilityAiApi, stabilityAiImageProperties.getOptions());
 	}
 
+	private void applyRestClientSettings(RestClient.Builder builder, HttpClientSettings httpClientSettings,
+			ClientHttpRequestFactoryBuilder<?> factoryBuilder) {
+		ClientHttpRequestFactory requestFactory = factoryBuilder.build(httpClientSettings);
+		builder.requestFactory(requestFactory);
+	}
+
 }

auto-configurations/models/spring-ai-autoconfigure-model-stability-ai/src/test/java/org/springframework/ai/model/stabilityai/autoconfigure/StabilityAiAutoConfigurationIT.java

@@ -60,4 +60,29 @@ void generate() {
 		});
 	}
 
+	@Test
+	void generateWithCustomTimeout() {
+		this.contextRunner
+			.withPropertyValues("spring.ai.stabilityai.connect-timeout=1s", "spring.ai.stabilityai.read-timeout=1s")
+			.withConfiguration(AutoConfigurations.of(StabilityAiImageAutoConfiguration.class))
+			.run(context -> {
+				ImageModel imageModel = context.getBean(ImageModel.class);
+				StabilityAiImageOptions imageOptions = StabilityAiImageOptions.builder()
+					.stylePreset(StyleEnum.PHOTOGRAPHIC)
+					.build();
+
+				var instructions = """
+						A light cream colored mini golden doodle.
+						""";
+
+				ImagePrompt imagePrompt = new ImagePrompt(instructions, imageOptions);
+				ImageResponse imageResponse = imageModel.call(imagePrompt);
+
+				ImageGeneration imageGeneration = imageResponse.getResult();
+				Image image = imageGeneration.getOutput();
+
+				assertThat(image.getB64Json()).isNotEmpty();
+			});
+	}
+
 }

auto-configurations/models/spring-ai-autoconfigure-model-stability-ai/src/test/java/org/springframework/ai/model/stabilityai/autoconfigure/StabilityAiImagePropertiesTests.java

@@ -16,6 +16,8 @@
 
 package org.springframework.ai.model.stabilityai.autoconfigure;
 
+import java.time.Duration;
+
 import org.junit.jupiter.api.Test;
 
 import org.springframework.ai.stabilityai.StabilityAiImageModel;
@@ -48,12 +50,16 @@ public void chatPropertiesTest() {
 				"spring.ai.stabilityai.image.options.sampler=K_EULER",
 				"spring.ai.stabilityai.image.options.seed=0",
 				"spring.ai.stabilityai.image.options.steps=30",
-				"spring.ai.stabilityai.image.options.style-preset=neon-punk"
+				"spring.ai.stabilityai.image.options.style-preset=neon-punk",
+
+				"spring.ai.stabilityai.connect-timeout=5s",
+				"spring.ai.stabilityai.read-timeout=30s"
 				)
 			// @formatter:on
 			.withConfiguration(AutoConfigurations.of(StabilityAiImageAutoConfiguration.class))
 			.run(context -> {
 				var chatProperties = context.getBean(StabilityAiImageProperties.class);
+				var connectionProperties = context.getBean(StabilityAiConnectionProperties.class);
 
 				assertThat(chatProperties.getBaseUrl()).isEqualTo("ENDPOINT");
 				assertThat(chatProperties.getApiKey()).isEqualTo("API_KEY");
@@ -69,6 +75,9 @@ public void chatPropertiesTest() {
 				assertThat(chatProperties.getOptions().getSeed()).isEqualTo(0);
 				assertThat(chatProperties.getOptions().getSteps()).isEqualTo(30);
 				assertThat(chatProperties.getOptions().getStylePreset()).isEqualTo("neon-punk");
+
+				assertThat(connectionProperties.getConnectTimeout()).isEqualTo(Duration.ofSeconds(5));
+				assertThat(connectionProperties.getReadTimeout()).isEqualTo(Duration.ofSeconds(30));
 			});
 	}