From 8070f75ebff05c97915fe6a8f9814fed5168e572 Mon Sep 17 00:00:00 2001 From: carlosdelest Date: Wed, 14 Feb 2024 15:52:31 +0100 Subject: [PATCH] Fix merge from main --- .../bulk/BulkShardRequestInferenceProvider.java | 13 ++++++++++--- .../services/cohere/CohereServiceSettings.java | 2 +- .../huggingface/HuggingFaceServiceSettings.java | 2 +- .../embeddings/OpenAiEmbeddingsServiceSettings.java | 2 +- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/action/bulk/BulkShardRequestInferenceProvider.java b/server/src/main/java/org/elasticsearch/action/bulk/BulkShardRequestInferenceProvider.java index 0315b44ca2537..4e787e96df3bb 100644 --- a/server/src/main/java/org/elasticsearch/action/bulk/BulkShardRequestInferenceProvider.java +++ b/server/src/main/java/org/elasticsearch/action/bulk/BulkShardRequestInferenceProvider.java @@ -27,6 +27,7 @@ import org.elasticsearch.inference.ModelSettings; import java.util.ArrayList; +import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -256,9 +257,9 @@ private void performInferenceOnBulkItemRequest( } // TODO We won't batch until chunkedInfer has batching support for (String inferenceFieldName : inferenceFieldNames) { - ActionListener inferenceResultsListener = new ActionListener<>() { + ActionListener> inferenceResultsListener = new ActionListener<>() { @Override - public void onResponse(ChunkedInferenceServiceResults results) { + public void onResponse(List results) { if (results == null) { onBulkItemFailure.apply( bulkItemRequest, @@ -271,7 +272,13 @@ public void onResponse(ChunkedInferenceServiceResults results) { Map inferenceFieldResult = new LinkedHashMap<>(); inferenceFieldResult.putAll(new ModelSettings(inferenceProvider.model).asMap()); - inferenceFieldResult.put(INFERENCE_RESULTS, results.chunksAsMap()); + inferenceFieldResult.put( + INFERENCE_RESULTS, + results.stream() + .map(ChunkedInferenceServiceResults::chunksAsMap) + .flatMap(Collection::stream) + .collect(Collectors.toList()) + ); rootInferenceFieldMap.put(inferenceFieldName, inferenceFieldResult); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java index eb6dbc352d36d..97ad1b575caa9 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java @@ -17,9 +17,9 @@ import org.elasticsearch.core.Nullable; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.ServiceSettings; +import org.elasticsearch.inference.SimilarityMeasure; import org.elasticsearch.xcontent.ToXContentObject; import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xpack.inference.common.SimilarityMeasure; import java.io.IOException; import java.net.URI; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java index 92920e0b9224f..f176cf7580567 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java @@ -15,9 +15,9 @@ import org.elasticsearch.core.Nullable; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.ServiceSettings; +import org.elasticsearch.inference.SimilarityMeasure; import org.elasticsearch.xcontent.ToXContentObject; import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xpack.inference.common.SimilarityMeasure; import java.io.IOException; import java.net.URI; diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java index 229e45a024458..468e82d4f0866 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java @@ -15,9 +15,9 @@ import org.elasticsearch.core.Nullable; import org.elasticsearch.inference.ModelConfigurations; import org.elasticsearch.inference.ServiceSettings; +import org.elasticsearch.inference.SimilarityMeasure; import org.elasticsearch.xcontent.ToXContentObject; import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xpack.inference.common.SimilarityMeasure; import org.elasticsearch.xpack.inference.services.ServiceUtils; import org.elasticsearch.xpack.inference.services.openai.OpenAiParseContext;