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

Make BedrockCohereEmbeddingOptions implement EmbeddingOptions

Browse Source 
Resolves #389
This commit is contained in:
Christian Tzolov
2024-03-03 09:14:02 +01:00
parent cbb59ce616
commit 5e6815863c
2 changed files with 19 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

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

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

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