Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

AWS Bedrock Titan embedding model add amazon.titan-embed-text-v2:0 support.

Browse Source
This commit is contained in:
wmz7year
2024-05-01 13:03:23 +08:00
committed by Christian Tzolov
parent cae045d52b
commit 220ec7f631
3 changed files with 32 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -153,7 +153,11 @@ public class TitanEmbeddingBedrockApi extends
/**
* amazon.titan-embed-text-v1
*/
TITAN_EMBED_TEXT_V1("amazon.titan-embed-text-v1");
TITAN_EMBED_TEXT_V1("amazon.titan-embed-text-v1"),
/**
* amazon.titan-embed-text-v2
*/
TITAN_EMBED_TEXT_V2("amazon.titan-embed-text-v2:0");;
private final String id;

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

@@ -21,6 +21,8 @@ import java.util.Base64;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
import software.amazon.awssdk.regions.Region;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingModel;
@@ -28,20 +30,24 @@ import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEm
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse;
import org.springframework.core.io.DefaultResourceLoader;
import com.fasterxml.jackson.databind.ObjectMapper;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
* @author Wei Jiang
*/
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@EnabledIfEnvironmentVariable(named = "AWS_SECRET_ACCESS_KEY", matches = ".*")
public class TitanEmbeddingBedrockApiIT {
@Test
public void embedText() {
public void embedTextV1() {
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
TitanEmbeddingModel.TITAN_EMBED_TEXT_V1.id(), Region.US_EAST_1.id(), Duration.ofMinutes(2));
TitanEmbeddingModel.TITAN_EMBED_TEXT_V1.id(), EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(), new ObjectMapper(), Duration.ofMinutes(2));
TitanEmbeddingRequest request = TitanEmbeddingRequest.builder().withInputText("I like to eat apples.").build();
@@ -52,11 +58,28 @@ public class TitanEmbeddingBedrockApiIT {
assertThat(response.embedding()).hasSize(1536);
}
@Test
public void embedTextV2() {
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
TitanEmbeddingModel.TITAN_EMBED_TEXT_V2.id(), EnvironmentVariableCredentialsProvider.create(),
Region.US_WEST_2.id(), new ObjectMapper(), Duration.ofMinutes(2));
TitanEmbeddingRequest request = TitanEmbeddingRequest.builder().withInputText("I like to eat apples.").build();
TitanEmbeddingResponse response = titanEmbedApi.embedding(request);
assertThat(response).isNotNull();
assertThat(response.inputTextTokenCount()).isEqualTo(7);
assertThat(response.embedding()).hasSize(1024);
}
@Test
public void embedImage() throws IOException {
TitanEmbeddingBedrockApi titanEmbedApi = new TitanEmbeddingBedrockApi(
TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), Region.US_EAST_1.id(), Duration.ofMinutes(2));
TitanEmbeddingModel.TITAN_EMBED_IMAGE_V1.id(), EnvironmentVariableCredentialsProvider.create(),
Region.US_EAST_1.id(), new ObjectMapper(), Duration.ofMinutes(2));
byte[] image = new DefaultResourceLoader().getResource("classpath:/spring_framework.png")
.getContentAsByteArray();

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

@@ -78,7 +78,7 @@ The prefix `spring.ai.bedrock.titan.embedding` (defined in `BedrockTitanEmbeddin
| spring.ai.bedrock.titan.embedding.model | The model id to use. See the `TitanEmbeddingModel` for the supported models. | amazon.titan-embed-image-v1
|====
Supported values are: `amazon.titan-embed-image-v1` and `amazon.titan-embed-text-v1`.
Supported values are: `amazon.titan-embed-image-v1`, `amazon.titan-embed-text-v1` and `amazon.titan-embed-text-v2:0`.
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].
== Runtime Options [[embedding-options]]