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

fix: bedrock titan embeddings should return usage

Browse Source 
Auto-cherry-pick to 1.0.x

Signed-off-by: Gareth Evans <gareth@bryncynfelin.co.uk>
This commit is contained in:
Gareth Evans
2025-06-06 12:15:45 +01:00
committed by Ilayaperumal Gopinathan
parent 95b3df7a2b
commit c4e434a5e5
1 changed files with 15 additions and 1 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/titan/BedrockTitanEmbeddingModel.java
View File

@@ -28,12 +28,14 @@ import org.slf4j.LoggerFactory;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingRequest;
import org.springframework.ai.bedrock.titan.api.TitanEmbeddingBedrockApi.TitanEmbeddingResponse;
import org.springframework.ai.chat.metadata.DefaultUsage;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.AbstractEmbeddingModel;
import org.springframework.ai.embedding.Embedding;
import org.springframework.ai.embedding.EmbeddingOptions;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.ai.embedding.EmbeddingResponseMetadata;
import org.springframework.util.Assert;
/**
@@ -89,6 +91,7 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
List<Embedding> embeddings = new ArrayList<>();
var indexCounter = new AtomicInteger(0);
int tokenUsage = 0;
for (String inputContent : request.getInstructions()) {
var apiRequest = createTitanEmbeddingRequest(inputContent, request.getOptions());
@@ -111,6 +114,10 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
}
embeddings.add(new Embedding(response.embedding(), indexCounter.getAndIncrement()));
if (response.inputTextTokenCount() != null) {
tokenUsage += response.inputTextTokenCount();
}
}
catch (Exception ex) {
logger.error("Titan API embedding failed for input at index {}: {}", indexCounter.get(),
@@ -120,7 +127,10 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
}
}
return new EmbeddingResponse(embeddings);
EmbeddingResponseMetadata embeddingResponseMetadata = new EmbeddingResponseMetadata("",
getDefaultUsage(tokenUsage));
return new EmbeddingResponse(embeddings, embeddingResponseMetadata);
}
private TitanEmbeddingRequest createTitanEmbeddingRequest(String inputContent, EmbeddingOptions requestOptions) {
@@ -155,6 +165,10 @@ public class BedrockTitanEmbeddingModel extends AbstractEmbeddingModel {
return input.length() > 100 ? input.substring(0, 100) + "..." : input;
}
private DefaultUsage getDefaultUsage(int tokens) {
return new DefaultUsage(tokens, 0);
}
public enum InputType {
TEXT, IMAGE