Skip to content

Commit 5e68158

Browse files
committed
Make BedrockCohereEmbeddingOptions implement EmbeddingOptions
Resolves #389
1 parent cbb59ce commit 5e68158

File tree

2 files changed

+19
-1
lines changed

2 files changed

+19
-1
lines changed

models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingOptions.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,13 @@
2222

2323
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType;
2424
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.Truncate;
25+
import org.springframework.ai.embedding.EmbeddingOptions;
2526

2627
/**
2728
* @author Christian Tzolov
2829
*/
2930
@JsonInclude(Include.NON_NULL)
30-
public class BedrockCohereEmbeddingOptions {
31+
public class BedrockCohereEmbeddingOptions implements EmbeddingOptions {
3132

3233
// @formatter:off
3334
/**

models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingClientIT.java

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010

1111
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
1212
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingModel;
13+
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType;
14+
import org.springframework.ai.embedding.EmbeddingRequest;
1315
import org.springframework.ai.embedding.EmbeddingResponse;
1416
import org.springframework.beans.factory.annotation.Autowired;
1517
import org.springframework.boot.SpringBootConfiguration;
@@ -49,6 +51,21 @@ void batchEmbedding() {
4951
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
5052
}
5153

54+
@Test
55+
void embeddingWthOptions() {
56+
assertThat(embeddingClient).isNotNull();
57+
EmbeddingResponse embeddingResponse = embeddingClient
58+
.call(new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
59+
BedrockCohereEmbeddingOptions.builder().withInputType(InputType.SEARCH_DOCUMENT).build()));
60+
assertThat(embeddingResponse.getResults()).hasSize(2);
61+
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
62+
assertThat(embeddingResponse.getResults().get(0).getIndex()).isEqualTo(0);
63+
assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty();
64+
assertThat(embeddingResponse.getResults().get(1).getIndex()).isEqualTo(1);
65+
66+
assertThat(embeddingClient.dimensions()).isEqualTo(1024);
67+
}
68+
5269
@SpringBootConfiguration
5370
public static class TestConfiguration {
5471

0 commit comments

Comments
 (0)