From 5e6815863c8520da8a0b29eb564bcb9978dcc641 Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Sun, 3 Mar 2024 09:14:02 +0100 Subject: [PATCH] Make BedrockCohereEmbeddingOptions implement EmbeddingOptions Resolves #389 --- .../cohere/BedrockCohereEmbeddingOptions.java | 3 ++- .../cohere/BedrockCohereEmbeddingClientIT.java | 17 +++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java index 68594d40a..c79c73bc8 100644 --- a/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java +++ b/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java @@ -22,12 +22,13 @@ import com.fasterxml.jackson.annotation.JsonProperty; import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType; import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.Truncate; +import org.springframework.ai.embedding.EmbeddingOptions; /** * @author Christian Tzolov */ @JsonInclude(Include.NON_NULL) -public class BedrockCohereEmbeddingOptions { +public class BedrockCohereEmbeddingOptions implements EmbeddingOptions { // @formatter:off /** diff --git a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java index 79a57290b..f00e2d742 100644 --- a/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java +++ b/models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java @@ -10,6 +10,8 @@ import software.amazon.awssdk.regions.Region; import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi; import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel; +import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType; +import org.springframework.ai.embedding.EmbeddingRequest; import org.springframework.ai.embedding.EmbeddingResponse; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.boot.SpringBootConfiguration; @@ -49,6 +51,21 @@ class BedrockCohereEmbeddingClientIT { assertThat(embeddingClient.dimensions()).isEqualTo(1024); } + @Test + void embeddingWthOptions() { + assertThat(embeddingClient).isNotNull(); + EmbeddingResponse embeddingResponse = embeddingClient + .call(new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"), + BedrockCohereEmbeddingOptions.builder().withInputType(InputType.SEARCH_DOCUMENT).build())); + assertThat(embeddingResponse.getResults()).hasSize(2); + assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty(); + assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0); + assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty(); + assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1); + + assertThat(embeddingClient.dimensions()).isEqualTo(1024); + } + @SpringBootConfiguration public static class TestConfiguration {