From 3497a1ec7b1aa57141652fb18b94e79063310e45 Mon Sep 17 00:00:00 2001 From: wmz7year Date: Mon, 6 May 2024 22:48:41 +0800 Subject: [PATCH] Fix Bedrock Cohere embedding truncate type types - fix compilation errors and javadoc --- .../cohere/api/CohereEmbeddingBedrockApi.java | 16 ++++++------ .../api/CohereEmbeddingBedrockApiIT.java | 26 +++++++++++++++++++ ...ockCohereEmbeddingAutoConfigurationIT.java | 4 +-- 3 files changed, 36 insertions(+), 10 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 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");