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

GH-2609: Fix the thread leak issue in VertexAiTextEmbeddingModel

Browse Source 
Fixes: #2609

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

The PredictionServiceClient was not being closed. Connections are kept open preventing resources from being disposed properly.

Signed-off-by: rmalara <rmalara@interactions.com>
Signed-off-by: Rodrigo Malara <rodrigomalara@gmail.com>
This commit is contained in:
rmalara
2025-03-31 17:48:51 -03:00
committed by Soby Chacko
parent 2294c5a807
commit 2517ca8020
1 changed files with 26 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2024 the original author or authors.
* Copyright 2023-2025 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -58,6 +58,7 @@ import org.springframework.util.StringUtils;
*
* @author Christian Tzolov
* @author Mark Pollack
* @author Rodrigo Malara
* @since 1.0.0
*/
public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel {
@@ -128,37 +129,38 @@ public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel {
.observation(this.observationConvention, DEFAULT_OBSERVATION_CONVENTION, () -> observationContext,
this.observationRegistry)
.observe(() -> {
PredictionServiceClient client = createPredictionServiceClient();
try (PredictionServiceClient client = createPredictionServiceClient()) {
EndpointName endpointName = this.connectionDetails.getEndpointName(finalOptions.getModel());
EndpointName endpointName = this.connectionDetails.getEndpointName(finalOptions.getModel());
PredictRequest.Builder predictRequestBuilder = getPredictRequestBuilder(request, endpointName,
finalOptions);
PredictRequest.Builder predictRequestBuilder = getPredictRequestBuilder(request, endpointName,
finalOptions);
PredictResponse embeddingResponse = this.retryTemplate
.execute(context -> getPredictResponse(client, predictRequestBuilder));
PredictResponse embeddingResponse = this.retryTemplate
.execute(context -> getPredictResponse(client, predictRequestBuilder));
int index = 0;
int totalTokenCount = 0;
List<Embedding> embeddingList = new ArrayList<>();
for (Value prediction : embeddingResponse.getPredictionsList()) {
Value embeddings = prediction.getStructValue().getFieldsOrThrow("embeddings");
Value statistics = embeddings.getStructValue().getFieldsOrThrow("statistics");
Value tokenCount = statistics.getStructValue().getFieldsOrThrow("token_count");
totalTokenCount = totalTokenCount + (int) tokenCount.getNumberValue();
int index = 0;
int totalTokenCount = 0;
List<Embedding> embeddingList = new ArrayList<>();
for (Value prediction : embeddingResponse.getPredictionsList()) {
Value embeddings = prediction.getStructValue().getFieldsOrThrow("embeddings");
Value statistics = embeddings.getStructValue().getFieldsOrThrow("statistics");
Value tokenCount = statistics.getStructValue().getFieldsOrThrow("token_count");
totalTokenCount = totalTokenCount + (int) tokenCount.getNumberValue();
Value values = embeddings.getStructValue().getFieldsOrThrow("values");
Value values = embeddings.getStructValue().getFieldsOrThrow("values");
float[] vectorValues = VertexAiEmbeddingUtils.toVector(values);
float[] vectorValues = VertexAiEmbeddingUtils.toVector(values);
embeddingList.add(new Embedding(vectorValues, index++));
embeddingList.add(new Embedding(vectorValues, index++));
}
EmbeddingResponse response = new EmbeddingResponse(embeddingList,
generateResponseMetadata(finalOptions.getModel(), totalTokenCount));
observationContext.setResponse(response);
return response;
}
EmbeddingResponse response = new EmbeddingResponse(embeddingList,
generateResponseMetadata(finalOptions.getModel(), totalTokenCount));
observationContext.setResponse(response);
return response;
});
}