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

Fix Bedrock Cohere embedding truncate type types

Browse Source 
- fix compilation errors and javadoc
This commit is contained in:
wmz7year
2024-05-06 22:48:41 +08:00
committed by Christian Tzolov
parent f955fd7277
commit 3497a1ec7b
3 changed files with 36 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java
View File

@@ -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
* <a href="https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-cohere.html#model-parameters-embed">AWS Bedrock Cohere Embedding API</a>
* Based on the <a href="https://docs.cohere.com/reference/embed">Cohere Embedding API</a>
*
* @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
}
}

26
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApiIT.java
View File

@@ -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);
}
}

4
spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/bedrock/cohere/BedrockCohereEmbeddingAutoConfigurationIT.java
View File

@@ -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");