Skip to content

Commit fd727b4

Browse files
tzolovtzolov
committed
Add retry backoff config. for auto-configuraitons
1 parent 3d05830 commit fd727b4

File tree

5 files changed

+143
-34
lines changed

5 files changed

+143
-34
lines changed

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/azure/openai/AzureOpenAiAutoConfiguration.java

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,6 @@
1616

1717
package org.springframework.ai.autoconfigure.azure.openai;
1818

19-
import java.time.Duration;
20-
2119
import com.azure.ai.openai.OpenAIClient;
2220
import com.azure.ai.openai.OpenAIClientBuilder;
2321
import com.azure.core.credential.AzureKeyCredential;
@@ -68,7 +66,7 @@ public AiClient azureOpenAiClient(OpenAIClient msoftSdkOpenAiClient, AzureOpenAi
6866
azureOpenAiClient.setTemperature(this.azureOpenAiProperties.getTemperature());
6967
azureOpenAiClient.setModel(this.azureOpenAiProperties.getModel());
7068

71-
return (azureOpenAiProperties.isRetryEnabled()) ? new RetryAiClient(retryTemplate, azureOpenAiClient)
69+
return (azureOpenAiProperties.getRetry().isEnabled()) ? new RetryAiClient(retryTemplate, azureOpenAiClient)
7270
: azureOpenAiClient;
7371
}
7472

@@ -78,16 +76,20 @@ public EmbeddingClient azureOpenAiEmbeddingClient(OpenAIClient msoftSdkOpenAiCli
7876
AzureOpenAiProperties azureOpenAiProperties, RetryTemplate retryTemplate) {
7977
var embeddingClient = new AzureOpenAiEmbeddingClient(msoftSdkOpenAiClient,
8078
this.azureOpenAiProperties.getEmbeddingModel());
81-
return (azureOpenAiProperties.isRetryEnabled()) ? new RetryEmbeddingClient(retryTemplate, embeddingClient)
79+
return (azureOpenAiProperties.getRetry().isEnabled()) ? new RetryEmbeddingClient(retryTemplate, embeddingClient)
8280
: embeddingClient;
8381
}
8482

8583
@Bean
8684
@ConditionalOnMissingBean
87-
public RetryTemplate retryTemplate() {
85+
public RetryTemplate retryTemplate(AzureOpenAiProperties openAiProperties) {
86+
var retry = openAiProperties.getRetry();
87+
// currentInterval = Math.min(initialInterval * Math.pow(multiplier, retryNum),
88+
// maxInterval)}
8889
return RetryTemplate.builder()
89-
.maxAttempts(10)
90-
.exponentialBackoff(Duration.ofSeconds(2), 5, Duration.ofMinutes(2))
90+
.maxAttempts(retry.getMaxAttempts())
91+
.exponentialBackoff(retry.getInitialInterval(), retry.getBackoffIntervalMultiplier(),
92+
retry.getMaximumBackoffDuration())
9193
.build();
9294
}
9395

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/azure/openai/AzureOpenAiProperties.java

Lines changed: 61 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
package org.springframework.ai.autoconfigure.azure.openai;
1818

19+
import java.time.Duration;
20+
1921
import org.springframework.boot.context.properties.ConfigurationProperties;
2022

2123
@ConfigurationProperties(AzureOpenAiProperties.CONFIG_PREFIX)
@@ -35,7 +37,65 @@ public class AzureOpenAiProperties {
3537

3638
private String embeddingModel = "text-embedding-ada-002";
3739

38-
private boolean retryEnabled = false;
40+
private final Retry retry = new Retry();
41+
42+
public Retry getRetry() {
43+
return retry;
44+
}
45+
46+
public static class Retry {
47+
48+
private boolean enabled = false;
49+
50+
private int maxAttempts = 10;
51+
52+
private Duration initialInterval = Duration.ofSeconds(2);
53+
54+
private double backoffIntervalMultiplier = 5.0;
55+
56+
private Duration maximumBackoffDuration = Duration.ofMinutes(2);
57+
58+
public boolean isEnabled() {
59+
return enabled;
60+
}
61+
62+
public void setEnabled(boolean enabled) {
63+
this.enabled = enabled;
64+
}
65+
66+
public int getMaxAttempts() {
67+
return maxAttempts;
68+
}
69+
70+
public void setMaxAttempts(int maxAttempts) {
71+
this.maxAttempts = maxAttempts;
72+
}
73+
74+
public Duration getInitialInterval() {
75+
return initialInterval;
76+
}
77+
78+
public void setInitialInterval(Duration initialInterval) {
79+
this.initialInterval = initialInterval;
80+
}
81+
82+
public double getBackoffIntervalMultiplier() {
83+
return backoffIntervalMultiplier;
84+
}
85+
86+
public void setBackoffIntervalMultiplier(double backoffIntervalMultiplier) {
87+
this.backoffIntervalMultiplier = backoffIntervalMultiplier;
88+
}
89+
90+
public Duration getMaximumBackoffDuration() {
91+
return maximumBackoffDuration;
92+
}
93+
94+
public void setMaximumBackoffDuration(Duration maximumBackoffDuration) {
95+
this.maximumBackoffDuration = maximumBackoffDuration;
96+
}
97+
98+
}
3999

40100
public String getEndpoint() {
41101
return endpoint;
@@ -81,12 +141,4 @@ public void setEmbeddingModel(String embeddingModel) {
81141
this.embeddingModel = embeddingModel;
82142
}
83143

84-
public boolean isRetryEnabled() {
85-
return retryEnabled;
86-
}
87-
88-
public void setRetryEnabled(boolean retryEnabled) {
89-
this.retryEnabled = retryEnabled;
90-
}
91-
92144
}

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,8 @@ public AiClient openAiClient(OpenAiProperties openAiProperties, RetryTemplate re
6060
openAiClient.setTemperature(openAiProperties.getTemperature());
6161
openAiClient.setModel(openAiProperties.getModel());
6262

63-
return (openAiProperties.isRetryEnabled()) ? new RetryAiClient(retryTemplate, openAiClient) : openAiClient;
63+
return (openAiProperties.getRetry().isEnabled()) ? new RetryAiClient(retryTemplate, openAiClient)
64+
: openAiClient;
6465
}
6566

6667
@Bean
@@ -72,7 +73,7 @@ public EmbeddingClient openAiEmbeddingClient(OpenAiProperties openAiProperties,
7273

7374
var embeddingClient = new OpenAiEmbeddingClient(openAiService, openAiProperties.getEmbedding().getModel());
7475

75-
return (openAiProperties.isRetryEnabled()) ? new RetryEmbeddingClient(retryTemplate, embeddingClient)
76+
return (openAiProperties.getRetry().isEnabled()) ? new RetryEmbeddingClient(retryTemplate, embeddingClient)
7677
: embeddingClient;
7778
}
7879

@@ -109,10 +110,14 @@ private OpenAiService theoOpenAiService(OpenAiProperties properties, String base
109110

110111
@Bean
111112
@ConditionalOnMissingBean
112-
public RetryTemplate retryTemplate() {
113+
public RetryTemplate retryTemplate(OpenAiProperties openAiProperties) {
114+
var retry = openAiProperties.getRetry();
115+
// currentInterval = Math.min(initialInterval * Math.pow(multiplier, retryNum),
116+
// maxInterval)}
113117
return RetryTemplate.builder()
114-
.maxAttempts(10)
115-
.exponentialBackoff(Duration.ofSeconds(2), 5, Duration.ofMinutes(2))
118+
.maxAttempts(retry.getMaxAttempts())
119+
.exponentialBackoff(retry.getInitialInterval(), retry.getBackoffIntervalMultiplier(),
120+
retry.getMaximumBackoffDuration())
116121
.build();
117122
}
118123

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

Lines changed: 59 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,65 @@ public class OpenAiProperties {
4141

4242
private String baseUrl = "https://api.openai.com";
4343

44-
private boolean retryEnabled = false;
44+
private final Retry retry = new Retry();
45+
46+
public Retry getRetry() {
47+
return retry;
48+
}
49+
50+
public static class Retry {
51+
52+
private boolean enabled = false;
53+
54+
private int maxAttempts = 10;
55+
56+
private Duration initialInterval = Duration.ofSeconds(2);
57+
58+
private double backoffIntervalMultiplier = 5.0;
59+
60+
private Duration maximumBackoffDuration = Duration.ofMinutes(2);
61+
62+
public boolean isEnabled() {
63+
return enabled;
64+
}
65+
66+
public void setEnabled(boolean enabled) {
67+
this.enabled = enabled;
68+
}
69+
70+
public int getMaxAttempts() {
71+
return maxAttempts;
72+
}
73+
74+
public void setMaxAttempts(int maxAttempts) {
75+
this.maxAttempts = maxAttempts;
76+
}
77+
78+
public Duration getInitialInterval() {
79+
return initialInterval;
80+
}
81+
82+
public void setInitialInterval(Duration initialInterval) {
83+
this.initialInterval = initialInterval;
84+
}
85+
86+
public double getBackoffIntervalMultiplier() {
87+
return backoffIntervalMultiplier;
88+
}
89+
90+
public void setBackoffIntervalMultiplier(double backoffIntervalMultiplier) {
91+
this.backoffIntervalMultiplier = backoffIntervalMultiplier;
92+
}
93+
94+
public Duration getMaximumBackoffDuration() {
95+
return maximumBackoffDuration;
96+
}
97+
98+
public void setMaximumBackoffDuration(Duration maximumBackoffDuration) {
99+
this.maximumBackoffDuration = maximumBackoffDuration;
100+
}
101+
102+
}
45103

46104
public String getApiKey() {
47105
return this.apiKey;
@@ -91,14 +149,6 @@ public Metadata getMetadata() {
91149
return this.metadata;
92150
}
93151

94-
public void setRetryEnabled(boolean retry) {
95-
this.retryEnabled = retry;
96-
}
97-
98-
public boolean isRetryEnabled() {
99-
return retryEnabled;
100-
}
101-
102152
public static class Embedding {
103153

104154
private final OpenAiProperties openAiProperties;

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/OpenAiIT.java

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,9 @@ public class OpenAiIT {
5050
@ParameterizedTest(name = "{0} : {displayName} ")
5151
@ValueSource(booleans = { false, true })
5252
public void embeddingClient(boolean retryEnabled) {
53-
contextRunner.withPropertyValues("spring.ai.openai.retryEnabled=" + retryEnabled).run(context -> {
53+
contextRunner.withPropertyValues("spring.ai.openai.retry.enabled=" + retryEnabled).run(context -> {
5454
OpenAiProperties properties = context.getBean(OpenAiProperties.class);
55-
assertThat(properties.isRetryEnabled()).isEqualTo(retryEnabled);
55+
assertThat(properties.getRetry().isEnabled()).isEqualTo(retryEnabled);
5656

5757
EmbeddingClient embeddingClient = context.getBean(EmbeddingClient.class);
5858
if (retryEnabled) {
@@ -85,9 +85,9 @@ public void embeddingClient(boolean retryEnabled) {
8585
@ParameterizedTest(name = "{0} : {displayName} ")
8686
@ValueSource(booleans = { false, true })
8787
public void aiClient(boolean retryEnabled) {
88-
contextRunner.withPropertyValues("spring.ai.openai.retryEnabled=" + retryEnabled).run(context -> {
88+
contextRunner.withPropertyValues("spring.ai.openai.retry.enabled=" + retryEnabled).run(context -> {
8989
OpenAiProperties properties = context.getBean(OpenAiProperties.class);
90-
assertThat(properties.isRetryEnabled()).isEqualTo(retryEnabled);
90+
assertThat(properties.getRetry().isEnabled()).isEqualTo(retryEnabled);
9191

9292
AiClient aiClient = context.getBean(AiClient.class);
9393
if (retryEnabled) {

0 commit comments

Comments
 (0)