Support configurable chunking in semantic_text fields

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -208,7 +208,10 @@ static TransportVersion def(int id) {
     public static final TransportVersion RESCORE_VECTOR_ALLOW_ZERO = def(9_039_0_00);
     public static final TransportVersion PROJECT_ID_IN_SNAPSHOT = def(9_040_0_00);
     public static final TransportVersion INDEX_STATS_AND_METADATA_INCLUDE_PEAK_WRITE_LOAD = def(9_041_0_00);
-    public static final TransportVersion SEMANTIC_TEXT_CHUNKING_CONFIG = def(9_042_00_0);
+    public static final TransportVersion REPOSITORIES_METADATA_AS_PROJECT_CUSTOM = def(9_042_0_00);
+    public static final TransportVersion BATCHED_QUERY_PHASE_VERSION = def(9_043_0_00);
+    public static final TransportVersion REMOTE_EXCEPTION = def(9_044_0_00);
+    public static final TransportVersion SEMANTIC_TEXT_CHUNKING_CONFIG = def(9_045_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestDenseInferenceServiceExtension.java

@@ -180,16 +180,14 @@ private List<ChunkedInference> makeChunkedResults(List<ChunkInferenceInput> inpu
             var results = new ArrayList<ChunkedInference>();
             for (ChunkInferenceInput input : inputs) {
                 List<ChunkedInput> chunkedInput = chunkInputs(input);
-                List<TextEmbeddingFloatResults.Chunk> chunks = new ArrayList<>(
-                    chunkedInput.stream()
-                        .map(
-                            c -> new TextEmbeddingFloatResults.Chunk(
-                                makeResults(List.of(c.input()), serviceSettings).embeddings().get(0),
-                                new ChunkedInference.TextOffset(c.startOffset(), c.endOffset())
-                            )
+                List<TextEmbeddingFloatResults.Chunk> chunks = chunkedInput.stream()
+                    .map(
+                        c -> new TextEmbeddingFloatResults.Chunk(
+                            makeResults(List.of(c.input()), serviceSettings).embeddings().get(0),
+                            new ChunkedInference.TextOffset(c.startOffset(), c.endOffset())
                         )
-                        .toList()
-                );
+                    )
+                    .toList();
                 ChunkedInferenceEmbedding chunkedInferenceEmbedding = new ChunkedInferenceEmbedding(chunks);
                 results.add(chunkedInferenceEmbedding);
             }

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestSparseInferenceServiceExtension.java

@@ -170,16 +170,15 @@ private SparseEmbeddingResults makeResults(List<String> input) {
         private List<ChunkedInference> makeChunkedResults(List<ChunkInferenceInput> inputs) {
             List<ChunkedInference> results = new ArrayList<>();
             for (ChunkInferenceInput chunkInferenceInput : inputs) {
-                String input = chunkInferenceInput.input();
                 List<ChunkedInput> chunkedInput = chunkInputs(chunkInferenceInput);
-                List<SparseEmbeddingResults.Chunk> chunks = new ArrayList<>(chunkedInput.stream().map(c -> {
+                List<SparseEmbeddingResults.Chunk> chunks = chunkedInput.stream().map(c -> {
                     var tokens = new ArrayList<WeightedToken>();
                     for (int i = 0; i < 5; i++) {
                         tokens.add(new WeightedToken("feature_" + i, generateEmbedding(c.input(), i)));
                     }
                     var embeddings = new SparseEmbeddingResults.Embedding(tokens, false);
                     return new SparseEmbeddingResults.Chunk(embeddings, new ChunkedInference.TextOffset(c.startOffset(), c.endOffset()));
-                }).toList());
+                }).toList();
                 ChunkedInferenceEmbedding chunkedInferenceEmbedding = new ChunkedInferenceEmbedding(chunks);
                 results.add(chunkedInferenceEmbedding);
             }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunker.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.inference.ChunkInferenceInput;
 import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.inference.ChunkingSettings;
+import org.elasticsearch.inference.ChunkingStrategy;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xpack.core.inference.results.ChunkedInferenceEmbedding;
@@ -22,6 +23,8 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReferenceArray;
 import java.util.function.Supplier;
@@ -48,11 +51,8 @@ public String chunkText() {
     }
 
     public record BatchRequest(List<Request> requests) {
-        public List<String> inputs() {
-            return requests.stream()
-                .map(r -> new ChunkInferenceInput(r.chunkText(), r.inputs().getFirst().chunkingSettings()))
-                .map(ChunkInferenceInput::input)
-                .collect(Collectors.toList());
+        public Supplier<List<String>> inputs() {
+            return () -> requests.stream().map(Request::chunkText).collect(Collectors.toList());
         }
     }
 
@@ -97,13 +97,20 @@ public EmbeddingRequestChunker(
             defaultChunkingSettings = DEFAULT_CHUNKING_SETTINGS;
         }
 
+        Map<ChunkingStrategy, Chunker> chunkers = inputs.stream()
+            .map(ChunkInferenceInput::chunkingSettings)
+            .filter(Objects::nonNull)
+            .map(ChunkingSettings::getChunkingStrategy)
+            .distinct()
+            .collect(Collectors.toMap(chunkingStrategy -> chunkingStrategy, ChunkerBuilder::fromChunkingStrategy));
+
         List<Request> allRequests = new ArrayList<>();
         for (int inputIndex = 0; inputIndex < inputs.size(); inputIndex++) {
             ChunkingSettings chunkingSettings = inputs.get(inputIndex).chunkingSettings();
             if (chunkingSettings == null) {
                 chunkingSettings = defaultChunkingSettings;
             }
-            Chunker chunker = ChunkerBuilder.fromChunkingStrategy(chunkingSettings.getChunkingStrategy());
+            Chunker chunker = chunkers.get(chunkingSettings.getChunkingStrategy());
             List<ChunkOffset> chunks = chunker.chunk(inputs.get(inputIndex).input(), chunkingSettings);
             int resultCount = Math.min(chunks.size(), MAX_CHUNKS);
             resultEmbeddings.add(new AtomicReferenceArray<>(resultCount));

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/EmbeddingsInput.java

@@ -14,6 +14,7 @@
 
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 public class EmbeddingsInput extends InferenceInputs {
@@ -26,14 +27,19 @@ public static EmbeddingsInput of(InferenceInputs inferenceInputs) {
         return (EmbeddingsInput) inferenceInputs;
     }
 
-    private final List<ChunkInferenceInput> input;
-
+    private final Supplier<List<ChunkInferenceInput>> listSupplier;
     private final InputType inputType;
 
     public EmbeddingsInput(List<ChunkInferenceInput> input, @Nullable InputType inputType) {
         this(input, inputType, false);
     }
 
+    public EmbeddingsInput(Supplier<List<ChunkInferenceInput>> inputSupplier, @Nullable InputType inputType) {
+        super(false);
+        this.listSupplier = Objects.requireNonNull(inputSupplier);
+        this.inputType = inputType;
+    }
+
     public EmbeddingsInput(List<String> input, @Nullable ChunkingSettings chunkingSettings, @Nullable InputType inputType) {
         this(input.stream().map(i -> new ChunkInferenceInput(i, chunkingSettings)).collect(Collectors.toList()), inputType, false);
     }
@@ -45,16 +51,16 @@ public EmbeddingsInput(List<ChunkInferenceInput> input, @Nullable InputType inpu
         this.inputType = inputType;
     }
 
-    public static EmbeddingsInput fromStrings(List<String> input, @Nullable InputType inputType) {
-        return new EmbeddingsInput(input, null, inputType);
+    public List<ChunkInferenceInput> getInputs() {
+        return this.listSupplier.get();
     }
 
-    public List<ChunkInferenceInput> getInputs() {
-        return this.input;
+    public static EmbeddingsInput fromStrings(List<String> input, @Nullable InputType inputType) {
+        return new EmbeddingsInput(input, null, inputType);
     }
 
     public List<String> getStringInputs() {
-        return this.input.stream().map(ChunkInferenceInput::input).collect(Collectors.toList());
+        return getInputs().stream().map(ChunkInferenceInput::input).collect(Collectors.toList());
     }
 
     public InputType getInputType() {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java

@@ -317,13 +317,14 @@ private void validateServiceSettings(MinimalServiceSettings settings) {
          * @return A mapper with the copied settings applied
          */
         private SemanticTextFieldMapper copySettings(SemanticTextFieldMapper mapper, MapperMergeContext mapperMergeContext) {
-            SemanticTextFieldMapper returnedMapper;
-            Builder builder = from(mapper);
+            SemanticTextFieldMapper returnedMapper = mapper;
             if (mapper.fieldType().getModelSettings() == null) {
+                Builder builder = from(mapper);
                 builder.setModelSettings(modelSettings.getValue());
+                returnedMapper = builder.build(mapperMergeContext.getMapperBuilderContext());
             }
-            builder.setChunkingSettings(mapper.fieldType().getChunkingSettings());
-            return builder.build(mapperMergeContext.getMapperBuilderContext());
+
+            return returnedMapper;
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchService.java

@@ -344,7 +344,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = alibabaCloudSearchModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockService.java

@@ -170,7 +170,7 @@ protected void doChunkedInfer(
 
             for (var request : batchedRequests) {
                 var action = baseAmazonBedrockModel.accept(actionCreator, taskSettings);
-                action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+                action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
             }
         } else {
             listener.onFailure(createInvalidModelException(model));

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureaistudio/AzureAiStudioService.java

@@ -141,7 +141,7 @@ protected void doChunkedInfer(
 
             for (var request : batchedRequests) {
                 var action = baseAzureAiStudioModel.accept(actionCreator, taskSettings);
-                action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+                action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
             }
         } else {
             listener.onFailure(createInvalidModelException(model));

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/azureopenai/AzureOpenAiService.java

@@ -294,7 +294,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = azureOpenAiModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java

@@ -307,7 +307,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = cohereModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googleaistudio/GoogleAiStudioService.java

@@ -364,7 +364,13 @@ protected void doChunkedInfer(
         ).batchRequestsWithListeners(listener);
 
         for (var request : batchedRequests) {
-            doInfer(model, EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), taskSettings, timeout, request.listener());
+            doInfer(
+                model,
+                EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType),
+                taskSettings,
+                timeout,
+                request.listener()
+            );
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java

@@ -250,7 +250,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = googleVertexAiModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java

@@ -135,7 +135,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = huggingFaceModel.accept(actionCreator);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ibmwatsonx/IbmWatsonxService.java

@@ -319,7 +319,7 @@ protected void doChunkedInfer(
         ).batchRequestsWithListeners(listener);
         for (var request : batchedRequests) {
             var action = ibmWatsonxModel.accept(getActionCreator(getSender(), getServiceComponents()), taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/jinaai/JinaAIService.java

@@ -288,7 +288,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = jinaaiModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/mistral/MistralService.java

@@ -123,7 +123,7 @@ protected void doChunkedInfer(
 
             for (var request : batchedRequests) {
                 var action = mistralEmbeddingsModel.accept(actionCreator, taskSettings);
-                action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+                action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
             }
         } else {
             listener.onFailure(createInvalidModelException(model));

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java

@@ -347,7 +347,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = openAiModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/VoyageAIService.java

@@ -308,7 +308,7 @@ protected void doChunkedInfer(
 
         for (var request : batchedRequests) {
             var action = voyageaiModel.accept(actionCreator, taskSettings);
-            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs(), inputType), timeout, request.listener());
+            action.execute(EmbeddingsInput.fromStrings(request.batch().inputs().get(), inputType), timeout, request.listener());
         }
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/chunking/EmbeddingRequestChunkerTests.java

@@ -98,7 +98,7 @@ public void testShortInputsAreSingleBatch() {
         ChunkInferenceInput input = new ChunkInferenceInput("one chunk");
         var batches = new EmbeddingRequestChunker<>(List.of(input), 100, 100, 10).batchRequestsWithListeners(testListener());
         assertThat(batches, hasSize(1));
-        assertThat(batches.get(0).batch().inputs(), contains(input.input()));
+        assertThat(batches.get(0).batch().inputs().get(), contains(input.input()));
     }
 
     public void testMultipleShortInputsAreSingleBatch() {
@@ -110,7 +110,7 @@ public void testMultipleShortInputsAreSingleBatch() {
         var batches = new EmbeddingRequestChunker<>(inputs, 100, 100, 10).batchRequestsWithListeners(testListener());
         assertThat(batches, hasSize(1));
         EmbeddingRequestChunker.BatchRequest batch = batches.getFirst().batch();
-        assertEquals(batch.inputs(), ChunkInferenceInput.inputs(inputs));
+        assertEquals(batch.inputs().get(), ChunkInferenceInput.inputs(inputs));
         for (int i = 0; i < inputs.size(); i++) {
             var request = batch.requests().get(i);
             assertThat(request.chunkText(), equalTo(inputs.get(i).input()));

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/SingleInputSenderExecutableActionTests.java

@@ -65,7 +65,7 @@ public void testOneInputIsValid() {
 
     public void testMoreThanOneInput() {
         var badInput = mock(EmbeddingsInput.class);
-        var input = List.of(new ChunkInferenceInput("one", null), new ChunkInferenceInput("two", null));
+        var input = List.of(new ChunkInferenceInput("one"), new ChunkInferenceInput("two"));
         when(badInput.getInputs()).thenReturn(input);
         when(badInput.isSingleInput()).thenReturn(false);
         var actualException = new AtomicReference<Exception>();

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilderTests.java

@@ -60,7 +60,6 @@
 import org.elasticsearch.xpack.core.ml.search.WeightedToken;
 import org.elasticsearch.xpack.inference.InferencePlugin;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextField;
-import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldTests;
 import org.junit.Before;
 import org.junit.BeforeClass;
 
@@ -370,12 +369,7 @@ private static SourceToParse buildSemanticTextFieldWithInferenceResults(
                 useLegacyFormat,
                 SEMANTIC_TEXT_FIELD,
                 null,
-                new SemanticTextField.InferenceResult(
-                    INFERENCE_ID,
-                    modelSettings,
-                    SemanticTextFieldTests.generateRandomChunkingSettings(),
-                    Map.of(SEMANTIC_TEXT_FIELD, List.of())
-                ),
+                new SemanticTextField.InferenceResult(INFERENCE_ID, modelSettings, null, Map.of(SEMANTIC_TEXT_FIELD, List.of())),
                 XContentType.JSON
             );
 

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/25_semantic_text_field_mapping_chunking.yml

@@ -274,13 +274,14 @@ setup:
             fields:
               inference_field:
                 type: "semantic"
-                number_of_fragments: 2
+                number_of_fragments: 3
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_4" }
-  - length: { hits.hits.0.highlight.inference_field: 2 }
+  - length: { hits.hits.0.highlight.inference_field: 3 }
   - match: { hits.hits.0.highlight.inference_field.0: "Elasticsearch is an open source, distributed, RESTful, search engine which" }
   - match: { hits.hits.0.highlight.inference_field.1: " which is built on top of Lucene internally and enjoys" }
+  - match: { hits.hits.0.highlight.inference_field.2: " enjoys all the features it provides." }
 
 ---
 "We respect multiple semantic_text fields with different chunking configurations":

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/25_semantic_text_field_mapping_chunking_bwc.yml

@@ -286,13 +286,14 @@ setup:
             fields:
               inference_field:
                 type: "semantic"
-                number_of_fragments: 2
+                number_of_fragments: 3
 
   - match: { hits.total.value: 1 }
   - match: { hits.hits.0._id: "doc_4" }
-  - length: { hits.hits.0.highlight.inference_field: 2 }
+  - length: { hits.hits.0.highlight.inference_field: 3 }
   - match: { hits.hits.0.highlight.inference_field.0: "Elasticsearch is an open source, distributed, RESTful, search engine which" }
   - match: { hits.hits.0.highlight.inference_field.1: " which is built on top of Lucene internally and enjoys" }
+  - match: { hits.hits.0.highlight.inference_field.2: " enjoys all the features it provides." }
 
 ---
 "We respect multiple semantic_text fields with different chunking configurations":