Add more timeout configuration options for AWS Bedrock

When submitting large requests or receiving large response, AWS Bedrock request
may experience latency that exceeds our default settings. Here we add additional
configuration options to allow increasing these timeout settings.

Signed-off-by: Christopher Miles <[email protected]>

Author: cmiles74

auto-configurations/models/spring-ai-autoconfigure-model-bedrock-ai/src/main/java/org/springframework/ai/model/bedrock/autoconfigure/BedrockAwsConnectionProperties.java

@@ -53,10 +53,30 @@ public class BedrockAwsConnectionProperties {
 	private String sessionToken;
 
 	/**
-	 * Set model timeout, Defaults 5 min.
+	 * Maximum duration of the entire API call operation.
 	 */
 	private Duration timeout = Duration.ofMinutes(5L);
 
+	/**
+	 * Maximum time to wait while establishing connection with AWS service.
+	 */
+	private Duration connectionTimeout = Duration.ofSeconds(5L);
+
+	/**
+	 * Maximum duration spent reading response data.
+	 */
+	private Duration asyncReadTimeout = Duration.ofSeconds(30L);
+
+	/**
+	 * Maximum time to wait for a new connection from the pool.
+	 */
+	private Duration connectionAcquisitionTimeout = Duration.ofSeconds(30L);
+
+	/**
+	 * Maximum time to wait for response data.
+	 */
+	private Duration socketTimeout = Duration.ofSeconds(480L);
+
 	public String getRegion() {
 		return this.region;
 	}
@@ -89,6 +109,38 @@ public void setTimeout(Duration timeout) {
 		this.timeout = timeout;
 	}
 
+	public Duration getConnectionTimeout() {
+		return connectionTimeout;
+	}
+
+	public void setConnectionTimeout(Duration connectionTimeout) {
+		this.connectionTimeout = connectionTimeout;
+	}
+
+	public Duration getAsyncReadTimeout() {
+		return asyncReadTimeout;
+	}
+
+	public void setAsyncReadTimeout(Duration asyncReadTimeout) {
+		this.asyncReadTimeout = asyncReadTimeout;
+	}
+
+	public Duration getConnectionAcquisitionTimeout() {
+		return connectionAcquisitionTimeout;
+	}
+
+	public void setConnectionAcquisitionTimeout(Duration connectionAcquisitionTimeout) {
+		this.connectionAcquisitionTimeout = connectionAcquisitionTimeout;
+	}
+
+	public Duration getSocketTimeout() {
+		return socketTimeout;
+	}
+
+	public void setSocketTimeout(Duration socketTimeout) {
+		this.socketTimeout = socketTimeout;
+	}
+
 	public String getSessionToken() {
 		return this.sessionToken;
 	}

auto-configurations/models/spring-ai-autoconfigure-model-bedrock-ai/src/main/java/org/springframework/ai/model/bedrock/converse/autoconfigure/BedrockConverseProxyChatAutoConfiguration.java

@@ -76,6 +76,10 @@ public BedrockProxyChatModel bedrockProxyChatModel(AwsCredentialsProvider creden
 			.credentialsProvider(credentialsProvider)
 			.region(regionProvider.getRegion())
 			.timeout(connectionProperties.getTimeout())
+			.connectionTimeout(connectionProperties.getConnectionTimeout())
+			.asyncReadTimeout(connectionProperties.getAsyncReadTimeout())
+			.connectionAcquisitionTimeout(connectionProperties.getConnectionAcquisitionTimeout())
+			.socketTimeout(connectionProperties.getSocketTimeout())
 			.defaultOptions(chatProperties.getOptions())
 			.observationRegistry(observationRegistry.getIfUnique(() -> ObservationRegistry.NOOP))
 			.toolCallingManager(toolCallingManager)

models/spring-ai-bedrock-converse/pom.xml

@@ -76,6 +76,12 @@
 			<version>${bedrockruntime.version}</version>
 		</dependency>
 
+		<dependency>
+			<groupId>software.amazon.awssdk</groupId>
+			<artifactId>apache-client</artifactId>
+			<version>${bedrockruntime.version}</version>
+		</dependency>
+
 		<!-- test dependencies -->
 		<dependency>
 			<groupId>org.springframework.ai</groupId>

models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java

@@ -41,6 +41,7 @@
 import software.amazon.awssdk.core.SdkBytes;
 import software.amazon.awssdk.core.document.Document;
 import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.http.apache.ApacheHttpClient;
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.regions.providers.DefaultAwsRegionProviderChain;
@@ -781,7 +782,15 @@ public static final class Builder {
 
 		private Region region = Region.US_EAST_1;
 
-		private Duration timeout = Duration.ofMinutes(10);
+		private Duration timeout = Duration.ofMinutes(5L);
+
+		private Duration connectionTimeout = Duration.ofSeconds(5L);
+
+		private Duration asyncReadTimeout = Duration.ofSeconds(30L);
+
+		private Duration connectionAcquisitionTimeout = Duration.ofSeconds(30L);
+
+		private Duration socketTimeout = Duration.ofSeconds(30L);
 
 		private ToolCallingManager toolCallingManager;
 
@@ -835,6 +844,30 @@ public Builder timeout(Duration timeout) {
 			return this;
 		}
 
+		public Builder connectionTimeout(Duration connectionTimeout) {
+			Assert.notNull(connectionTimeout, "'connectionTimeout' must not be null.");
+			this.connectionTimeout = connectionTimeout;
+			return this;
+		}
+
+		public Builder asyncReadTimeout(Duration asyncReadTimeout) {
+			Assert.notNull(asyncReadTimeout, "'asyncReadTimeout' must not be null.");
+			this.asyncReadTimeout = asyncReadTimeout;
+			return this;
+		}
+
+		public Builder connectionAcquisitionTimeout(Duration connectionAcquisitionTimeout) {
+			Assert.notNull(connectionAcquisitionTimeout, "'connectionAcquisitionTimeout' must not be null.");
+			this.connectionAcquisitionTimeout = connectionAcquisitionTimeout;
+			return this;
+		}
+
+		public Builder socketTimeout(Duration socketTimeout) {
+			Assert.notNull(socketTimeout, "'socketTimeout' must not be null.");
+			this.socketTimeout = socketTimeout;
+			return this;
+		}
+
 		public Builder defaultOptions(ToolCallingChatOptions defaultOptions) {
 			Assert.notNull(defaultOptions, "'defaultOptions' must not be null.");
 			this.defaultOptions = defaultOptions;
@@ -864,22 +897,28 @@ public Builder bedrockRuntimeAsyncClient(BedrockRuntimeAsyncClient bedrockRuntim
 		}
 
 		public BedrockProxyChatModel build() {
-
 			if (this.bedrockRuntimeClient == null) {
+
+				var httpClientBuilder = ApacheHttpClient.builder()
+					.connectionAcquisitionTimeout(connectionAcquisitionTimeout)
+					.connectionTimeout(this.connectionTimeout)
+					.socketTimeout(this.socketTimeout);
+
 				this.bedrockRuntimeClient = BedrockRuntimeClient.builder()
 					.region(this.region)
-					.httpClientBuilder(null)
+					.httpClientBuilder(httpClientBuilder)
 					.credentialsProvider(this.credentialsProvider)
 					.overrideConfiguration(c -> c.apiCallTimeout(this.timeout))
 					.build();
 			}
 
 			if (this.bedrockRuntimeAsyncClient == null) {
 
-				// TODO: Is it ok to configure the NettyNioAsyncHttpClient explicitly???
 				var httpClientBuilder = NettyNioAsyncHttpClient.builder()
 					.tcpKeepAlive(true)
-					.connectionAcquisitionTimeout(Duration.ofSeconds(30))
+					.readTimeout(this.asyncReadTimeout)
+					.connectionTimeout(this.connectionTimeout)
+					.connectionAcquisitionTimeout(this.connectionAcquisitionTimeout)
 					.maxConcurrency(200);
 
 				var builder = BedrockRuntimeAsyncClient.builder()

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chat/bedrock-converse.adoc

@@ -73,11 +73,14 @@ The prefix `spring.ai.bedrock.aws` is the property prefix to configure the conne
 |====
 | Property | Description | Default
 
-| spring.ai.bedrock.aws.region     | AWS region to use.  | us-east-1
-| spring.ai.bedrock.aws.timeout    | AWS timeout to use. | 5m
-| spring.ai.bedrock.aws.access-key | AWS access key.  | -
-| spring.ai.bedrock.aws.secret-key | AWS secret key.  | -
-| spring.ai.bedrock.aws.session-token | AWS session token for temporary credentials. | -
+| spring.ai.bedrock.aws.region     | AWS region to use  | us-east-1
+| spring.ai.bedrock.aws.timeout    | AWS max duration for entire API call | 5m
+| spring.ai.bedrock.aws.connectionTimeout | Max duration to wait while establishing connection | 5s
+| spring.ai.bedrock.aws.connectionAcquisitionTimeout | Max duration to wait for new connection from the pool | 30s
+| spring.ai.bedrock.aws.asyncReadTimeout | Max duration spent reading asynchronous responses | 30s
+| spring.ai.bedrock.aws.access-key | AWS access key  | -
+| spring.ai.bedrock.aws.secret-key | AWS secret key  | -
+| spring.ai.bedrock.aws.session-token | AWS session token for temporary credentials | -
 |====
 
 [NOTE]