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 13752cc44..1ae424227 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
@@ -32,7 +32,8 @@ import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.Coher
/**
* Cohere Embedding API.
- * https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere.html#model-parameters-embed
+ * AWS Bedrock Cohere Embedding API
+ * Based on the Cohere Embedding API
*
* @author Christian Tzolov
* @author Wei Jiang
@@ -151,22 +152,21 @@ public class CohereEmbeddingBedrockApi extends
}
/**
- * 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.
+ * 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 {
/**
- * (Default) Returns an error when the input exceeds the maximum input token length.
+ * Returns an error when the input exceeds the maximum input token length.
*/
NONE,
/**
- * Discard the start of the input.
+ * Discards the start of the input.
*/
- LEFT,
+ START,
/**
- * Discards the end of the input.
+ * (default) Discards the end of the input.
*/
- RIGHT
+ END
}
}
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 f96269fed..83afec90d 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
@@ -32,6 +32,7 @@ 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 = ".*")
@@ -56,4 +57,29 @@ public class CohereEmbeddingBedrockApiIT {
assertThat(response.embeddings().get(0)).hasSize(1024);
}
+ @Test
+ public void embedTextWithTruncate() {
+
+ CohereEmbeddingRequest request = new CohereEmbeddingRequest(
+ List.of("I like to eat apples", "I like to eat oranges"),
+ CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT, CohereEmbeddingRequest.Truncate.START);
+
+ CohereEmbeddingResponse response = api.embedding(request);
+
+ assertThat(response).isNotNull();
+ assertThat(response.texts()).isEqualTo(request.texts());
+ assertThat(response.embeddings()).hasSize(2);
+ assertThat(response.embeddings().get(0)).hasSize(1024);
+
+ request = new CohereEmbeddingRequest(List.of("I like to eat apples", "I like to eat oranges"),
+ CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT, CohereEmbeddingRequest.Truncate.END);
+
+ response = api.embedding(request);
+
+ assertThat(response).isNotNull();
+ assertThat(response.texts()).isEqualTo(request.texts());
+ assertThat(response.embeddings()).hasSize(2);
+ assertThat(response.embeddings().get(0)).hasSize(1024);
+ }
+
}
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 040dc25b1..49498f7c5 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
@@ -90,7 +90,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
"spring.ai.bedrock.aws.region=" + Region.EU_CENTRAL_1.id(),
"spring.ai.bedrock.cohere.embedding.model=MODEL_XYZ",
"spring.ai.bedrock.cohere.embedding.options.inputType=CLASSIFICATION",
- "spring.ai.bedrock.cohere.embedding.options.truncate=RIGHT")
+ "spring.ai.bedrock.cohere.embedding.options.truncate=START")
.withConfiguration(AutoConfigurations.of(BedrockCohereEmbeddingAutoConfiguration.class))
.run(context -> {
var properties = context.getBean(BedrockCohereEmbeddingProperties.class);
@@ -101,7 +101,7 @@ public class BedrockCohereEmbeddingAutoConfigurationIT {
assertThat(properties.getModel()).isEqualTo("MODEL_XYZ");
assertThat(properties.getOptions().getInputType()).isEqualTo(InputType.CLASSIFICATION);
- assertThat(properties.getOptions().getTruncate()).isEqualTo(CohereEmbeddingRequest.Truncate.RIGHT);
+ assertThat(properties.getOptions().getTruncate()).isEqualTo(CohereEmbeddingRequest.Truncate.START);
assertThat(awsProperties.getAccessKey()).isEqualTo("ACCESS_KEY");
assertThat(awsProperties.getSecretKey()).isEqualTo("SECRET_KEY");