From 2517ca80200d47d07ee77ca4d2e30ed4c439dcac Mon Sep 17 00:00:00 2001 From: rmalara Date: Mon, 31 Mar 2025 17:48:51 -0300 Subject: [PATCH] GH-2609: Fix the thread leak issue in VertexAiTextEmbeddingModel 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 Signed-off-by: Rodrigo Malara --- .../text/VertexAiTextEmbeddingModel.java | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java index c2cc0ea5c..836137095 100644 --- a/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java +++ b/models/spring-ai-vertex-ai-embedding/src/main/java/org/springframework/ai/vertexai/embedding/text/VertexAiTextEmbeddingModel.java @@ -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 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 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; }); }