Skip to content

Commit 30f3973

Browse files
tzolovtzolov
committed
Add ops retry support open-ai and azure-openai auto-config.
 Add retryEnabled (default to false) propety to the auto-config properties. When the retryEnabled=true the auto-config wrapps the response into corresponding RetryXXXX decorator uisng a default RestTemplate instance.
1 parent 404eaf2 commit 30f3973

File tree

6 files changed

+79
-30
lines changed

6 files changed

+79
-30
lines changed

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

Lines changed: 27 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,16 +16,23 @@
1616

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

19+
import java.time.Duration;
20+
1921
import com.azure.ai.openai.OpenAIClient;
2022
import com.azure.ai.openai.OpenAIClientBuilder;
2123
import com.azure.core.credential.AzureKeyCredential;
2224
import org.springframework.ai.azure.openai.client.AzureOpenAiClient;
2325
import org.springframework.ai.azure.openai.embedding.AzureOpenAiEmbeddingClient;
26+
import org.springframework.ai.client.AiClient;
27+
import org.springframework.ai.client.RetryAiClient;
28+
import org.springframework.ai.embedding.EmbeddingClient;
29+
import org.springframework.ai.embedding.RetryEmbeddingClient;
2430
import org.springframework.boot.autoconfigure.AutoConfiguration;
2531
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2632
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
2733
import org.springframework.boot.context.properties.EnableConfigurationProperties;
2834
import org.springframework.context.annotation.Bean;
35+
import org.springframework.retry.support.RetryTemplate;
2936
import org.springframework.util.StringUtils;
3037

3138
@AutoConfiguration
@@ -53,16 +60,32 @@ public OpenAIClient msoftSdkOpenAiClient(AzureOpenAiProperties azureOpenAiProper
5360
}
5461

5562
@Bean
56-
public AzureOpenAiClient azureOpenAiClient(OpenAIClient msoftSdkOpenAiClient) {
63+
public AiClient azureOpenAiClient(OpenAIClient msoftSdkOpenAiClient, AzureOpenAiProperties azureOpenAiProperties,
64+
RetryTemplate retryTemplate) {
5765
AzureOpenAiClient azureOpenAiClient = new AzureOpenAiClient(msoftSdkOpenAiClient);
5866
azureOpenAiClient.setTemperature(this.azureOpenAiProperties.getTemperature());
5967
azureOpenAiClient.setModel(this.azureOpenAiProperties.getModel());
60-
return azureOpenAiClient;
68+
69+
return (azureOpenAiProperties.isRetryEnabled()) ? new RetryAiClient(retryTemplate, azureOpenAiClient)
70+
: azureOpenAiClient;
71+
}
72+
73+
@Bean
74+
public EmbeddingClient azureOpenAiEmbeddingClient(OpenAIClient msoftSdkOpenAiClient,
75+
AzureOpenAiProperties azureOpenAiProperties, RetryTemplate retryTemplate) {
76+
var embeddingClient = new AzureOpenAiEmbeddingClient(msoftSdkOpenAiClient,
77+
this.azureOpenAiProperties.getEmbeddingModel());
78+
return (azureOpenAiProperties.isRetryEnabled()) ? new RetryEmbeddingClient(retryTemplate, embeddingClient)
79+
: embeddingClient;
6180
}
6281

6382
@Bean
64-
public AzureOpenAiEmbeddingClient azureOpenAiEmbeddingClient(OpenAIClient msoftSdkOpenAiClient) {
65-
return new AzureOpenAiEmbeddingClient(msoftSdkOpenAiClient, this.azureOpenAiProperties.getEmbeddingModel());
83+
@ConditionalOnMissingBean
84+
public RetryTemplate retryTemplate() {
85+
return RetryTemplate.builder()
86+
.maxAttempts(10)
87+
.exponentialBackoff(Duration.ofSeconds(2), 5, Duration.ofMinutes(2))
88+
.build();
6689
}
6790

6891
}

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

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ public class AzureOpenAiProperties {
3535

3636
private String embeddingModel = "text-embedding-ada-002";
3737

38+
private boolean retryEnabled = false;
39+
3840
public String getEndpoint() {
3941
return endpoint;
4042
}
@@ -79,4 +81,12 @@ public void setEmbeddingModel(String embeddingModel) {
7981
this.embeddingModel = embeddingModel;
8082
}
8183

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

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/huggingface/HuggingfaceAutoConfiguration.java

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

1717
package org.springframework.ai.autoconfigure.huggingface;
1818

19+
import org.springframework.ai.client.AiClient;
1920
import org.springframework.ai.huggingface.client.HuggingfaceAiClient;
2021
import org.springframework.boot.autoconfigure.AutoConfiguration;
2122
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -27,14 +28,8 @@
2728
@EnableConfigurationProperties(HuggingfaceProperties.class)
2829
public class HuggingfaceAutoConfiguration {
2930

30-
private final HuggingfaceProperties huggingfaceProperties;
31-
32-
public HuggingfaceAutoConfiguration(HuggingfaceProperties huggingfaceProperties) {
33-
this.huggingfaceProperties = huggingfaceProperties;
34-
}
35-
3631
@Bean
37-
public HuggingfaceAiClient huggingfaceAiClient(HuggingfaceProperties huggingfaceProperties) {
32+
public AiClient huggingfaceAiClient(HuggingfaceProperties huggingfaceProperties) {
3833
return new HuggingfaceAiClient(huggingfaceProperties.getApiKey(), huggingfaceProperties.getUrl());
3934
}
4035

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/ollama/OllamaAutoConfiguration.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.ai.autoconfigure.ollama;
1818

19+
import org.springframework.ai.client.AiClient;
1920
import org.springframework.ai.ollama.client.OllamaClient;
2021
import org.springframework.boot.autoconfigure.AutoConfiguration;
2122
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -30,7 +31,7 @@ public class OllamaAutoConfiguration {
3031

3132
@Bean
3233
@ConditionalOnMissingBean
33-
public OllamaClient ollamaClient(OllamaProperties properties) {
34+
public AiClient ollamaClient(OllamaProperties properties) {
3435
return new OllamaClient(properties.getBaseUrl(), properties.getModel());
3536
}
3637

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

Lines changed: 27 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -16,33 +16,33 @@
1616

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

19-
import static org.springframework.ai.autoconfigure.openai.OpenAiProperties.CONFIG_PREFIX;
20-
2119
import java.time.Duration;
2220

2321
import com.fasterxml.jackson.databind.ObjectMapper;
2422
import com.theokanning.openai.client.OpenAiApi;
2523
import com.theokanning.openai.service.OpenAiService;
24+
import okhttp3.OkHttpClient;
25+
import retrofit2.Retrofit;
26+
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
27+
import retrofit2.converter.jackson.JacksonConverterFactory;
2628

2729
import org.springframework.ai.autoconfigure.NativeHints;
30+
import org.springframework.ai.client.AiClient;
31+
import org.springframework.ai.client.RetryAiClient;
2832
import org.springframework.ai.embedding.EmbeddingClient;
33+
import org.springframework.ai.embedding.RetryEmbeddingClient;
2934
import org.springframework.ai.openai.client.OpenAiClient;
30-
import org.springframework.ai.openai.metadata.support.OpenAiHttpResponseHeadersInterceptor;
3135
import org.springframework.ai.openai.embedding.OpenAiEmbeddingClient;
36+
import org.springframework.ai.openai.metadata.support.OpenAiHttpResponseHeadersInterceptor;
3237
import org.springframework.boot.autoconfigure.AutoConfiguration;
3338
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
3439
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
3540
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3641
import org.springframework.context.annotation.Bean;
3742
import org.springframework.context.annotation.ImportRuntimeHints;
43+
import org.springframework.retry.support.RetryTemplate;
3844
import org.springframework.util.StringUtils;
3945

40-
import okhttp3.OkHttpClient;
41-
import retrofit2.Retrofit;
42-
import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory;
43-
import retrofit2.converter.jackson.JacksonConverterFactory;
44-
import retrofit2.http.HEAD;
45-
4646
@AutoConfiguration
4747
@ConditionalOnClass(OpenAiService.class)
4848
@EnableConfigurationProperties(OpenAiProperties.class)
@@ -51,7 +51,7 @@ public class OpenAiAutoConfiguration {
5151

5252
@Bean
5353
@ConditionalOnMissingBean
54-
public OpenAiClient openAiClient(OpenAiProperties openAiProperties) {
54+
public AiClient openAiClient(OpenAiProperties openAiProperties, RetryTemplate retryTemplate) {
5555

5656
OpenAiService openAiService = theoOpenAiService(openAiProperties, openAiProperties.getBaseUrl(),
5757
openAiProperties.getApiKey(), openAiProperties.getDuration());
@@ -60,25 +60,28 @@ public OpenAiClient openAiClient(OpenAiProperties openAiProperties) {
6060
openAiClient.setTemperature(openAiProperties.getTemperature());
6161
openAiClient.setModel(openAiProperties.getModel());
6262

63-
return openAiClient;
63+
return (openAiProperties.isRetryEnabled()) ? new RetryAiClient(retryTemplate, openAiClient) : openAiClient;
6464
}
6565

6666
@Bean
6767
@ConditionalOnMissingBean
68-
public EmbeddingClient openAiEmbeddingClient(OpenAiProperties openAiProperties) {
68+
public EmbeddingClient openAiEmbeddingClient(OpenAiProperties openAiProperties, RetryTemplate retryTemplate) {
6969

7070
OpenAiService openAiService = theoOpenAiService(openAiProperties, openAiProperties.getEmbedding().getBaseUrl(),
7171
openAiProperties.getEmbedding().getApiKey(), openAiProperties.getDuration());
7272

73-
return new OpenAiEmbeddingClient(openAiService, openAiProperties.getEmbedding().getModel());
73+
var embeddingClient = new OpenAiEmbeddingClient(openAiService, openAiProperties.getEmbedding().getModel());
74+
75+
return (openAiProperties.isRetryEnabled()) ? new RetryEmbeddingClient(retryTemplate, embeddingClient)
76+
: embeddingClient;
7477
}
7578

7679
private OpenAiService theoOpenAiService(OpenAiProperties properties, String baseUrl, String apiKey,
7780
Duration duration) {
7881

7982
if ("https://api.openai.com".equals(baseUrl) && !StringUtils.hasText(apiKey)) {
80-
throw new IllegalArgumentException(
81-
"You must provide an API key with the property name " + CONFIG_PREFIX + ".api-key");
83+
throw new IllegalArgumentException("You must provide an API key with the property name "
84+
+ OpenAiProperties.CONFIG_PREFIX + ".api-key");
8285
}
8386

8487
ObjectMapper mapper = OpenAiService.defaultObjectMapper();
@@ -104,4 +107,13 @@ private OpenAiService theoOpenAiService(OpenAiProperties properties, String base
104107
return new OpenAiService(api);
105108
}
106109

110+
@Bean
111+
@ConditionalOnMissingBean
112+
public RetryTemplate retryTemplate() {
113+
return RetryTemplate.builder()
114+
.maxAttempts(10)
115+
.exponentialBackoff(Duration.ofSeconds(2), 5, Duration.ofMinutes(2))
116+
.build();
117+
}
118+
107119
}

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

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,13 @@
1616

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

19-
import static org.springframework.ai.autoconfigure.openai.OpenAiProperties.CONFIG_PREFIX;
20-
2119
import java.time.Duration;
2220

2321
import org.springframework.boot.context.properties.ConfigurationProperties;
2422
import org.springframework.util.Assert;
2523
import org.springframework.util.StringUtils;
2624

27-
@ConfigurationProperties(CONFIG_PREFIX)
25+
@ConfigurationProperties(OpenAiProperties.CONFIG_PREFIX)
2826
public class OpenAiProperties {
2927

3028
public static final String CONFIG_PREFIX = "spring.ai.openai";
@@ -43,6 +41,8 @@ public class OpenAiProperties {
4341

4442
private String baseUrl = "https://api.openai.com";
4543

44+
private boolean retryEnabled = false;
45+
4646
public String getApiKey() {
4747
return this.apiKey;
4848
}
@@ -91,6 +91,14 @@ public Metadata getMetadata() {
9191
return this.metadata;
9292
}
9393

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

96104
private final OpenAiProperties openAiProperties;

0 commit comments

Comments
 (0)