Add hardcoded max token length

Author: kderusso

server/src/main/java/org/elasticsearch/action/search/RankFeaturePhase.java

@@ -173,7 +173,8 @@ public void onFailure(Exception e) {
                     queryResult.getContextId(),
                     queryResult.getShardSearchRequest(),
                     entry,
-                    rankFeaturePhaseRankCoordinatorContext.snippets()
+                    rankFeaturePhaseRankCoordinatorContext.snippets(),
+                    rankFeaturePhaseRankCoordinatorContext.tokenSizeLimit()
                 ),
                 context.getTask(),
                 listener

server/src/main/java/org/elasticsearch/search/rank/context/RankFeaturePhaseRankCoordinatorContext.java

@@ -55,6 +55,8 @@ public RerankSnippetInput snippets() {
         return snippets;
     }
 
+    public abstract Integer tokenSizeLimit();
+
     /**
      * Computes the updated scores for a list of features (i.e. document-based data). We also pass along an ActionListener
      * that should be called with the new scores, and will continue execution to the next phase

server/src/main/java/org/elasticsearch/search/rank/feature/RankFeatureShardPhase.java

@@ -66,7 +66,7 @@ public static void prepareForFetch(SearchContext searchContext, RankFeatureShard
                     if (snippets.numSnippets() != null) {
                         highlightBuilder.numOfFragments(snippets.numSnippets());
                     }
-                    highlightBuilder.fragmentSize(20); // TODO use model limit
+                    highlightBuilder.fragmentSize(request.getTokenSizeLimit());
                     SearchHighlightContext searchHighlightContext = highlightBuilder.build(searchContext.getSearchExecutionContext());
                     searchContext.highlight(searchHighlightContext);
                 }

server/src/main/java/org/elasticsearch/search/rank/feature/RankFeatureShardRequest.java

@@ -41,19 +41,22 @@ public class RankFeatureShardRequest extends AbstractTransportRequest implements
     private final int[] docIds;
 
     private final RerankSnippetInput snippets;
+    private final int tokenSizeLimit;
 
     public RankFeatureShardRequest(
         OriginalIndices originalIndices,
         ShardSearchContextId contextId,
         ShardSearchRequest shardSearchRequest,
         List<Integer> docIds,
-        @Nullable RerankSnippetInput snippets
+        @Nullable RerankSnippetInput snippets,
+        int tokenSizeLimit
     ) {
         this.originalIndices = originalIndices;
         this.shardSearchRequest = shardSearchRequest;
         this.docIds = docIds.stream().flatMapToInt(IntStream::of).toArray();
         this.contextId = contextId;
         this.snippets = snippets;
+        this.tokenSizeLimit = tokenSizeLimit;
     }
 
     public RankFeatureShardRequest(StreamInput in) throws IOException {
@@ -64,8 +67,10 @@ public RankFeatureShardRequest(StreamInput in) throws IOException {
         contextId = in.readOptionalWriteable(ShardSearchContextId::new);
         if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             snippets = in.readOptionalWriteable(RerankSnippetInput::new);
+            this.tokenSizeLimit = in.readVInt();
         } else {
             snippets = null;
+            this.tokenSizeLimit = 0;
         }
     }
 
@@ -78,6 +83,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalWriteable(contextId);
         if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             out.writeOptionalWriteable(snippets);
+            out.writeVInt(tokenSizeLimit);
         }
     }
 
@@ -113,6 +119,10 @@ public RerankSnippetInput snippets() {
         return snippets;
     }
 
+    public int getTokenSizeLimit() {
+        return tokenSizeLimit;
+    }
+
     @Override
     public SearchShardTask createTask(long id, String type, String action, TaskId parentTaskId, Map<String, String> headers) {
         return new SearchShardTask(id, type, action, getDescription(), parentTaskId, headers);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/random/RandomRankFeaturePhaseRankCoordinatorContext.java

@@ -26,6 +26,11 @@ public RandomRankFeaturePhaseRankCoordinatorContext(int size, int from, int rank
         this.seed = seed;
     }
 
+    @Override
+    public Integer tokenSizeLimit() {
+        return 0;
+    }
+
     @Override
     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
         // Generate random scores seeded by doc

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

@@ -28,13 +28,18 @@
 
 import static org.elasticsearch.xpack.core.ClientHelper.INFERENCE_ORIGIN;
 import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin;
+import static org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService.DEFAULT_RERANK_ID;
+import static org.elasticsearch.xpack.inference.services.elasticsearch.ElasticsearchInternalService.RERANKER_ID;
 
 /**
  * A {@code RankFeaturePhaseRankCoordinatorContext} that performs a rerank inference call to determine relevance scores for documents within
  * the provided rank window.
  */
 public class TextSimilarityRankFeaturePhaseRankCoordinatorContext extends RankFeaturePhaseRankCoordinatorContext {
 
+    private static final int RERANK_TOKEN_SIZE_LIMIT = 512;
+    private static final int DEFAULT_TOKEN_SIZE_LIMIT = 4096;
+
     protected final Client client;
     protected final String inferenceId;
     protected final String inferenceText;
@@ -58,6 +63,20 @@ public TextSimilarityRankFeaturePhaseRankCoordinatorContext(
         this.minScore = minScore;
     }
 
+    /**
+     * @return The token size limit to apply to this rerank context.
+     * This is not yet available so we are hardcoding it for now.
+     * See: https://github.com/elastic/ml-team/issues/1622
+     */
+    @Override
+    public Integer tokenSizeLimit() {
+        if (inferenceId.equals(DEFAULT_RERANK_ID) || inferenceId.equals(RERANKER_ID)) {
+            return RERANK_TOKEN_SIZE_LIMIT;
+        }
+
+        return DEFAULT_TOKEN_SIZE_LIMIT;
+    }
+
     @Override
     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {