Add Hugging Face Rerank support

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

@@ -176,7 +176,9 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_REPORT_SHARD_PARTITIONING_8_19 = def(8_841_0_29);
     public static final TransportVersion ESQL_DRIVER_TASK_DESCRIPTION_8_19 = def(8_841_0_30);
     public static final TransportVersion ML_INFERENCE_HUGGING_FACE_CHAT_COMPLETION_ADDED_8_19 = def(8_841_0_31);
-    public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED_8_19 = def(8_841_0_32);
+    public static final TransportVersion V_8_19_FIELD_CAPS_ADD_CLUSTER_ALIAS = def(8_841_0_32);
+    public static final TransportVersion ESQL_HASH_OPERATOR_STATUS_OUTPUT_TIME_8_19 = def(8_841_0_34);
+    public static final TransportVersion ML_INFERENCE_HUGGING_FACE_RERANK_ADDED_8_19 = def(8_841_0_35);
     public static final TransportVersion V_9_0_0 = def(9_000_0_09);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_1 = def(9_000_0_10);
     public static final TransportVersion INITIAL_ELASTICSEARCH_9_0_2 = def(9_000_0_11);

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/results/RankedDocsResultsTests.java

@@ -82,4 +82,20 @@ private List<RankedDocsResults.RankedDoc> rankedDocsNullStringToEmpty(List<Ranke
     protected RankedDocsResults doParseInstance(XContentParser parser) throws IOException {
         return RankedDocsResults.createParser(true).apply(parser, null);
     }
+
+    public record RerankExpectation(Map<String, Object> rankedDocFields) {}
+
+    public static Map<String, Object> buildExpectationRerank(List<RerankExpectation> rerank) {
+        return Map.of(
+            RankedDocsResults.RERANK,
+            rerank.stream()
+                .map(
+                    rerankExpectation -> Map.of(
+                        RankedDocsResults.RankedDoc.NAME,
+                        rerankExpectation.rankedDocFields
+                    )
+                )
+                .toList()
+        );
+    }
 }

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

@@ -26,7 +26,6 @@
 import org.elasticsearch.xpack.inference.services.ServiceUtils;
 import org.elasticsearch.xpack.inference.services.huggingface.action.HuggingFaceActionCreator;
 
-import java.util.Collections;
 import java.util.Map;
 
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.createInvalidModelException;
@@ -58,11 +57,7 @@ public void parseRequestConfig(
     ) {
         try {
             Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
-            Map<String, Object> taskSettingsMap = Collections.emptyMap();
-
-            if (TaskType.RERANK.equals(taskType)) {
-                taskSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.TASK_SETTINGS);
-            }
+            Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
 
             ChunkingSettings chunkingSettings = null;
             if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
@@ -102,12 +97,8 @@ public HuggingFaceModel parsePersistedConfigWithSecrets(
         Map<String, Object> secrets
     ) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
+        Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
         Map<String, Object> secretSettingsMap = removeFromMapOrThrowIfNull(secrets, ModelSecrets.SECRET_SETTINGS);
-        Map<String, Object> taskSettingsMap = Collections.emptyMap();
-
-        if (TaskType.RERANK.equals(taskType)) {
-            taskSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.TASK_SETTINGS);
-        }
 
         ChunkingSettings chunkingSettings = null;
         if (TaskType.TEXT_EMBEDDING.equals(taskType)) {
@@ -131,11 +122,7 @@ public HuggingFaceModel parsePersistedConfigWithSecrets(
     @Override
     public HuggingFaceModel parsePersistedConfig(String inferenceEntityId, TaskType taskType, Map<String, Object> config) {
         Map<String, Object> serviceSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.SERVICE_SETTINGS);
-        Map<String, Object> taskSettingsMap = Collections.emptyMap();
-
-        if (TaskType.RERANK.equals(taskType)) {
-            taskSettingsMap = removeFromMapOrThrowIfNull(config, ModelConfigurations.TASK_SETTINGS);
-        }
+        Map<String, Object> taskSettingsMap = removeFromMapOrDefaultEmpty(config, ModelConfigurations.TASK_SETTINGS);
 
         ChunkingSettings chunkingSettings = null;
         if (TaskType.TEXT_EMBEDDING.equals(taskType)) {

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

@@ -76,43 +76,43 @@ public HuggingFaceService(HttpRequestSender.Factory factory, ServiceComponents s
     }
 
     @Override
-    protected HuggingFaceModel createModel(HuggingFaceModelParameters input) {
-        return switch (input.taskType()) {
+    protected HuggingFaceModel createModel(HuggingFaceModelParameters params) {
+        return switch (params.taskType()) {
             case RERANK -> new HuggingFaceRerankModel(
-                input.inferenceEntityId(),
-                input.taskType(),
+                params.inferenceEntityId(),
+                params.taskType(),
                 NAME,
-                input.serviceSettings(),
-                input.taskSettings(),
-                input.secretSettings(),
-                input.context()
+                params.serviceSettings(),
+                params.taskSettings(),
+                params.secretSettings(),
+                params.context()
             );
             case TEXT_EMBEDDING -> new HuggingFaceEmbeddingsModel(
-                input.inferenceEntityId(),
-                input.taskType(),
+                params.inferenceEntityId(),
+                params.taskType(),
                 NAME,
-                input.serviceSettings(),
-                input.chunkingSettings(),
-                input.secretSettings(),
-                input.context()
+                params.serviceSettings(),
+                params.chunkingSettings(),
+                params.secretSettings(),
+                params.context()
             );
             case SPARSE_EMBEDDING -> new HuggingFaceElserModel(
-                input.inferenceEntityId(),
-                input.taskType(),
+                params.inferenceEntityId(),
+                params.taskType(),
                 NAME,
-                input.serviceSettings(),
-                input.secretSettings(),
-                input.context()
+                params.serviceSettings(),
+                params.secretSettings(),
+                params.context()
             );
             case CHAT_COMPLETION, COMPLETION -> new HuggingFaceChatCompletionModel(
-                input.inferenceEntityId(),
-                input.taskType(),
+                params.inferenceEntityId(),
+                params.taskType(),
                 NAME,
-                input.serviceSettings(),
-                input.secretSettings(),
-                input.context()
+                params.serviceSettings(),
+                params.secretSettings(),
+                params.context()
             );
-            default -> throw new ElasticsearchStatusException(input.failureMessage(), RestStatus.BAD_REQUEST);
+            default -> throw new ElasticsearchStatusException(params.failureMessage(), RestStatus.BAD_REQUEST);
         };
     }
 

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

@@ -51,9 +51,12 @@ public class HuggingFaceActionCreator implements HuggingFaceActionVisitor {
         OpenAiChatCompletionResponseEntity::fromResponse
     );
     private static final ResponseHandler RERANK_HANDLER = new HuggingFaceResponseHandler("hugging face rerank", (request, response) -> {
-        var errorMessage = format(INVALID_REQUEST_TYPE_MESSAGE, "RERANK", request != null ? request.getClass().getName() : "null");
-
         if ((request instanceof HuggingFaceRerankRequest) == false) {
+            var errorMessage = format(
+                INVALID_REQUEST_TYPE_MESSAGE,
+                "RERANK",
+                request != null ? request.getClass().getSimpleName() : "null"
+            );
             throw new IllegalArgumentException(errorMessage);
         }
         return HuggingFaceRerankResponseEntity.fromResponse((HuggingFaceRerankRequest) request, response);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/request/rerank/HuggingFaceRerankRequest.java

@@ -62,20 +62,19 @@ public HttpRequest createHttpRequest() {
                     input,
                     returnDocuments,
                     topN != null ? topN : model.getTaskSettings().getTopNDocumentsOnly(),
-                    model.getTaskSettings(),
-                    model.getServiceSettings().modelId()
+                    model.getTaskSettings()
                 )
             ).getBytes(StandardCharsets.UTF_8)
         );
         httpPost.setEntity(byteEntity);
-        httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
+        httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaTypeWithoutParameters());
 
         decorateWithAuth(httpPost);
 
         return new HttpRequest(httpPost, getInferenceEntityId());
     }
 
-    public void decorateWithAuth(HttpPost httpPost) {
+    void decorateWithAuth(HttpPost httpPost) {
         httpPost.setHeader(createAuthBearerHeader(model.apiKey()));
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/request/rerank/HuggingFaceRerankRequestEntity.java

@@ -17,7 +17,6 @@
 import java.util.Objects;
 
 public record HuggingFaceRerankRequestEntity(
-    String model,
     String query,
     List<String> documents,
     @Nullable Boolean returnDocuments,
@@ -34,17 +33,6 @@ public record HuggingFaceRerankRequestEntity(
         Objects.requireNonNull(taskSettings);
     }
 
-    public HuggingFaceRerankRequestEntity(
-        String query,
-        List<String> input,
-        @Nullable Boolean returnDocuments,
-        @Nullable Integer topN,
-        HuggingFaceRerankTaskSettings taskSettings,
-        String model
-    ) {
-        this(model, query, input, returnDocuments, topN, taskSettings);
-    }
-
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();

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

@@ -44,7 +44,7 @@ public HuggingFaceRerankModel(
     }
 
     // Should only be used directly for testing
-    public HuggingFaceRerankModel(
+    HuggingFaceRerankModel(
         String inferenceEntityId,
         TaskType taskType,
         String service,

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

@@ -23,16 +23,11 @@
 import java.util.Comparator;
 import java.util.List;
 
-import static org.elasticsearch.common.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.elasticsearch.common.xcontent.XContentParserUtils.parseList;
-import static org.elasticsearch.core.Strings.format;
 import static org.elasticsearch.xpack.inference.external.response.XContentUtils.moveToFirstToken;
 
 public class HuggingFaceRerankResponseEntity {
 
-    private static final String FAILED_TO_FIND_FIELD_TEMPLATE = "Failed to find required field [%s] in Hugging Face rerank response";
-    private static final String INVALID_ID_FIELD_FORMAT_TEMPLATE = "Expected numeric value for record ID field in Hugging Face rerank ";
-
     /**
      * Parses the Hugging Face rerank response.
 
@@ -41,10 +36,9 @@ public class HuggingFaceRerankResponseEntity {
      * <pre>
      *     <code>
      *         {
-     *              "input": ["luke", "like", "leia", "chewy","r2d2", "star", "wars"],
+     *              "texts": ["luke", "leia"],
      *              "query": "star wars main character",
-     *              "return_documents": true,
-     *              "top_n": 1
+     *              "return_text": true
      *          }
      *     </code>
      * </pre>
@@ -53,15 +47,18 @@ public class HuggingFaceRerankResponseEntity {
 
      * <pre>
      *     <code>
-     *         {
-     *              "rerank": [
-     *                  {
-     *                       "index": 5,
-     *                       "relevance_score": -0.06920313,
-     *                       "text": "star"
-     *                   }
-     *               ]
-     *          }
+     *         [
+     *              {
+     *                   "index": 0,
+     *                   "score": -0.07996220886707306,
+     *                   "text": "luke"
+     *               },
+     *               {
+     *                  "index": 1,
+     *                  "score": -0.08393221348524094,
+     *                  "text": "leia"
+     *              }
+     *         ]
      *     </code>
      * </pre>
      */
@@ -71,10 +68,6 @@ public static RankedDocsResults fromResponse(HuggingFaceRerankRequest request, H
 
         try (XContentParser jsonParser = XContentFactory.xContent(XContentType.JSON).createParser(parserConfig, response.body())) {
             moveToFirstToken(jsonParser);
-
-            XContentParser.Token token = jsonParser.currentToken();
-            ensureExpectedToken(XContentParser.Token.START_ARRAY, token, jsonParser);
-
             var rankedDocs = doParse(jsonParser);
             var rankedDocsByRelevanceStream = rankedDocs.stream()
                 .sorted(Comparator.comparingDouble(RankedDocsResults.RankedDoc::relevanceScore).reversed());
@@ -88,24 +81,11 @@ public static RankedDocsResults fromResponse(HuggingFaceRerankRequest request, H
     private static List<RankedDocsResults.RankedDoc> doParse(XContentParser parser) throws IOException {
         return parseList(parser, (listParser, index) -> {
             var parsedRankedDoc = HuggingFaceRerankResponseEntity.RankedDocEntry.parse(parser);
-
-            if (parsedRankedDoc.id == null) {
-                throw new IllegalStateException(format(FAILED_TO_FIND_FIELD_TEMPLATE, RankedDocEntry.ID.getPreferredName()));
-            }
-
-            if (parsedRankedDoc.score == null) {
-                throw new IllegalStateException(format(FAILED_TO_FIND_FIELD_TEMPLATE, RankedDocEntry.SCORE.getPreferredName()));
-            }
-
-            try {
-                return new RankedDocsResults.RankedDoc(parsedRankedDoc.id, parsedRankedDoc.score, parsedRankedDoc.text);
-            } catch (NumberFormatException e) {
-                throw new IllegalStateException(format(INVALID_ID_FIELD_FORMAT_TEMPLATE, parsedRankedDoc.id));
-            }
+            return new RankedDocsResults.RankedDoc(parsedRankedDoc.id, parsedRankedDoc.score, parsedRankedDoc.text);
         });
     }
 
-    private record RankedDocEntry(@Nullable Integer id, @Nullable Float score, @Nullable String text) {
+    private record RankedDocEntry(Integer id, Float score, @Nullable String text) {
 
         private static final ParseField TEXT = new ParseField("text");
         private static final ParseField SCORE = new ParseField("score");

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/action/HuggingFaceActionCreatorTests.java

@@ -358,11 +358,14 @@ public void testSend_FailsFromInvalidResponseFormat_ForRerankAction() throws IOE
 
     private void assertRerankActionCreator(List<String> documents, String query, int topN, boolean returnText) throws IOException {
         assertThat(webServer.requests(), hasSize(1));
-        assertNull(webServer.requests().getFirst().getUri().getQuery());
-        assertThat(webServer.requests().getFirst().getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType()));
-        assertThat(webServer.requests().getFirst().getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret"));
+        assertNull(webServer.requests().get(0).getUri().getQuery());
+        assertThat(
+            webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE),
+            equalTo(XContentType.JSON.mediaTypeWithoutParameters())
+        );
+        assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret"));
 
-        var requestMap = entityAsMap(webServer.requests().getFirst().getBody());
+        var requestMap = entityAsMap(webServer.requests().get(0).getBody());
         assertThat(requestMap.size(), is(4));
         assertThat(requestMap.get("texts"), is(documents));
         assertThat(requestMap.get("query"), is(query));

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/request/rerank/HuggingFaceRerankRequestEntityTests.java

@@ -18,12 +18,11 @@
 import java.io.IOException;
 import java.util.List;
 
-import static org.elasticsearch.xpack.inference.MatchersUtils.equalToIgnoringWhitespaceInJsonString;
+import static org.elasticsearch.common.xcontent.XContentHelper.stripWhitespace;
 
 public class HuggingFaceRerankRequestEntityTests extends ESTestCase {
     private static final String INPUT = "texts";
     private static final String QUERY = "query";
-    private static final String INFERENCE_ID = "model";
     private static final Integer TOP_N = 8;
     private static final Boolean RETURN_DOCUMENTS = false;
 
@@ -33,36 +32,28 @@ public void testXContent_WritesAllFields_WhenTheyAreDefined() throws IOException
             List.of(INPUT),
             Boolean.TRUE,
             TOP_N,
-            new HuggingFaceRerankTaskSettings(TOP_N, RETURN_DOCUMENTS),
-            INFERENCE_ID
+            new HuggingFaceRerankTaskSettings(TOP_N, RETURN_DOCUMENTS)
         );
 
         XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
         entity.toXContent(builder, ToXContent.EMPTY_PARAMS);
         String xContentResult = Strings.toString(builder);
-
-        assertThat(xContentResult, equalToIgnoringWhitespaceInJsonString("""
+        String expected = """
             {"texts":["texts"],
             "query":"query",
             "return_text":true,
-            "top_n":8}"""));
+            "top_n":8}""";
+        assertEquals(stripWhitespace(expected), xContentResult);
     }
 
     public void testXContent_WritesMinimalFields() throws IOException {
-        var entity = new HuggingFaceRerankRequestEntity(
-            QUERY,
-            List.of(INPUT),
-            null,
-            null,
-            new HuggingFaceRerankTaskSettings(null, null),
-            INFERENCE_ID
-        );
+        var entity = new HuggingFaceRerankRequestEntity(QUERY, List.of(INPUT), null, null, new HuggingFaceRerankTaskSettings(null, null));
 
         XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
         entity.toXContent(builder, ToXContent.EMPTY_PARAMS);
         String xContentResult = Strings.toString(builder);
-
-        assertThat(xContentResult, equalToIgnoringWhitespaceInJsonString("""
-            {"texts":["texts"],"query":"query"}"""));
+        String expected = """
+            {"texts":["texts"],"query":"query"}""";
+        assertEquals(stripWhitespace(expected), xContentResult);
     }
 }