GH-2609: Fix the thread leak issue in VertexAiTextEmbeddingModel

Fixes: #2609

#2609

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

Signed-off-by: rmalara <[email protected]>
Signed-off-by: Rodrigo Malara <[email protected]>

Author: rmalara

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 @@
  *
  * @author Christian Tzolov
  * @author Mark Pollack
+ * @author Rodrigo Malara
  * @since 1.0.0
  */
 public class VertexAiTextEmbeddingModel extends AbstractEmbeddingModel {
@@ -128,37 +129,38 @@ public EmbeddingResponse call(EmbeddingRequest request) {
 			.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++));
-				}
-				EmbeddingResponse response = new EmbeddingResponse(embeddingList,
-						generateResponseMetadata(finalOptions.getModel(), totalTokenCount));
+						embeddingList.add(new Embedding(vectorValues, index++));
+					}
+					EmbeddingResponse response = new EmbeddingResponse(embeddingList,
+							generateResponseMetadata(finalOptions.getModel(), totalTokenCount));
 
-				observationContext.setResponse(response);
+					observationContext.setResponse(response);
 
-				return response;
+					return response;
+				}
 			});
 	}