Fix merge from main

carlosdelest · Feb 14, 2024 · 8070f75 · 8070f75
1 parent e278205
commit 8070f75
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 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
@@ -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<ChunkedInferenceServiceResults> inferenceResultsListener = new ActionListener<>() {
+                    ActionListener<List<ChunkedInferenceServiceResults>> inferenceResultsListener = new ActionListener<>() {
                         @Override
-                        public void onResponse(ChunkedInferenceServiceResults results) {
+                        public void onResponse(List<ChunkedInferenceServiceResults> results) {
                             if (results == null) {
                                 onBulkItemFailure.apply(
                                     bulkItemRequest,
@@ -271,7 +272,13 @@ public void onResponse(ChunkedInferenceServiceResults results) {
 
                             Map<String, Object> 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/...rc/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceSettings.java b/...rc/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/...va/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceSettings.java b/...va/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/...ticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsServiceSettings.java b/...ticsearch/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;