Skip to content

Commit 220ec7f

Browse files
wmz7yeartzolov
authored andcommitted
AWS Bedrock Titan embedding model add amazon.titan-embed-text-v2:0 support.
1 parent cae045d commit 220ec7f

File tree

3 files changed

+32
-5
lines changed

3 files changed

+32
-5
lines changed

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,11 @@ public enum TitanEmbeddingModel {
153153
/**
154154
* amazon.titan-embed-text-v1
155155
*/
156-
TITAN_EMBED_TEXT_V1("amazon.titan-embed-text-v1");
156+
TITAN_EMBED_TEXT_V1("amazon.titan-embed-text-v1"),
157+
/**
158+
* amazon.titan-embed-text-v2
159+
*/
160+
TITAN_EMBED_TEXT_V2("amazon.titan-embed-text-v2:0");;
157161

158162
private final String id;
159163

models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/titan/api/TitanEmbeddingBedrockApiIT.java

Lines changed: 26 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,27 +21,33 @@
2121

2222
import org.junit.jupiter.api.Test;
2323
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
24+
25+
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
2426
import software.amazon.awssdk.regions.Region;
2527

2628
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel;
2729
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest;
2830
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse;
2931
import org.springframework.core.io.DefaultResourceLoader;
3032

33+
import com.fasterxml.jackson.databind.ObjectMapper;
34+
3135
import static org.assertj.core.api.Assertions.assertThat;
3236

3337
/**
3438
* @author Christian Tzolov
39+
* @author Wei Jiang
3540
*/
3641
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
3742
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
3843
public class TitanEmbeddingBedrockApiIT {
3944

4045
@Test
41-
public void embedText() {
46+
public void embedTextV1() {
4247

4348
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
44-
TitanEmbeddingModel.TITAN_EMBED_TEXT_V1.id(), Region.US_EAST_1.id(), Duration.ofMinutes(2));
49+
TitanEmbeddingModel.TITAN_EMBED_TEXT_V1.id(), EnvironmentVariableCredentialsProvider.create(),
50+
Region.US_EAST_1.id(), new ObjectMapper(), Duration.ofMinutes(2));
4551

4652
TitanEmbeddingRequest request = TitanEmbeddingRequest.builder().withInputText("I like to eat apples.").build();
4753

@@ -52,11 +58,28 @@ public void embedText() {
5258
assertThat(response.embedding()).hasSize(1536);
5359
}
5460

61+
@Test
62+
public void embedTextV2() {
63+
64+
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
65+
TitanEmbeddingModel.TITAN_EMBED_TEXT_V2.id(), EnvironmentVariableCredentialsProvider.create(),
66+
Region.US_WEST_2.id(), new ObjectMapper(), Duration.ofMinutes(2));
67+
68+
TitanEmbeddingRequest request = TitanEmbeddingRequest.builder().withInputText("I like to eat apples.").build();
69+
70+
TitanEmbeddingResponse response = titanEmbedApi.embedding(request);
71+
72+
assertThat(response).isNotNull();
73+
assertThat(response.inputTextTokenCount()).isEqualTo(7);
74+
assertThat(response.embedding()).hasSize(1024);
75+
}
76+
5577
@Test
5678
public void embedImage() throws IOException {
5779

5880
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
59-
TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id(), Duration.ofMinutes(2));
81+
TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), EnvironmentVariableCredentialsProvider.create(),
82+
Region.US_EAST_1.id(), new ObjectMapper(), Duration.ofMinutes(2));
6083

6184
byte[] image = new DefaultResourceLoader().getResource("classpath:/spring_framework.png")
6285
.getContentAsByteArray();

spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/bedrock-titan-embedding.adoc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ The prefix `spring.ai.bedrock.titan.embedding` (defined in `BedrockTitanEmbeddin
7878
| spring.ai.bedrock.titan.embedding.model | The model id to use. See the `TitanEmbeddingModel` for the supported models. | amazon.titan-embed-image-v1
7979
|====
8080

81-
Supported values are: `amazon.titan-embed-image-v1` and `amazon.titan-embed-text-v1`.
81+
Supported values are: `amazon.titan-embed-image-v1`, `amazon.titan-embed-text-v1` and `amazon.titan-embed-text-v2:0`.
8282
Model ID values can also be found in the https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html[AWS Bedrock documentation for base model IDs].
8383

8484
== Runtime Options [[embedding-options]]

0 commit comments

Comments
 (0)