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

Ensure matching enum id to value for multilingual cohere embedding model

Browse Source
This commit is contained in:
Bruno Oliveira
2024-12-05 14:24:50 +01:00
committed by Mark Pollack
parent 1453198c5b
commit 2aaa57fc7a
5 changed files with 72 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java
View File

@@ -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.
* <a href="https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere.html#model-parameters-embed">AWS Bedrock Cohere Embedding API</a>
* Based on the <a href="https://docs.cohere.com/reference/embed">Cohere Embedding API</a>
* Cohere Embedding API. <a href=
* "https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere.html#model-parameters-embed">AWS
* Bedrock Cohere Embedding API</a> Based on the
* <a href="https://docs.cohere.com/reference/embed">Cohere Embedding API</a>
*
* @author Christian Tzolov
* @author Wei Jiang
* @since 0.8.0
*/
public class CohereEmbeddingBedrockApi extends
AbstractBedrockApi<CohereEmbeddingRequest, CohereEmbeddingResponse, CohereEmbeddingResponse> {
public class CohereEmbeddingBedrockApi
extends AbstractBedrockApi<CohereEmbeddingRequest, CohereEmbeddingResponse, CohereEmbeddingResponse> {
/**
* 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<String> texts,
@JsonProperty("input_type") InputType inputType,
@JsonProperty("truncate") Truncate truncate) {
public record CohereEmbeddingRequest(@JsonProperty("texts") List<String> 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<float[]> embeddings,
@JsonProperty("texts") List<String> texts,
public record CohereEmbeddingResponse(@JsonProperty("id") String id,
@JsonProperty("embeddings") List<float[]> embeddings, @JsonProperty("texts") List<String> 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

2
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModelIT.java
View File

@@ -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));
}

2
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java
View File

@@ -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

2
spring-ai-spring-boot-autoconfigure/src/main/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingProperties.java
View File

@@ -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()

2
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java
View File

@@ -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));