Skip to content

Commit 216df35

Browse files
author
wmz7year
committed
AWS Bedrock embedding models add exponential backoff support.
1 parent 470d00e commit 216df35

File tree

7 files changed

+88
-36
lines changed

7 files changed

+88
-36
lines changed

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
import org.springframework.ai.embedding.EmbeddingRequest;
2929
import org.springframework.ai.embedding.EmbeddingResponse;
3030
import org.springframework.ai.model.ModelOptionsUtils;
31+
import org.springframework.ai.retry.RetryUtils;
32+
import org.springframework.retry.support.RetryTemplate;
3133
import org.springframework.util.Assert;
3234

3335
/**
@@ -44,6 +46,11 @@ public class BedrockCohereEmbeddingModel extends AbstractEmbeddingModel {
4446

4547
private final BedrockCohereEmbeddingOptions defaultOptions;
4648

49+
/**
50+
* The retry template used to retry the Bedrock API calls.
51+
*/
52+
private final RetryTemplate retryTemplate;
53+
4754
// private CohereEmbeddingRequest.InputType inputType =
4855
// CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT;
4956

@@ -60,10 +67,18 @@ public BedrockCohereEmbeddingModel(CohereEmbeddingBedrockApi cohereEmbeddingBedr
6067

6168
public BedrockCohereEmbeddingModel(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi,
6269
BedrockCohereEmbeddingOptions options) {
70+
this(cohereEmbeddingBedrockApi, options, RetryUtils.DEFAULT_RETRY_TEMPLATE);
71+
}
72+
73+
public BedrockCohereEmbeddingModel(CohereEmbeddingBedrockApi cohereEmbeddingBedrockApi,
74+
BedrockCohereEmbeddingOptions options, RetryTemplate retryTemplate) {
6375
Assert.notNull(cohereEmbeddingBedrockApi, "CohereEmbeddingBedrockApi must not be null");
64-
Assert.notNull(options, "BedrockCohereEmbeddingOptions must not be null");
76+
Assert.notNull(options, "DefaultOptions must not be null");
77+
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
78+
6579
this.embeddingApi = cohereEmbeddingBedrockApi;
6680
this.defaultOptions = options;
81+
this.retryTemplate = retryTemplate;
6782
}
6883

6984
// /**
@@ -104,13 +119,16 @@ public EmbeddingResponse call(EmbeddingRequest request) {
104119

105120
var apiRequest = new CohereEmbeddingRequest(request.getInstructions(), optionsToUse.getInputType(),
106121
optionsToUse.getTruncate());
107-
CohereEmbeddingResponse apiResponse = this.embeddingApi.embedding(apiRequest);
108-
var indexCounter = new AtomicInteger(0);
109-
List<Embedding> embeddings = apiResponse.embeddings()
110-
.stream()
111-
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
112-
.toList();
113-
return new EmbeddingResponse(embeddings);
122+
123+
return this.retryTemplate.execute(ctx -> {
124+
CohereEmbeddingResponse apiResponse = this.embeddingApi.embedding(apiRequest);
125+
var indexCounter = new AtomicInteger(0);
126+
List<Embedding> embeddings = apiResponse.embeddings()
127+
.stream()
128+
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
129+
.toList();
130+
return new EmbeddingResponse(embeddings);
131+
});
114132
}
115133

116134
/**

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/BedrockTitanEmbeddingModel.java‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
import org.springframework.ai.embedding.EmbeddingOptions;
3232
import org.springframework.ai.embedding.EmbeddingRequest;
3333
import org.springframework.ai.embedding.EmbeddingResponse;
34+
import org.springframework.ai.retry.RetryUtils;
35+
import org.springframework.retry.support.RetryTemplate;
3436
import org.springframework.util.Assert;
3537

3638
/**
@@ -50,6 +52,11 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
5052

5153
private final TitanEmbeddingBedrockApi embeddingApi;
5254

55+
/**
56+
* The retry template used to retry the Bedrock API calls.
57+
*/
58+
private final RetryTemplate retryTemplate;
59+
5360
public enum InputType {
5461

5562
TEXT, IMAGE
@@ -62,7 +69,15 @@ public enum InputType {
6269
private InputType inputType = InputType.TEXT;
6370

6471
public BedrockTitanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingBedrockApi) {
72+
this(titanEmbeddingBedrockApi, RetryUtils.DEFAULT_RETRY_TEMPLATE);
73+
}
74+
75+
public BedrockTitanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingBedrockApi, RetryTemplate retryTemplate) {
76+
Assert.notNull(titanEmbeddingBedrockApi, "TitanEmbeddingBedrockApi must not be null");
77+
Assert.notNull(retryTemplate, "RetryTemplate must not be null");
78+
6579
this.embeddingApi = titanEmbeddingBedrockApi;
80+
this.retryTemplate = retryTemplate;
6681
}
6782

6883
/**
@@ -87,17 +102,19 @@ public EmbeddingResponse call(EmbeddingRequest request) {
87102
"Titan Embedding does not support batch embedding. Will make multiple API calls to embed(Document)");
88103
}
89104

90-
List<List<Double>> embeddingList = new ArrayList<>();
91-
for (String inputContent : request.getInstructions()) {
92-
var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
93-
TitanEmbeddingResponse response = this.embeddingApi.embedding(apiRequest);
94-
embeddingList.add(response.embedding());
95-
}
96-
var indexCounter = new AtomicInteger(0);
97-
List<Embedding> embeddings = embeddingList.stream()
98-
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
99-
.toList();
100-
return new EmbeddingResponse(embeddings);
105+
return this.retryTemplate.execute(ctx -> {
106+
List<List<Double>> embeddingList = new ArrayList<>();
107+
for (String inputContent : request.getInstructions()) {
108+
var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
109+
TitanEmbeddingResponse response = this.embeddingApi.embedding(apiRequest);
110+
embeddingList.add(response.embedding());
111+
}
112+
var indexCounter = new AtomicInteger(0);
113+
List<Embedding> embeddings = embeddingList.stream()
114+
.map(e -> new Embedding(e, indexCounter.getAndIncrement()))
115+
.toList();
116+
return new EmbeddingResponse(embeddings);
117+
});
101118
}
102119

103120
private TitanEmbeddingRequest createTitanEmbeddingRequest(String inputContent, EmbeddingOptions requestOptions) {

‎models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApi.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -173,7 +173,7 @@ public enum TitanEmbeddingModel {
173173
/**
174174
* amazon.titan-embed-text-v2
175175
*/
176-
TITAN_EMBED_TEXT_V2("amazon.titan-embed-text-v2:0");;
176+
TITAN_EMBED_TEXT_V2("amazon.titan-embed-text-v2:0");
177177

178178
private final String id;
179179

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfiguration.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionConfiguration;
2323
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
24+
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
2425
import org.springframework.ai.bedrock.cohere.BedrockCohereEmbeddingModel;
2526
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
2627
import org.springframework.boot.autoconfigure.AutoConfiguration;
@@ -31,6 +32,7 @@
3132
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3233
import org.springframework.context.annotation.Bean;
3334
import org.springframework.context.annotation.Import;
35+
import org.springframework.retry.support.RetryTemplate;
3436

3537
/**
3638
* {@link AutoConfiguration Auto-configuration} for Bedrock Cohere Embedding Client.
@@ -39,7 +41,7 @@
3941
* @author Wei Jiang
4042
* @since 0.8.0
4143
*/
42-
@AutoConfiguration
44+
@AutoConfiguration(after = SpringAiRetryAutoConfiguration.class)
4345
@ConditionalOnClass(CohereEmbeddingBedrockApi.class)
4446
@EnableConfigurationProperties({ BedrockCohereEmbeddingProperties.class, BedrockAwsConnectionProperties.class })
4547
@ConditionalOnProperty(prefix = BedrockCohereEmbeddingProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@@ -59,9 +61,9 @@ public CohereEmbeddingBedrockApi cohereEmbeddingApi(BedrockCohereEmbeddingProper
5961
@ConditionalOnMissingBean
6062
@ConditionalOnBean(CohereEmbeddingBedrockApi.class)
6163
public BedrockCohereEmbeddingModel cohereEmbeddingModel(CohereEmbeddingBedrockApi cohereEmbeddingApi,
62-
BedrockCohereEmbeddingProperties properties) {
64+
BedrockCohereEmbeddingProperties properties, RetryTemplate retryTemplate) {
6365

64-
return new BedrockCohereEmbeddingModel(cohereEmbeddingApi, properties.getOptions());
66+
return new BedrockCohereEmbeddingModel(cohereEmbeddingApi, properties.getOptions(), retryTemplate);
6567
}
6668

6769
}

‎spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfiguration.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionConfiguration;
2323
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
24+
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
2425
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingModel;
2526
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
2627
import org.springframework.boot.autoconfigure.AutoConfiguration;
@@ -31,6 +32,7 @@
3132
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3233
import org.springframework.context.annotation.Bean;
3334
import org.springframework.context.annotation.Import;
35+
import org.springframework.retry.support.RetryTemplate;
3436

3537
/**
3638
* {@link AutoConfiguration Auto-configuration} for Bedrock Titan Embedding Client.
@@ -39,7 +41,7 @@
3941
* @author Wei Jiang
4042
* @since 0.8.0
4143
*/
42-
@AutoConfiguration
44+
@AutoConfiguration(after = SpringAiRetryAutoConfiguration.class)
4345
@ConditionalOnClass(TitanEmbeddingBedrockApi.class)
4446
@EnableConfigurationProperties({ BedrockTitanEmbeddingProperties.class, BedrockAwsConnectionProperties.class })
4547
@ConditionalOnProperty(prefix = BedrockTitanEmbeddingProperties.CONFIG_PREFIX, name = "enabled", havingValue = "true")
@@ -59,9 +61,10 @@ public TitanEmbeddingBedrockApi titanEmbeddingBedrockApi(BedrockTitanEmbeddingPr
5961
@ConditionalOnMissingBean
6062
@ConditionalOnBean(TitanEmbeddingBedrockApi.class)
6163
public BedrockTitanEmbeddingModel titanEmbeddingModel(TitanEmbeddingBedrockApi titanEmbeddingApi,
62-
BedrockTitanEmbeddingProperties properties) {
64+
BedrockTitanEmbeddingProperties properties, RetryTemplate retryTemplate) {
6365

64-
return new BedrockTitanEmbeddingModel(titanEmbeddingApi).withInputType(properties.getInputType());
66+
return new BedrockTitanEmbeddingModel(titanEmbeddingApi, retryTemplate)
67+
.withInputType(properties.getInputType());
6568
}
6669

6770
}

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import org.junit.jupiter.api.Test;
1919
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
2020
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
21+
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
2122
import org.springframework.ai.bedrock.cohere.BedrockCohereEmbeddingModel;
2223
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
2324
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
@@ -47,7 +48,8 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
4748
"spring.ai.bedrock.cohere.embedding.model=" + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(),
4849
"spring.ai.bedrock.cohere.embedding.options.inputType=SEARCH_DOCUMENT",
4950
"spring.ai.bedrock.cohere.embedding.options.truncate=NONE")
50-
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class));
51+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
52+
BedrockCohereEmbeddingAutoConfiguration.class));
5153

5254
@Test
5355
public void singleEmbedding() {
@@ -91,7 +93,8 @@ public void propertiesTest() {
9193
"spring.ai.bedrock.cohere.embedding.model=MODEL_XYZ",
9294
"spring.ai.bedrock.cohere.embedding.options.inputType=CLASSIFICATION",
9395
"spring.ai.bedrock.cohere.embedding.options.truncate=START")
94-
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
96+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
97+
BedrockCohereEmbeddingAutoConfiguration.class))
9598
.run(context -> {
9699
var properties = context.getBean(BedrockCohereEmbeddingProperties.class);
97100
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
@@ -113,23 +116,26 @@ public void embeddingDisabled() {
113116

114117
// It is disabled by default
115118
new ApplicationContextRunner()
116-
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
119+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
120+
BedrockCohereEmbeddingAutoConfiguration.class))
117121
.run(context -> {
118122
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isEmpty();
119123
assertThat(context.getBeansOfType(BedrockCohereEmbeddingModel.class)).isEmpty();
120124
});
121125

122126
// Explicitly enable the embedding auto-configuration.
123127
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true")
124-
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
128+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
129+
BedrockCohereEmbeddingAutoConfiguration.class))
125130
.run(context -> {
126131
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isNotEmpty();
127132
assertThat(context.getBeansOfType(BedrockCohereEmbeddingModel.class)).isNotEmpty();
128133
});
129134

130135
// Explicitly disable the embedding auto-configuration.
131136
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=false")
132-
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
137+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
138+
BedrockCohereEmbeddingAutoConfiguration.class))
133139
.run(context -> {
134140
assertThat(context.getBeansOfType(BedrockCohereEmbeddingProperties.class)).isEmpty();
135141
assertThat(context.getBeansOfType(BedrockCohereEmbeddingModel.class)).isEmpty();

‎spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/titan/BedrockTitanEmbeddingAutoConfigurationIT.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import software.amazon.awssdk.regions.Region;
2424

2525
import org.springframework.ai.autoconfigure.bedrock.BedrockAwsConnectionProperties;
26+
import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration;
2627
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingModel;
2728
import org.springframework.ai.bedrock.titan.BedrockTitanEmbeddingModel.InputType;
2829
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel;
@@ -47,7 +48,8 @@ public class BedrockTitanEmbeddingAutoConfigurationIT {
4748
"spring.ai.bedrock.aws.secret-key=" + System.getenv("AWS_SECRET_ACCESS_KEY"),
4849
"spring.ai.bedrock.aws.region=" + Region.US_EAST_1.id(),
4950
"spring.ai.bedrock.titan.embedding.model=" + TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id())
50-
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class));
51+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
52+
BedrockTitanEmbeddingAutoConfiguration.class));
5153

5254
@Test
5355
public void singleTextEmbedding() {
@@ -87,7 +89,8 @@ public void propertiesTest() {
8789
"spring.ai.bedrock.aws.access-key=ACCESS_KEY", "spring.ai.bedrock.aws.secret-key=SECRET_KEY",
8890
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
8991
"spring.ai.bedrock.titan.embedding.model=MODEL_XYZ", "spring.ai.bedrock.titan.embedding.inputType=TEXT")
90-
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
92+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
93+
BedrockTitanEmbeddingAutoConfiguration.class))
9194
.run(context -> {
9295
var properties = context.getBean(BedrockTitanEmbeddingProperties.class);
9396
var awsProperties = context.getBean(BedrockAwsConnectionProperties.class);
@@ -108,23 +111,26 @@ public void embeddingDisabled() {
108111

109112
// It is disabled by default
110113
new ApplicationContextRunner()
111-
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
114+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
115+
BedrockTitanEmbeddingAutoConfiguration.class))
112116
.run(context -> {
113117
assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isEmpty();
114118
assertThat(context.getBeansOfType(BedrockTitanEmbeddingModel.class)).isEmpty();
115119
});
116120

117121
// Explicitly enable the embedding auto-configuration.
118122
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=true")
119-
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
123+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
124+
BedrockTitanEmbeddingAutoConfiguration.class))
120125
.run(context -> {
121126
assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isNotEmpty();
122127
assertThat(context.getBeansOfType(BedrockTitanEmbeddingModel.class)).isNotEmpty();
123128
});
124129

125130
// Explicitly disable the embedding auto-configuration.
126131
new ApplicationContextRunner().withPropertyValues("spring.ai.bedrock.titan.embedding.enabled=false")
127-
.withConfiguration(AutoConfigurations.of(BedrockTitanEmbeddingAutoConfiguration.class))
132+
.withConfiguration(AutoConfigurations.of(SpringAiRetryAutoConfiguration.class,
133+
BedrockTitanEmbeddingAutoConfiguration.class))
128134
.run(context -> {
129135
assertThat(context.getBeansOfType(BedrockTitanEmbeddingProperties.class)).isEmpty();
130136
assertThat(context.getBeansOfType(BedrockTitanEmbeddingModel.class)).isEmpty();

0 commit comments

Comments
 (0)