From 2aaa57fc7abf097d40dc2bedb5c0242620615b03 Mon Sep 17 00:00:00 2001 From: Bruno Oliveira Date: Thu, 5 Dec 2024 14:24:50 +0100 Subject: [PATCH] Ensure matching enum id to value for multilingual cohere embedding model --- .../cohere/api/CohereEmbeddingBedrockApi.java | 126 ++++++++++-------- .../cohere/BedrockCohereEmbeddingModelIT.java | 2 +- .../api/CohereEmbeddingBedrockApiIT.java | 2 +- .../BedrockCohereEmbeddingProperties.java | 2 +- ...ockCohereEmbeddingAutoConfigurationIT.java | 2 +- 5 files changed, 72 insertions(+), 62 deletions(-) diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java index 2765dd4a8..3fd2a956a 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java @@ -16,8 +16,6 @@ package org.springframework.ai.bedrock.cohere.api; -// @formatter:off - import java.time.Duration; import java.util.List; @@ -33,22 +31,23 @@ import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.Coher import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingResponse; /** - * Cohere Embedding API. - * AWS Bedrock Cohere Embedding API - * Based on the Cohere Embedding API + * Cohere Embedding API. AWS + * Bedrock Cohere Embedding API Based on the + * Cohere Embedding API * * @author Christian Tzolov * @author Wei Jiang * @since 0.8.0 */ -public class CohereEmbeddingBedrockApi extends - AbstractBedrockApi { +public class CohereEmbeddingBedrockApi + extends AbstractBedrockApi { /** - * Create a new CohereEmbeddingBedrockApi instance using the default credentials provider chain, the default object - * mapper, default temperature and topP values. - * - * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the supported models. + * Create a new CohereEmbeddingBedrockApi instance using the default credentials + * provider chain, the default object mapper, default temperature and topP values. + * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the + * supported models. * @param region The AWS region to use. */ public CohereEmbeddingBedrockApi(String modelId, String region) { @@ -56,13 +55,14 @@ public class CohereEmbeddingBedrockApi extends } /** - * Create a new CohereEmbeddingBedrockApi instance using the provided credentials provider, region and object - * mapper. - * - * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the supported models. + * Create a new CohereEmbeddingBedrockApi instance using the provided credentials + * provider, region and object mapper. + * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the + * supported models. * @param credentialsProvider The credentials provider to connect to AWS. * @param region The AWS region to use. - * @param objectMapper The object mapper to use for JSON serialization and deserialization. + * @param objectMapper The object mapper to use for JSON serialization and + * deserialization. */ public CohereEmbeddingBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region, ObjectMapper objectMapper) { @@ -70,10 +70,10 @@ public class CohereEmbeddingBedrockApi extends } /** - * Create a new CohereEmbeddingBedrockApi instance using the default credentials provider chain, the default object - * mapper, default temperature and topP values. - * - * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the supported models. + * Create a new CohereEmbeddingBedrockApi instance using the default credentials + * provider chain, the default object mapper, default temperature and topP values. + * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the + * supported models. * @param region The AWS region to use. * @param timeout The timeout to use. */ @@ -82,13 +82,14 @@ public class CohereEmbeddingBedrockApi extends } /** - * Create a new CohereEmbeddingBedrockApi instance using the provided credentials provider, region and object - * mapper. - * - * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the supported models. + * Create a new CohereEmbeddingBedrockApi instance using the provided credentials + * provider, region and object mapper. + * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the + * supported models. * @param credentialsProvider The credentials provider to connect to AWS. * @param region The AWS region to use. - * @param objectMapper The object mapper to use for JSON serialization and deserialization. + * @param objectMapper The object mapper to use for JSON serialization and + * deserialization. * @param timeout The timeout to use. */ public CohereEmbeddingBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, String region, @@ -97,13 +98,14 @@ public class CohereEmbeddingBedrockApi extends } /** - * Create a new CohereEmbeddingBedrockApi instance using the provided credentials provider, region and object - * mapper. - * - * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the supported models. + * Create a new CohereEmbeddingBedrockApi instance using the provided credentials + * provider, region and object mapper. + * @param modelId The model id to use. See the {@link CohereEmbeddingModel} for the + * supported models. * @param credentialsProvider The credentials provider to connect to AWS. * @param region The AWS region to use. - * @param objectMapper The object mapper to use for JSON serialization and deserialization. + * @param objectMapper The object mapper to use for JSON serialization and + * deserialization. * @param timeout The timeout to use. */ public CohereEmbeddingBedrockApi(String modelId, AwsCredentialsProvider credentialsProvider, Region region, @@ -117,13 +119,15 @@ public class CohereEmbeddingBedrockApi extends } /** - * Cohere Embedding model ids. https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html + * Cohere Embedding model ids. + * https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html */ public enum CohereEmbeddingModel { + /** * cohere.embed-multilingual-v3 */ - COHERE_EMBED_MULTILINGUAL_V1("cohere.embed-multilingual-v3"), + COHERE_EMBED_MULTILINGUAL_V3("cohere.embed-multilingual-v3"), /** * cohere.embed-english-v3 */ @@ -147,20 +151,20 @@ public class CohereEmbeddingBedrockApi extends /** * The Cohere Embed model request. * - * @param texts An array of strings for the model to embed. For optimal performance, we recommend reducing the - * length of each text to less than 512 tokens. 1 token is about 4 characters. - * @param inputType Prepends special tokens to differentiate each type from one another. You should not mix - * different types together, except when mixing types for search and retrieval. In this case, embed your corpus - * with the search_document type and embedded queries with type search_query type. - * @param truncate Specifies how the API handles inputs longer than the maximum token length. If you specify LEFT or - * RIGHT, the model discards the input until the remaining input is exactly the maximum input token length for the - * model. + * @param texts An array of strings for the model to embed. For optimal performance, + * we recommend reducing the length of each text to less than 512 tokens. 1 token is + * about 4 characters. + * @param inputType Prepends special tokens to differentiate each type from one + * another. You should not mix different types together, except when mixing types for + * search and retrieval. In this case, embed your corpus with the search_document type + * and embedded queries with type search_query type. + * @param truncate Specifies how the API handles inputs longer than the maximum token + * length. If you specify LEFT or RIGHT, the model discards the input until the + * remaining input is exactly the maximum input token length for the model. */ @JsonInclude(Include.NON_NULL) - public record CohereEmbeddingRequest( - @JsonProperty("texts") List texts, - @JsonProperty("input_type") InputType inputType, - @JsonProperty("truncate") Truncate truncate) { + public record CohereEmbeddingRequest(@JsonProperty("texts") List texts, + @JsonProperty("input_type") InputType inputType, @JsonProperty("truncate") Truncate truncate) { /** * Cohere Embedding API input types. @@ -168,8 +172,8 @@ public class CohereEmbeddingBedrockApi extends public enum InputType { /** - * In search use-cases, use search_document when you encode documents for embeddings that you store in a - * vector database. + * In search use-cases, use search_document when you encode documents for + * embeddings that you store in a vector database. */ @JsonProperty("search_document") SEARCH_DOCUMENT, @@ -188,12 +192,17 @@ public class CohereEmbeddingBedrockApi extends */ @JsonProperty("clustering") CLUSTERING + } /** - * Specifies how the API handles inputs longer than the maximum token length. Passing START will discard the start of the input. END will discard the end of the input. In both cases, input is discarded until the remaining input is exactly the maximum input token length for the model. + * Specifies how the API handles inputs longer than the maximum token length. + * Passing START will discard the start of the input. END will discard the end of + * the input. In both cases, input is discarded until the remaining input is + * exactly the maximum input token length for the model. */ public enum Truncate { + /** * Returns an error when the input exceeds the maximum input token length. */ @@ -206,6 +215,7 @@ public class CohereEmbeddingBedrockApi extends * (default) Discards the end of the input. */ END + } } @@ -213,22 +223,22 @@ public class CohereEmbeddingBedrockApi extends * Cohere Embedding response. * * @param id An identifier for the response. - * @param embeddings An array of embeddings, where each embedding is an array of floats with 1024 elements. The - * length of the embeddings array will be the same as the length of the original texts array. - * @param texts An array containing the text entries for which embeddings were returned. + * @param embeddings An array of embeddings, where each embedding is an array of + * floats with 1024 elements. The length of the embeddings array will be the same as + * the length of the original texts array. + * @param texts An array containing the text entries for which embeddings were + * returned. * @param responseType The type of the response. The value is always embeddings. - * @param amazonBedrockInvocationMetrics Bedrock invocation metrics. Currently bedrock doesn't return - * invocationMetrics for the cohere embedding model. + * @param amazonBedrockInvocationMetrics Bedrock invocation metrics. Currently bedrock + * doesn't return invocationMetrics for the cohere embedding model. */ @JsonInclude(Include.NON_NULL) - public record CohereEmbeddingResponse( - @JsonProperty("id") String id, - @JsonProperty("embeddings") List embeddings, - @JsonProperty("texts") List texts, + public record CohereEmbeddingResponse(@JsonProperty("id") String id, + @JsonProperty("embeddings") List embeddings, @JsonProperty("texts") List texts, @JsonProperty("response_type") String responseType, - // For future use: Currently bedrock doesn't return invocationMetrics for the cohere embedding model. + // For future use: Currently bedrock doesn't return invocationMetrics for the + // cohere embedding model. @JsonProperty("amazon-bedrock-invocationMetrics") AmazonBedrockInvocationMetrics amazonBedrockInvocationMetrics) { } } -// @formatter:on diff --git a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModelIT.java b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModelIT.java index 1de587388..04b76646b 100644 --- a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModelIT.java +++ b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModelIT.java @@ -169,7 +169,7 @@ class BedrockCohereEmbeddingModelIT { @Bean public CohereEmbeddingBedrockApi cohereEmbeddingApi() { - return new CohereEmbeddingBedrockApi(CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(), + return new CohereEmbeddingBedrockApi(CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V3.id(), EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper(), Duration.ofMinutes(2)); } diff --git a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java index b037e7598..b45266184 100644 --- a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java +++ b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java @@ -39,7 +39,7 @@ import static org.assertj.core.api.Assertions.assertThat; public class CohereEmbeddingBedrockApiIT { CohereEmbeddingBedrockApi api = new CohereEmbeddingBedrockApi( - CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(), EnvironmentVariableCredentialsProvider.create(), + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V3.id(), EnvironmentVariableCredentialsProvider.create(), Region.US_EAST_1.id(), new ObjectMapper(), Duration.ofMinutes(2)); @Test diff --git a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java index 9e2752dfc..f6f6b3a95 100644 --- a/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java +++ b/spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java @@ -43,7 +43,7 @@ public class BedrockCohereEmbeddingProperties { * Bedrock Cohere Embedding generative name. Defaults to * 'cohere.embed-multilingual-v3'. */ - private String model = CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(); + private String model = CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V3.id(); @NestedConfigurationProperty private BedrockCohereEmbeddingOptions options = BedrockCohereEmbeddingOptions.builder() diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java index 90f26cd50..0ed6cddf4 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java @@ -44,7 +44,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT { private final ApplicationContextRunner contextRunner = BedrockTestUtils.getContextRunner() .withPropertyValues("spring.ai.bedrock.cohere.embedding.enabled=true", - "spring.ai.bedrock.cohere.embedding.model=" + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V1.id(), + "spring.ai.bedrock.cohere.embedding.model=" + CohereEmbeddingModel.COHERE_EMBED_MULTILINGUAL_V3.id(), "spring.ai.bedrock.cohere.embedding.options.inputType=SEARCH_DOCUMENT", "spring.ai.bedrock.cohere.embedding.options.truncate=NONE") .withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class));