Add RerankRequestChunker

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/action/TransportInferenceAction.java

@@ -73,6 +73,7 @@ protected void doInference(
         InferenceService service,
         ActionListener<InferenceServiceResults> listener
     ) {
+
         service.infer(
             model,
             request.getQuery(),

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

@@ -17,6 +17,9 @@ public class ChunkingSettingsBuilder {
     public static final SentenceBoundaryChunkingSettings DEFAULT_SETTINGS = new SentenceBoundaryChunkingSettings(250, 1);
     // Old settings used for backward compatibility for endpoints created before 8.16 when default was changed
     public static final WordBoundaryChunkingSettings OLD_DEFAULT_SETTINGS = new WordBoundaryChunkingSettings(250, 100);
+    public static final int ELASTIC_RERANKER_TOKEN_LIMIT = 512;
+    public static final int ELASTIC_RERANKER_EXTRA_TOKEN_COUNT = 3;
+    public static final float WORDS_PER_TOKEN = 0.75f;
 
     public static ChunkingSettings fromMap(Map<String, Object> settings) {
         return fromMap(settings, true);
@@ -51,4 +54,17 @@ public static ChunkingSettings fromMap(Map<String, Object> settings, boolean ret
             case RECURSIVE -> RecursiveChunkingSettings.fromMap(new HashMap<>(settings));
         };
     }
+
+    public static ChunkingSettings buildChunkingSettingsForElasticRerank(int queryWordCount) {
+        var queryTokenCount = Math.ceil(queryWordCount / WORDS_PER_TOKEN);
+        var chunkSizeTokenCountWithFullQuery = (ELASTIC_RERANKER_TOKEN_LIMIT - ELASTIC_RERANKER_EXTRA_TOKEN_COUNT - queryTokenCount);
+
+        var maxChunkSizeTokenCount = Math.floor((float) ELASTIC_RERANKER_TOKEN_LIMIT / 2);
+        if (chunkSizeTokenCountWithFullQuery > maxChunkSizeTokenCount) {
+            maxChunkSizeTokenCount = chunkSizeTokenCountWithFullQuery;
+        }
+
+        var maxChunkSizeWordCount = (int) (maxChunkSizeTokenCount * WORDS_PER_TOKEN);
+        return new SentenceBoundaryChunkingSettings(maxChunkSizeWordCount, 1);
+    }
 }

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

@@ -0,0 +1,104 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.chunking;
+
+import com.ibm.icu.text.BreakIterator;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.ChunkingSettings;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+public class RerankRequestChunker {
+    private final List<String> inputs;
+    private final List<RerankChunks> rerankChunks;
+
+    public RerankRequestChunker(String query, List<String> inputs, Integer maxChunksPerDoc) {
+        this.inputs = inputs;
+        this.rerankChunks = chunk(inputs, buildChunkingSettingsForElasticRerank(query), maxChunksPerDoc);
+    }
+
+    private List<RerankChunks> chunk(List<String> inputs, ChunkingSettings chunkingSettings, Integer maxChunksPerDoc) {
+        var chunker = ChunkerBuilder.fromChunkingStrategy(chunkingSettings.getChunkingStrategy());
+        var chunks = new ArrayList<RerankChunks>();
+        for (int i = 0; i < inputs.size(); i++) {
+            var chunksForInput = chunker.chunk(inputs.get(i), chunkingSettings);
+            if (maxChunksPerDoc != null && chunksForInput.size() > maxChunksPerDoc) {
+                var limitedChunks = chunksForInput.subList(0, maxChunksPerDoc - 1);
+                var lastChunk = limitedChunks.getLast();
+                limitedChunks.add(new Chunker.ChunkOffset(lastChunk.end(), inputs.get(i).length()));
+                chunksForInput = limitedChunks;
+            }
+
+            for (var chunk : chunksForInput) {
+                chunks.add(new RerankChunks(i, inputs.get(i).substring(chunk.start(), chunk.end())));
+            }
+        }
+        return chunks;
+    }
+
+    public List<String> getChunkedInputs() {
+        List<String> chunkedInputs = new ArrayList<>();
+        for (RerankChunks chunk : rerankChunks) {
+            chunkedInputs.add(chunk.chunkString());
+        }
+
+        // TODO: Score the inputs here and only return the top N chunks for each document
+        return chunkedInputs;
+    }
+
+    public ActionListener<InferenceServiceResults> parseChunkedRerankResultsListener(ActionListener<InferenceServiceResults> listener) {
+        return ActionListener.wrap(results -> {
+            if (results instanceof RankedDocsResults rankedDocsResults) {
+                listener.onResponse(parseRankedDocResultsForChunks(rankedDocsResults));
+
+            } else {
+                listener.onFailure(new IllegalArgumentException("Expected RankedDocsResults but got: " + results.getClass()));
+            }
+
+        }, listener::onFailure);
+    }
+
+    // TODO: Can we assume the rankeddocsresults are always sorted by relevance score?
+    // TODO: Should we short circuit if no chunking was done?
+    private RankedDocsResults parseRankedDocResultsForChunks(RankedDocsResults rankedDocsResults) {
+        List<RankedDocsResults.RankedDoc> updatedRankedDocs = new ArrayList<>();
+        Set<Integer> docIndicesSeen = new HashSet<>();
+        for (RankedDocsResults.RankedDoc rankedDoc : rankedDocsResults.getRankedDocs()) {
+            int chunkIndex = rankedDoc.index();
+            int docIndex = rerankChunks.get(chunkIndex).docIndex();
+
+            if (docIndicesSeen.contains(docIndex) == false) {
+                // Create a ranked doc with the full input string and the index for the document instead of the chunk
+                RankedDocsResults.RankedDoc updatedRankedDoc = new RankedDocsResults.RankedDoc(
+                    docIndex,
+                    rankedDoc.relevanceScore(),
+                    inputs.get(docIndex)
+                );
+                updatedRankedDocs.add(updatedRankedDoc);
+                docIndicesSeen.add(docIndex);
+            }
+        }
+
+        return new RankedDocsResults(updatedRankedDocs);
+    }
+
+    public record RerankChunks(int docIndex, String chunkString) {};
+
+    private ChunkingSettings buildChunkingSettingsForElasticRerank(String query) {
+        var wordIterator = BreakIterator.getWordInstance();
+        wordIterator.setText(query);
+        var queryWordCount = ChunkerUtils.countWords(0, query.length(), wordIterator);
+        return ChunkingSettingsBuilder.buildChunkingSettingsForElasticRerank(queryWordCount);
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankFeaturePhaseRankCoordinatorContext.java

@@ -178,6 +178,7 @@ protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs, boolean rer
     }
 
     protected InferenceAction.Request generateRequest(List<String> docFeatures) {
+        // TODO: Try running the RerankRequestChunker here.
         return new InferenceAction.Request(
             TaskType.RERANK,
             inferenceId,

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

@@ -9,23 +9,49 @@
 
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.core.ml.inference.assignment.AdaptiveAllocationsSettings;
 
 import java.io.IOException;
+import java.util.EnumSet;
+import java.util.Locale;
 import java.util.Map;
 
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalEnum;
+import static org.elasticsearch.xpack.inference.services.ServiceUtils.extractOptionalPositiveInteger;
 import static org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService.RERANKER_ID;
 
 public class ElasticRerankerServiceSettings extends ElasticsearchInternalServiceSettings {
 
     public static final String NAME = "elastic_reranker_service_settings";
 
+    private static final String LONG_DOCUMENT_HANDLING_STRATEGY = "long_document_handling_strategy";
+    private static final String MAX_CHUNKS_PER_DOC = "max_chunks_per_doc";
+
+    private final LongDocumentHandlingStrategy longDocumentHandlingStrategy;
+    private final Integer maxChunksPerDoc;
+
     public static ElasticRerankerServiceSettings defaultEndpointSettings() {
         return new ElasticRerankerServiceSettings(null, 1, RERANKER_ID, new AdaptiveAllocationsSettings(Boolean.TRUE, 0, 32));
     }
 
     public ElasticRerankerServiceSettings(ElasticsearchInternalServiceSettings other) {
         super(other);
+        this.longDocumentHandlingStrategy = null;
+        this.maxChunksPerDoc = null;
+    }
+
+    public ElasticRerankerServiceSettings(
+        ElasticsearchInternalServiceSettings other,
+        LongDocumentHandlingStrategy longDocumentHandlingStrategy,
+        Integer maxChunksPerDoc
+    ) {
+        super(other);
+        this.longDocumentHandlingStrategy = longDocumentHandlingStrategy;
+        this.maxChunksPerDoc = maxChunksPerDoc;
+
     }
 
     private ElasticRerankerServiceSettings(
@@ -35,10 +61,15 @@ private ElasticRerankerServiceSettings(
         AdaptiveAllocationsSettings adaptiveAllocationsSettings
     ) {
         super(numAllocations, numThreads, modelId, adaptiveAllocationsSettings, null);
+        this.longDocumentHandlingStrategy = null;
+        this.maxChunksPerDoc = null;
     }
 
     public ElasticRerankerServiceSettings(StreamInput in) throws IOException {
         super(in);
+        // TODO: Add transport version here
+        this.longDocumentHandlingStrategy = in.readOptionalEnum(LongDocumentHandlingStrategy.class);
+        this.maxChunksPerDoc = in.readOptionalInt();
     }
 
     /**
@@ -48,21 +79,89 @@ public ElasticRerankerServiceSettings(StreamInput in) throws IOException {
      * {@link ValidationException} is thrown.
      *
      * @param map Source map containing the config
-     * @return The builder
+     * @return Parsed and validated service settings
      */
-    public static Builder fromRequestMap(Map<String, Object> map) {
+    public static ElasticRerankerServiceSettings fromMap(Map<String, Object> map) {
         ValidationException validationException = new ValidationException();
         var baseSettings = ElasticsearchInternalServiceSettings.fromMap(map, validationException);
 
+        LongDocumentHandlingStrategy longDocumentHandlingStrategy = extractOptionalEnum(
+            map,
+            LONG_DOCUMENT_HANDLING_STRATEGY,
+            ModelConfigurations.SERVICE_SETTINGS,
+            LongDocumentHandlingStrategy::fromString,
+            EnumSet.allOf(LongDocumentHandlingStrategy.class),
+            validationException
+        );
+
+        Integer maxChunksPerDoc = extractOptionalPositiveInteger(
+            map,
+            MAX_CHUNKS_PER_DOC,
+            ModelConfigurations.SERVICE_SETTINGS,
+            validationException
+        );
+
+        if (maxChunksPerDoc != null
+            && (longDocumentHandlingStrategy == null || longDocumentHandlingStrategy == LongDocumentHandlingStrategy.TRUNCATE)) {
+            validationException.addValidationError(
+                "The [" + MAX_CHUNKS_PER_DOC + "] setting requires [" + LONG_DOCUMENT_HANDLING_STRATEGY + "] to be set to [chunk]"
+            );
+        }
+
         if (validationException.validationErrors().isEmpty() == false) {
             throw validationException;
         }
 
-        return baseSettings;
+        return new ElasticRerankerServiceSettings(baseSettings.build(), longDocumentHandlingStrategy, maxChunksPerDoc);
+    }
+
+    public LongDocumentHandlingStrategy getLongDocumentHandlingStrategy() {
+        return longDocumentHandlingStrategy;
+    }
+
+    public Integer getMaxChunksPerDoc() {
+        return maxChunksPerDoc;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        super.writeTo(out);
+        // TODO: Add transport version here
+        out.writeOptionalEnum(longDocumentHandlingStrategy);
+        out.writeOptionalInt(maxChunksPerDoc);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+        addInternalSettingsToXContent(builder, params);
+        if (longDocumentHandlingStrategy != null) {
+            builder.field(LONG_DOCUMENT_HANDLING_STRATEGY, longDocumentHandlingStrategy.strategyName);
+        }
+        if (maxChunksPerDoc != null) {
+            builder.field(MAX_CHUNKS_PER_DOC, maxChunksPerDoc);
+        }
+        builder.endObject();
+        return builder;
     }
 
     @Override
     public String getWriteableName() {
         return ElasticRerankerServiceSettings.NAME;
     }
+
+    public enum LongDocumentHandlingStrategy {
+        CHUNK("chunk"),
+        TRUNCATE("truncate");
+
+        public final String strategyName;
+
+        LongDocumentHandlingStrategy(String strategyName) {
+            this.strategyName = strategyName;
+        }
+
+        public static LongDocumentHandlingStrategy fromString(String name) {
+            return valueOf(name.trim().toUpperCase(Locale.ROOT));
+        }
+    }
 }

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

@@ -58,6 +58,7 @@
 import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TextSimilarityConfigUpdate;
 import org.elasticsearch.xpack.inference.chunking.ChunkingSettingsBuilder;
 import org.elasticsearch.xpack.inference.chunking.EmbeddingRequestChunker;
+import org.elasticsearch.xpack.inference.chunking.RerankRequestChunker;
 import org.elasticsearch.xpack.inference.services.ConfigurationParseContext;
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
 
@@ -349,19 +350,13 @@ private void rerankerCase(
         ActionListener<Model> modelListener
     ) {
 
-        var esServiceSettingsBuilder = ElasticsearchInternalServiceSettings.fromRequestMap(serviceSettingsMap);
+        var serviceSettings = ElasticRerankerServiceSettings.fromMap(serviceSettingsMap);
 
         throwIfNotEmptyMap(config, name());
         throwIfNotEmptyMap(serviceSettingsMap, name());
 
         modelListener.onResponse(
-            new ElasticRerankerModel(
-                inferenceEntityId,
-                taskType,
-                NAME,
-                new ElasticRerankerServiceSettings(esServiceSettingsBuilder.build()),
-                RerankTaskSettings.fromMap(taskSettingsMap)
-            )
+            new ElasticRerankerModel(inferenceEntityId, taskType, NAME, serviceSettings, RerankTaskSettings.fromMap(taskSettingsMap))
         );
     }
 
@@ -535,7 +530,7 @@ public Model parsePersistedConfig(String inferenceEntityId, TaskType taskType, M
                 inferenceEntityId,
                 taskType,
                 NAME,
-                new ElasticRerankerServiceSettings(ElasticsearchInternalServiceSettings.fromPersistedMap(serviceSettingsMap)),
+                ElasticRerankerServiceSettings.fromMap(serviceSettingsMap),
                 RerankTaskSettings.fromMap(taskSettingsMap)
             );
         } else {
@@ -688,7 +683,25 @@ public void inferRerank(
         Map<String, Object> requestTaskSettings,
         ActionListener<InferenceServiceResults> listener
     ) {
-        var request = buildInferenceRequest(model.mlNodeDeploymentId(), new TextSimilarityConfigUpdate(query), inputs, inputType, timeout);
+        var chunkedInputs = inputs;
+        var resultsListener = listener;
+        if (model instanceof ElasticRerankerModel elasticRerankerModel) {
+            var serviceSettings = elasticRerankerModel.getServiceSettings();
+            var longDocumentHandlingStrategy = serviceSettings.getLongDocumentHandlingStrategy();
+            if (longDocumentHandlingStrategy == ElasticRerankerServiceSettings.LongDocumentHandlingStrategy.CHUNK) {
+                var rerankChunker = new RerankRequestChunker(query, inputs, serviceSettings.getMaxChunksPerDoc());
+                chunkedInputs = rerankChunker.getChunkedInputs();
+                resultsListener = rerankChunker.parseChunkedRerankResultsListener(listener);
+            }
+
+        }
+        var request = buildInferenceRequest(
+            model.mlNodeDeploymentId(),
+            new TextSimilarityConfigUpdate(query),
+            chunkedInputs,
+            inputType,
+            timeout
+        );
 
         var returnDocs = Boolean.TRUE;
         if (returnDocuments != null) {
@@ -698,9 +711,9 @@ public void inferRerank(
             returnDocs = RerankTaskSettings.of(modelSettings, requestSettings).returnDocuments();
         }
 
-        Function<Integer, String> inputSupplier = returnDocs == Boolean.TRUE ? inputs::get : i -> null;
+        Function<Integer, String> inputSupplier = returnDocs == Boolean.TRUE ? chunkedInputs::get : i -> null;
 
-        ActionListener<InferModelAction.Response> mlResultsListener = listener.delegateFailureAndWrap(
+        ActionListener<InferModelAction.Response> mlResultsListener = resultsListener.delegateFailureAndWrap(
             (l, inferenceResult) -> l.onResponse(
                 textSimilarityResultsToRankedDocs(inferenceResult.getInferenceResults(), inputSupplier, topN)
             )