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

Add truncation support for Cohere embeddings

Browse Source 
Fixes: #1753

https://github.com/spring-projects/spring-ai/issues/1753

- Add character-based truncation (max 2048 chars) for Cohere embedding requests
- Support both START and END truncation strategies
- Add unit tests verifying truncation behavior for both strategies

Truncation is applied before sending requests to Bedrock API to avoid
ValidationException when text exceeds maximum length. The END strategy
(default) keeps the first 2048 characters while START keeps the last
2048 characters.
This commit is contained in:
Soby Chacko
2024-11-19 21:01:55 -05:00
committed by Mark Pollack
parent 551206feb0
commit 1abfd9ab96
3 changed files with 133 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModel.java
View File

@@ -18,6 +18,7 @@ package org.springframework.ai.bedrock.cohere;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi;
import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.CohereEmbeddingRequest;
@@ -37,10 +38,13 @@ import org.springframework.util.Assert;
* this API. If this change in the future we will add it as metadata.
*
* @author Christian Tzolov
* @author Soby Chacko
* @since 0.8.0
*/
public class BedrockCohereEmbeddingModel extends AbstractEmbeddingModel {
private static final int COHERE_MAX_CHARACTERS = 2048;
private final CohereEmbeddingBedrockApi embeddingApi;
private final BedrockCohereEmbeddingOptions defaultOptions;
@@ -74,11 +78,34 @@ public class BedrockCohereEmbeddingModel extends AbstractEmbeddingModel {
@Override
public EmbeddingResponse call(EmbeddingRequest request) {
Assert.notEmpty(request.getInstructions(), "At least one text is required!");
List<String> instructions = request.getInstructions();
Assert.notEmpty(instructions, "At least one text is required!");
final BedrockCohereEmbeddingOptions optionsToUse = this.mergeOptions(request.getOptions());
var apiRequest = new CohereEmbeddingRequest(request.getInstructions(), optionsToUse.getInputType(),
List<String> truncatedInstructions = instructions.stream().map(text -> {
if (text == null || text.isEmpty()) {
return text;
}
if (text.length() <= COHERE_MAX_CHARACTERS) {
return text;
}
// Handle truncation based on option
return switch (optionsToUse.getTruncate()) {
case END -> text.substring(0, COHERE_MAX_CHARACTERS); // Keep first 2048 chars
case START -> text.substring(text.length() - COHERE_MAX_CHARACTERS); // Keep
// last
// 2048
// chars
default -> text.substring(0, COHERE_MAX_CHARACTERS); // Default to END
// behavior
};
}).collect(Collectors.toList());
var apiRequest = new CohereEmbeddingRequest(truncatedInstructions, optionsToUse.getInputType(),
optionsToUse.getTruncate());
CohereEmbeddingResponse apiResponse = this.embeddingApi.embedding(apiRequest);
var indexCounter = new AtomicInteger(0);

104
models/spring-ai-bedrock/src/test/java/org/springframework/ai/bedrock/cohere/BedrockCohereEmbeddingModelIT.java
View File

@@ -22,6 +22,7 @@ import java.util.List;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.mockito.ArgumentCaptor;
import software.amazon.awssdk.auth.credentials.EnvironmentVariableCredentialsProvider;
import software.amazon.awssdk.regions.Region;
@@ -31,11 +32,14 @@ import org.springframework.ai.bedrock.cohere.api.CohereEmbeddingBedrockApi.Coher
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.mock.mockito.SpyBean;
import org.springframework.context.annotation.Bean;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.verify;
@SpringBootTest
@EnabledIfEnvironmentVariable(named = "AWS_ACCESS_KEY_ID", matches = ".*")
@@ -45,6 +49,13 @@ class BedrockCohereEmbeddingModelIT {
@Autowired
private BedrockCohereEmbeddingModel embeddingModel;
@SpyBean
private CohereEmbeddingBedrockApi embeddingApi;
@Autowired
@Qualifier("embeddingModelStartTruncate")
private BedrockCohereEmbeddingModel embeddingModelStartTruncate;
@Test
void singleEmbedding() {
assertThat(this.embeddingModel).isNotNull();
@@ -54,6 +65,77 @@ class BedrockCohereEmbeddingModelIT {
assertThat(this.embeddingModel.dimensions()).isEqualTo(1024);
}
@Test
void truncatesLongText() {
String longText = "Hello World".repeat(300);
assertThat(longText.length()).isGreaterThan(2048);
EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(longText));
assertThat(embeddingResponse.getResults()).hasSize(1);
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
assertThat(this.embeddingModel.dimensions()).isEqualTo(1024);
}
@Test
void truncatesMultipleLongTexts() {
String longText1 = "Hello World".repeat(300);
String longText2 = "Another Text".repeat(300);
EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(longText1, longText2));
assertThat(embeddingResponse.getResults()).hasSize(2);
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
assertThat(embeddingResponse.getResults().get(1).getOutput()).isNotEmpty();
assertThat(this.embeddingModel.dimensions()).isEqualTo(1024);
}
@Test
void verifyExactTruncationLength() {
String longText = "x".repeat(3000);
ArgumentCaptor<CohereEmbeddingBedrockApi.CohereEmbeddingRequest> requestCaptor = ArgumentCaptor
.forClass(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.class);
EmbeddingResponse embeddingResponse = embeddingModel.embedForResponse(List.of(longText));
verify(embeddingApi).embedding(requestCaptor.capture());
CohereEmbeddingBedrockApi.CohereEmbeddingRequest capturedRequest = requestCaptor.getValue();
assertThat(capturedRequest.texts()).hasSize(1);
assertThat(capturedRequest.texts().get(0).length()).isLessThanOrEqualTo(2048);
assertThat(embeddingResponse.getResults()).hasSize(1);
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
}
@Test
void truncatesLongTextFromStart() {
String startMarker = "START_MARKER_";
String endMarker = "_END_MARKER";
String middlePadding = "x".repeat(2500); // Long enough to force truncation
String longText = startMarker + middlePadding + endMarker;
assertThat(longText.length()).isGreaterThan(2048);
ArgumentCaptor<CohereEmbeddingBedrockApi.CohereEmbeddingRequest> requestCaptor = ArgumentCaptor
.forClass(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.class);
EmbeddingResponse embeddingResponse = this.embeddingModelStartTruncate.embedForResponse(List.of(longText));
// Verify truncation behavior
verify(embeddingApi).embedding(requestCaptor.capture());
String truncatedText = requestCaptor.getValue().texts().get(0);
assertThat(truncatedText.length()).isLessThanOrEqualTo(2048);
assertThat(truncatedText).doesNotContain(startMarker);
assertThat(truncatedText).endsWith(endMarker);
// Verify embedding response
assertThat(embeddingResponse.getResults()).hasSize(1);
assertThat(embeddingResponse.getResults().get(0).getOutput()).isNotEmpty();
assertThat(this.embeddingModelStartTruncate.dimensions()).isEqualTo(1024);
}
@Test
void batchEmbedding() {
assertThat(this.embeddingModel).isNotNull();
@@ -93,9 +175,27 @@ class BedrockCohereEmbeddingModelIT {
Duration.ofMinutes(2));
}
@Bean
@Bean("embeddingModel")
public BedrockCohereEmbeddingModel cohereAiEmbedding(CohereEmbeddingBedrockApi cohereEmbeddingApi) {
return new BedrockCohereEmbeddingModel(cohereEmbeddingApi);
// custom model that uses the END truncation strategy, instead of the default
// NONE.
return new BedrockCohereEmbeddingModel(cohereEmbeddingApi,
BedrockCohereEmbeddingOptions.builder()
.withInputType(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT)
.withTruncate(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.Truncate.END)
.build());
}
@Bean("embeddingModelStartTruncate")
public BedrockCohereEmbeddingModel cohereAiEmbeddingStartTruncate(
CohereEmbeddingBedrockApi cohereEmbeddingApi) {
// custom model that uses the START truncation strategy, instead of the
// default NONE.
return new BedrockCohereEmbeddingModel(cohereEmbeddingApi,
BedrockCohereEmbeddingOptions.builder()
.withInputType(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.InputType.SEARCH_DOCUMENT)
.withTruncate(CohereEmbeddingBedrockApi.CohereEmbeddingRequest.Truncate.START)
.build());
}
}

2
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/embeddings/bedrock-cohere-embedding.adoc
View File

@@ -73,6 +73,8 @@ The prefix `spring.ai.bedrock.cohere.embedding` (defined in `BedrockCohereEmbedd
| spring.ai.bedrock.cohere.embedding.options.truncate | 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. | NONE
|====
NOTE: When accessing Cohere via Amazon Bedrock, the functionality of truncating is not available. This is an issue with Amazon Bedrock. The Spring AI class `BedrockCohereEmbeddingModel` will truncate to 2048 character length, which is the maximum supported by the model.
Look at the https://github.com/spring-projects/spring-ai/blob/056b95a00efa5b014a1f488329fbd07a46c02378/models/spring-ai-bedrock/src/main/java/org/springframework/ai/bedrock/cohere/api/CohereEmbeddingBedrockApi.java#L150[CohereEmbeddingModel] for other model IDs.
Supported values are: `cohere.embed-multilingual-v3` and `cohere.embed-english-v3`.
Model ID values can also be found in the https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids-arns.html[AWS Bedrock documentation for base model IDs].