Add Azure AI Rerank support

Author: Evgenii-Kazannik

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

@@ -50,6 +50,8 @@
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsTaskSettings;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankServiceSettings;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.azureopenai.AzureOpenAiSecretSettings;
 import org.elasticsearch.xpack.inference.services.azureopenai.completion.AzureOpenAiCompletionServiceSettings;
 import org.elasticsearch.xpack.inference.services.azureopenai.completion.AzureOpenAiCompletionTaskSettings;
@@ -311,6 +313,17 @@ private static void addAzureAiStudioNamedWriteables(List<NamedWriteableRegistry.
                 AzureAiStudioChatCompletionTaskSettings::new
             )
         );
+
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                AzureAiStudioRerankServiceSettings.NAME,
+                AzureAiStudioRerankServiceSettings::new
+            )
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(TaskSettings.class, AzureAiStudioRerankTaskSettings.NAME, AzureAiStudioRerankTaskSettings::new)
+        );
     }
 
     private static void addAzureOpenAiNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {

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

@@ -10,6 +10,7 @@
 public class AzureAiStudioConstants {
     public static final String EMBEDDINGS_URI_PATH = "/v1/embeddings";
     public static final String COMPLETIONS_URI_PATH = "/v1/chat/completions";
+    public static final String RERANK_URI_PATH = "/v1/rerank";
 
     // common service settings fields
     public static final String TARGET_FIELD = "target";
@@ -35,5 +36,9 @@ public class AzureAiStudioConstants {
     public static final Double MIN_TEMPERATURE_TOP_P = 0.0;
     public static final Double MAX_TEMPERATURE_TOP_P = 2.0;
 
+    // rerank task settings fields
+    public static final String RETURN_DOCUMENTS_FIELD = "return_documents";
+    public static final String TOP_N_FIELD = "top_n";
+
     private AzureAiStudioConstants() {}
 }

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

@@ -0,0 +1,68 @@
+/*
+ * 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.azureaistudio;
+
+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.external.http.retry.RequestSender;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+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.external.response.ErrorMessageResponseEntity;
+import org.elasticsearch.xpack.inference.services.azureaistudio.request.AzureAiStudioRerankRequest;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankModel;
+import org.elasticsearch.xpack.inference.services.azureaistudio.response.AzureAiStudioRerankResponseEntity;
+import org.elasticsearch.xpack.inference.services.azureopenai.response.AzureMistralOpenAiExternalResponseHandler;
+
+import java.util.function.Supplier;
+
+public class AzureAiStudioRerankRequestManager extends AzureAiStudioRequestManager {
+    private static final Logger logger = LogManager.getLogger(AzureAiStudioRerankRequestManager.class);
+
+    private static final ResponseHandler HANDLER = createRerankHandler();
+
+    private final AzureAiStudioRerankModel model;
+
+    public AzureAiStudioRerankRequestManager(AzureAiStudioRerankModel model, ThreadPool threadPool) {
+        super(threadPool, model);
+        this.model = model;
+    }
+
+    @Override
+    public void execute(
+        InferenceInputs inferenceInputs,
+        RequestSender requestSender,
+        Supplier<Boolean> hasRequestRerankFunction,
+        ActionListener<InferenceServiceResults> listener
+    ) {
+        var rerankInput = QueryAndDocsInputs.of(inferenceInputs);
+        AzureAiStudioRerankRequest request = new AzureAiStudioRerankRequest(
+            model,
+            rerankInput.getQuery(),
+            rerankInput.getChunks(),
+            rerankInput.getReturnDocuments(),
+            rerankInput.getTopN()
+        );
+
+        execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestRerankFunction, listener));
+    }
+
+    private static ResponseHandler createRerankHandler() {
+        return new AzureMistralOpenAiExternalResponseHandler(
+            "azure ai studio rerank",
+            new AzureAiStudioRerankResponseEntity(),
+            ErrorMessageResponseEntity::fromResponse,
+            true
+        );
+    }
+
+}

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

@@ -44,6 +44,7 @@
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionTaskSettings;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsServiceSettings;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankModel;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
@@ -307,6 +308,24 @@ private static AzureAiStudioModel createModel(
             return completionModel;
         }
 
+        if (taskType == TaskType.RERANK) {
+            var rerankModel = new AzureAiStudioRerankModel(
+                inferenceEntityId,
+                taskType,
+                NAME,
+                serviceSettings,
+                taskSettings,
+                secretSettings,
+                context
+            );
+            checkProviderAndEndpointTypeForTask(
+                TaskType.RERANK,
+                rerankModel.getServiceSettings().provider(),
+                rerankModel.getServiceSettings().endpointType()
+            );
+            return rerankModel;
+        }
+
         throw new ElasticsearchStatusException(failureMessage, RestStatus.BAD_REQUEST);
     }
 

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

@@ -13,8 +13,10 @@
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioChatCompletionRequestManager;
 import org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioEmbeddingsRequestManager;
+import org.elasticsearch.xpack.inference.services.azureaistudio.AzureAiStudioRerankRequestManager;
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsModel;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankModel;
 
 import java.util.Map;
 import java.util.Objects;
@@ -49,4 +51,12 @@ public ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map
         var errorMessage = constructFailedToSendRequestMessage("Azure AI Studio embeddings");
         return new SenderExecutableAction(sender, requestManager, errorMessage);
     }
+
+    @Override
+    public ExecutableAction create(AzureAiStudioRerankModel rerankModel, Map<String, Object> taskSettings) {
+        var overriddenModel = AzureAiStudioRerankModel.of(rerankModel, taskSettings);
+        var requestManager = new AzureAiStudioRerankRequestManager(overriddenModel, serviceComponents.threadPool());
+        var errorMessage = constructFailedToSendRequestMessage("Azure AI Studio rerank");
+        return new SenderExecutableAction(sender, requestManager, errorMessage);
+    }
 }

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

@@ -10,11 +10,14 @@
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.azureaistudio.embeddings.AzureAiStudioEmbeddingsModel;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankModel;
 
 import java.util.Map;
 
 public interface AzureAiStudioActionVisitor {
     ExecutableAction create(AzureAiStudioEmbeddingsModel embeddingsModel, Map<String, Object> taskSettings);
 
     ExecutableAction create(AzureAiStudioChatCompletionModel completionModel, Map<String, Object> taskSettings);
+
+    ExecutableAction create(AzureAiStudioRerankModel rerankModel, Map<String, Object> taskSettings);
 }

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

@@ -0,0 +1,80 @@
+/*
+ * 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.azureaistudio.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.core.Nullable;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.inference.external.request.HttpRequest;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankModel;
+
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Objects;
+
+public class AzureAiStudioRerankRequest extends AzureAiStudioRequest {
+    private final String query;
+    private final List<String> input;
+    private final Boolean returnDocuments;
+    private final Integer topN;
+    private final AzureAiStudioRerankModel rerankModel;
+
+    public AzureAiStudioRerankRequest(
+        AzureAiStudioRerankModel model,
+        String query,
+        List<String> input,
+        @Nullable Boolean returnDocuments,
+        @Nullable Integer topN
+    ) {
+        super(model);
+        this.rerankModel = Objects.requireNonNull(model);
+        this.query = query;
+        this.input = Objects.requireNonNull(input);
+        this.returnDocuments = returnDocuments;
+        this.topN = topN;
+    }
+
+    @Override
+    public HttpRequest createHttpRequest() {
+        HttpPost httpPost = new HttpPost(this.uri);
+
+        ByteArrayEntity byteEntity = new ByteArrayEntity(Strings.toString(createRequestEntity()).getBytes(StandardCharsets.UTF_8));
+        httpPost.setEntity(byteEntity);
+
+        httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
+        setAuthHeader(httpPost, rerankModel);
+
+        return new HttpRequest(httpPost, getInferenceEntityId());
+    }
+
+    @Override
+    public Request truncate() {
+        // no truncation
+        return this;
+    }
+
+    @Override
+    public boolean[] getTruncationInfo() {
+        // no truncation
+        return null;
+    }
+
+    public Integer getTopN() {
+        return topN != null ? topN : rerankModel.getTaskSettings().topN();
+    }
+
+    private AzureAiStudioRerankRequestEntity createRequestEntity() {
+        var taskSettings = rerankModel.getTaskSettings();
+        return new AzureAiStudioRerankRequestEntity(query, input, returnDocuments, topN, taskSettings);
+    }
+
+}

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

@@ -0,0 +1,59 @@
+/*
+ * 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.azureaistudio.request;
+
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xpack.inference.services.azureaistudio.rerank.AzureAiStudioRerankTaskSettings;
+import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankTaskSettings;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+public record AzureAiStudioRerankRequestEntity(
+    String query,
+    List<String> input,
+    @Nullable Boolean returnDocuments,
+    @Nullable Integer topN,
+    AzureAiStudioRerankTaskSettings taskSettings
+) implements ToXContentObject {
+
+    private static final String RETURN_TEXT = "return_text";
+    private static final String DOCUMENTS_FIELD = "texts";
+    private static final String QUERY = "query";
+
+    public AzureAiStudioRerankRequestEntity {
+        Objects.requireNonNull(query);
+        Objects.requireNonNull(input);
+        Objects.requireNonNull(taskSettings);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+
+        builder.field(DOCUMENTS_FIELD, input);
+        builder.field(QUERY, query);
+
+        if (returnDocuments != null) {
+            builder.field(RETURN_TEXT, returnDocuments);
+        } else if (taskSettings.returnDocuments() != null) {
+            builder.field(RETURN_TEXT, taskSettings.returnDocuments());
+        }
+
+        if (topN != null) {
+            builder.field(HuggingFaceRerankTaskSettings.TOP_N_DOCS_ONLY, topN);
+        } else if (taskSettings.topN() != null) {
+            builder.field(HuggingFaceRerankTaskSettings.TOP_N_DOCS_ONLY, taskSettings.topN());
+        }
+        builder.endObject();
+        return builder;
+    }
+}