Extend huggingface with rerank

Author: Evgenii-Kazannik

server/src/test/java/org/elasticsearch/inference/SettingsConfigurationTestUtils.java

@@ -20,8 +20,8 @@
 public class SettingsConfigurationTestUtils {
 
     public static SettingsConfiguration getRandomSettingsConfigurationField() {
-        return new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.SPARSE_EMBEDDING)).setDefaultValue(
-            randomAlphaOfLength(10)
+        return new SettingsConfiguration.Builder(EnumSet.of(TaskType.TEXT_EMBEDDING, TaskType.SPARSE_EMBEDDING, TaskType.RERANK))
+            .setDefaultValue(randomAlphaOfLength(10)
         )
             .setDescription(randomAlphaOfLength(10))
             .setLabel(randomAlphaOfLength(10))

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceBaseRestTest.java

@@ -171,6 +171,22 @@ static String mockDenseServiceModelConfig() {
             """;
     }
 
+    static String mockRerankServiceModelConfig() {
+        return """
+            {
+              "task_type": "rerank",
+              "service": "rerank_test_service",
+              "service_settings": {
+                "model": "rerank_model",
+                "api_key": "abc64"
+              },
+              "task_settings": {
+                 "return_documents": true
+              }
+            }
+            """;
+    }
+
     static void deleteModel(String modelId) throws IOException {
         var request = new Request("DELETE", "_inference/" + modelId);
         var response = client().performRequest(request);
@@ -484,6 +500,10 @@ private String jsonBody(List<String> input, @Nullable String query) {
     @SuppressWarnings("unchecked")
     protected void assertNonEmptyInferenceResults(Map<String, Object> resultMap, int expectedNumberOfResults, TaskType taskType) {
         switch (taskType) {
+            case RERANK -> {
+                var results = (List<Map<String, Object>>) resultMap.get(TaskType.RERANK.toString());
+                assertThat(results, hasSize(expectedNumberOfResults));
+            }
             case SPARSE_EMBEDDING -> {
                 var results = (List<Map<String, Object>>) resultMap.get(TaskType.SPARSE_EMBEDDING.toString());
                 assertThat(results, hasSize(expectedNumberOfResults));

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java

@@ -47,6 +47,9 @@ public class InferenceCrudIT extends InferenceBaseRestTest {
 
     @SuppressWarnings("unchecked")
     public void testCRUD() throws IOException {
+        for (int i = 0; i < 6; i++) {
+            putModel("r_model_" + i, mockRerankServiceModelConfig(), TaskType.RERANK);
+        }
         for (int i = 0; i < 5; i++) {
             putModel("se_model_" + i, mockSparseServiceModelConfig(), TaskType.SPARSE_EMBEDDING);
         }
@@ -55,9 +58,16 @@ public void testCRUD() throws IOException {
         }
 
         var getAllModels = getAllModels();
-        int numModels = 12;
+        int numModels = 18;
         assertThat(getAllModels, hasSize(numModels));
 
+        var getRerankModels = getModels("_all", TaskType.RERANK);
+        int numRerankModels = 6;
+        assertThat(getRerankModels, hasSize(numRerankModels));
+        for (var rerankModel : getRerankModels) {
+            assertEquals("rerank", rerankModel.get("task_type"));
+        }
+
         var getSparseModels = getModels("_all", TaskType.SPARSE_EMBEDDING);
         int numSparseModels = 6;
         assertThat(getSparseModels, hasSize(numSparseModels));
@@ -94,6 +104,9 @@ public void testCRUD() throws IOException {
             assertNotEquals(oldApiKey, newApiKey);
             assertEquals(updatedEndpoint, singleModel.get(0));
         }
+        for (int i = 0; i < 6; i++) {
+            deleteModel("r_model_" + i, TaskType.RERANK);
+        }
         for (int i = 0; i < 5; i++) {
             deleteModel("se_model_" + i, TaskType.SPARSE_EMBEDDING);
         }

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

@@ -79,6 +79,7 @@
 import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceServiceSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserServiceSettings;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.embeddings.IbmWatsonxEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankServiceSettings;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankTaskSettings;
@@ -353,6 +354,13 @@ private static void addHuggingFaceNamedWriteables(List<NamedWriteableRegistry.En
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(ServiceSettings.class, HuggingFaceServiceSettings.NAME, HuggingFaceServiceSettings::new)
         );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                HuggingFaceRerankServiceSettings.NAME,
+                HuggingFaceRerankServiceSettings::new
+            )
+        );
     }
 
     private static void addGoogleAiStudioNamedWritables(List<NamedWriteableRegistry.Entry> namedWriteables) {

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

@@ -128,6 +128,7 @@
 import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiService;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceService;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserService;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankService;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxService;
 import org.elasticsearch.xpack.inference.services.jinaai.JinaAIService;
 import org.elasticsearch.xpack.inference.services.mistral.MistralService;
@@ -361,6 +362,7 @@ public void loadExtensions(ExtensionLoader loader) {
     public List<InferenceServiceExtension.Factory> getInferenceServiceFactories() {
         return List.of(
             context -> new HuggingFaceElserService(httpFactory.get(), serviceComponents.get()),
+            context -> new HuggingFaceRerankService(httpFactory.get(), serviceComponents.get()),
             context -> new HuggingFaceService(httpFactory.get(), serviceComponents.get()),
             context -> new OpenAiService(httpFactory.get(), serviceComponents.get()),
             context -> new CohereService(httpFactory.get(), serviceComponents.get()),

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

@@ -0,0 +1,78 @@
+/*
+ * 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.services.huggingface;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.common.Truncator;
+import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+import org.elasticsearch.xpack.inference.external.http.sender.BaseRequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.ExecutableInferenceRequest;
+import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
+import org.elasticsearch.xpack.inference.external.http.sender.QueryAndDocsInputs;
+import org.elasticsearch.xpack.inference.services.huggingface.request.HuggingFaceInferenceRerankRequest;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Supplier;
+
+import static org.elasticsearch.xpack.inference.common.Truncator.truncate;
+
+public class HuggingFaceRequestRerankManager extends BaseRequestManager {
+    private static final Logger logger = LogManager.getLogger(HuggingFaceRequestRerankManager.class);
+
+    public static HuggingFaceRequestRerankManager of(
+        HuggingFaceModel model,
+        ResponseHandler responseHandler,
+        Truncator truncator,
+        ThreadPool threadPool
+    ) {
+        return new HuggingFaceRequestRerankManager(
+            Objects.requireNonNull(model),
+            Objects.requireNonNull(responseHandler),
+            Objects.requireNonNull(truncator),
+            Objects.requireNonNull(threadPool)
+        );
+    }
+
+    private final HuggingFaceModel model;
+    private final ResponseHandler responseHandler;
+    private final Truncator truncator;
+
+    private HuggingFaceRequestRerankManager(HuggingFaceModel model, ResponseHandler responseHandler, Truncator truncator, ThreadPool threadPool) {
+        super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
+        this.model = model;
+        this.responseHandler = responseHandler;
+        this.truncator = truncator;
+    }
+
+    @Override
+    public void execute(
+        InferenceInputs inferenceInputs,
+        RequestSender requestSender,
+        Supplier<Boolean> hasRequestCompletedFunction,
+        ActionListener<InferenceServiceResults> listener
+    ) {
+        List<String> docsInput = QueryAndDocsInputs.of(inferenceInputs).getChunks();
+        var truncatedInput = truncate(docsInput, model.getTokenLimit());
+        var request = new HuggingFaceInferenceRerankRequest(truncator, truncatedInput, model);
+
+        execute(new ExecutableInferenceRequest(requestSender, logger, request, responseHandler, hasRequestCompletedFunction, listener));
+    }
+
+    record RateLimitGrouping(int accountHash) {
+
+        public static RateLimitGrouping of(HuggingFaceModel model) {
+            return new RateLimitGrouping(new HuggingFaceAccount(model.rateLimitServiceSettings().uri(), model.apiKey()).hashCode());
+        }
+    }
+}

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

@@ -10,13 +10,16 @@
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.external.response.huggingface.HuggingFaceRerankResponseEntity;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceRequestManager;
+import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceRequestRerankManager;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceResponseHandler;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserModel;
 import org.elasticsearch.xpack.inference.services.huggingface.embeddings.HuggingFaceEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.huggingface.response.HuggingFaceElserResponseEntity;
 import org.elasticsearch.xpack.inference.services.huggingface.response.HuggingFaceEmbeddingsResponseEntity;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankModel;
 
 import java.util.Objects;
 
@@ -34,6 +37,23 @@ public HuggingFaceActionCreator(Sender sender, ServiceComponents serviceComponen
         this.serviceComponents = Objects.requireNonNull(serviceComponents);
     }
 
+    @Override
+    public ExecutableAction create(HuggingFaceRerankModel model) {
+        var responseHandler = new HuggingFaceResponseHandler("hugging face rerank", HuggingFaceRerankResponseEntity::fromResponse);
+        var requestCreator = HuggingFaceRequestRerankManager.of(
+            model,
+            responseHandler,
+            serviceComponents.truncator(),
+            serviceComponents.threadPool()
+        );
+        var errorMessage = format(
+            "Failed to send Hugging Face %s request from inference entity id [%s]",
+            "rerank",
+            model.getInferenceEntityId()
+        );
+        return new SenderExecutableAction(sender, requestCreator, errorMessage);
+    }
+
     @Override
     public ExecutableAction create(HuggingFaceEmbeddingsModel model) {
         var responseHandler = new HuggingFaceResponseHandler(

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

@@ -10,8 +10,11 @@
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserModel;
 import org.elasticsearch.xpack.inference.services.huggingface.embeddings.HuggingFaceEmbeddingsModel;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankModel;
 
 public interface HuggingFaceActionVisitor {
+    ExecutableAction create(HuggingFaceRerankModel model);
+
     ExecutableAction create(HuggingFaceEmbeddingsModel model);
 
     ExecutableAction create(HuggingFaceElserModel model);

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

@@ -0,0 +1,74 @@
+/*
+ * 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.services.huggingface.request;
+
+import org.apache.http.HttpHeaders;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.entity.ByteArrayEntity;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.inference.common.Truncator;
+import org.elasticsearch.xpack.inference.external.request.HttpRequest;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceAccount;
+import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceModel;
+
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.inference.external.request.RequestUtils.createAuthBearerHeader;
+
+public class HuggingFaceInferenceRerankRequest implements Request {
+
+    private final Truncator truncator;
+    private final HuggingFaceAccount account;
+    private final Truncator.TruncationResult truncationResult;
+    private final HuggingFaceModel model;
+
+    public HuggingFaceInferenceRerankRequest(Truncator truncator, Truncator.TruncationResult input, HuggingFaceModel model) {
+        this.truncator = Objects.requireNonNull(truncator);
+        this.account = HuggingFaceAccount.of(model);
+        this.truncationResult = Objects.requireNonNull(input);
+        this.model = Objects.requireNonNull(model);
+    }
+
+    public HttpRequest createHttpRequest() {
+        HttpPost httpPost = new HttpPost(account.uri());
+
+        ByteArrayEntity byteEntity = new ByteArrayEntity(
+            Strings.toString(new HuggingFaceInferenceRerankRequestEntity(truncationResult.input(), "default")).getBytes(StandardCharsets.UTF_8)
+        );
+        httpPost.setEntity(byteEntity);
+        httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaTypeWithoutParameters());
+        httpPost.setHeader(createAuthBearerHeader(account.apiKey()));
+
+        return new HttpRequest(httpPost, getInferenceEntityId());
+    }
+
+    public URI getURI() {
+        return account.uri();
+    }
+
+    @Override
+    public String getInferenceEntityId() {
+        return model.getInferenceEntityId();
+    }
+
+    @Override
+    public Request truncate() {
+        var truncateResult = truncator.truncate(truncationResult.input());
+
+        return new HuggingFaceInferenceRerankRequest(truncator, truncateResult, model);
+    }
+
+    @Override
+    public boolean[] getTruncationInfo() {
+        return truncationResult.truncated().clone();
+    }
+}

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

@@ -0,0 +1,37 @@
+/*
+ * 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.services.huggingface.request;
+
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+public record HuggingFaceInferenceRerankRequestEntity(List<String> documents, String query) implements ToXContentObject {
+
+    private static final String DOCUMENTS_FIELD = "documents";
+    private static final String QUERY_FIELD = "query";
+
+    public HuggingFaceInferenceRerankRequestEntity {
+        Objects.requireNonNull(documents);
+        Objects.requireNonNull(query);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+
+        builder.field(DOCUMENTS_FIELD, documents);
+        builder.field(QUERY_FIELD, query);
+
+        builder.endObject();
+        return builder;
+    }
+}