[Inference API] Add "rerank" task type to "elastic" provider

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

@@ -214,6 +214,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_REMOVE_AGGREGATE_TYPE = def(9_045_0_00);
     public static final TransportVersion ADD_PROJECT_ID_TO_DSL_ERROR_INFO = def(9_046_0_00);
     public static final TransportVersion SEMANTIC_TEXT_CHUNKING_CONFIG = def(9_047_00_0);
+    public static final TransportVersion ML_INFERENCE_ELASTIC_RERANK = def(9_048_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -66,6 +66,27 @@ protected RankedDocsResults mutateInstanceForVersion(RankedDocsResults instance,
         }
     }
 
+    public record RerankExpectation(int index, float relevanceScore) {}
+
+    public static Map<String, Object> buildExpectationRankedDocResults(List<RerankExpectation> rerankExpectations) {
+        return Map.of(
+            RankedDocsResults.RERANK,
+            rerankExpectations.stream()
+                .map(
+                    rerankExpectation -> Map.of(
+                        RankedDocsResults.RankedDoc.NAME,
+                        Map.of(
+                            RankedDocsResults.RankedDoc.INDEX,
+                            rerankExpectation.index,
+                            RankedDocsResults.RankedDoc.RELEVANCE_SCORE,
+                            rerankExpectation.relevanceScore
+                        )
+                    )
+                )
+                .toList()
+        );
+    }
+
     private List<RankedDocsResults.RankedDoc> rankedDocsNullStringToEmpty(List<RankedDocsResults.RankedDoc> rankedDocs) {
         var result = new ArrayList<RankedDocsResults.RankedDoc>(rankedDocs.size());
         for (var doc : rankedDocs) {

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

@@ -60,6 +60,8 @@
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.deepseek.DeepSeekChatCompletionModel;
 import org.elasticsearch.xpack.inference.services.elastic.completion.ElasticInferenceServiceCompletionServiceSettings;
+import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankServiceSettings;
+import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.elasticsearch.CustomElandInternalServiceSettings;
 import org.elasticsearch.xpack.inference.services.elasticsearch.CustomElandInternalTextEmbeddingServiceSettings;
@@ -147,7 +149,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         addAnthropicNamedWritables(namedWriteables);
         addAmazonBedrockNamedWriteables(namedWriteables);
         addAwsNamedWriteables(namedWriteables);
-        addEisNamedWriteables(namedWriteables);
+        addElasticNamedWriteables(namedWriteables);
         addAlibabaCloudSearchNamedWriteables(namedWriteables);
         addJinaAINamedWriteables(namedWriteables);
         addVoyageAINamedWriteables(namedWriteables);
@@ -646,20 +648,39 @@ private static void addVoyageAINamedWriteables(List<NamedWriteableRegistry.Entry
         );
     }
 
-    private static void addEisNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+    private static void addElasticNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {
+        // Sparse Text Embeddings
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(
                 ServiceSettings.class,
                 ElasticInferenceServiceSparseEmbeddingsServiceSettings.NAME,
                 ElasticInferenceServiceSparseEmbeddingsServiceSettings::new
             )
         );
+
+        // Completion
         namedWriteables.add(
             new NamedWriteableRegistry.Entry(
                 ServiceSettings.class,
                 ElasticInferenceServiceCompletionServiceSettings.NAME,
                 ElasticInferenceServiceCompletionServiceSettings::new
             )
         );
+
+        // Rerank
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                ServiceSettings.class,
+                ElasticInferenceServiceRerankServiceSettings.NAME,
+                ElasticInferenceServiceRerankServiceSettings::new
+            )
+        );
+        namedWriteables.add(
+            new NamedWriteableRegistry.Entry(
+                TaskSettings.class,
+                ElasticInferenceServiceRerankTaskSettings.NAME,
+                ElasticInferenceServiceRerankTaskSettings::new
+            )
+        );
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/elastic/rerank/ElasticInferenceServiceRerankRequest.java

@@ -0,0 +1,98 @@
+/*
+ * 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.external.request.elastic.rerank;
+
+import org.apache.http.HttpHeaders;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.client.methods.HttpRequestBase;
+import org.apache.http.entity.ByteArrayEntity;
+import org.apache.http.message.BasicHeader;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.external.request.elastic.ElasticInferenceServiceRequest;
+import org.elasticsearch.xpack.inference.external.request.elastic.ElasticInferenceServiceRequestMetadata;
+import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankModel;
+import org.elasticsearch.xpack.inference.telemetry.TraceContext;
+import org.elasticsearch.xpack.inference.telemetry.TraceContextHandler;
+
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Objects;
+
+public class ElasticInferenceServiceRerankRequest extends ElasticInferenceServiceRequest {
+
+    private final URI uri;
+    private final String query;
+    private final List<String> documents;
+    private final TraceContextHandler traceContextHandler;
+    private final ElasticInferenceServiceRerankModel model;
+
+    public ElasticInferenceServiceRerankRequest(
+        String query,
+        List<String> documents,
+        ElasticInferenceServiceRerankModel model,
+        TraceContext traceContext,
+        ElasticInferenceServiceRequestMetadata metadata
+    ) {
+        super(metadata);
+        this.query = query;
+        this.documents = documents;
+        this.model = Objects.requireNonNull(model);
+        this.uri = model.uri();
+        this.traceContextHandler = new TraceContextHandler(traceContext);
+    }
+
+    @Override
+    public HttpRequestBase createHttpRequestBase() {
+        var httpPost = new HttpPost(uri);
+        var requestEntity = Strings.toString(
+            new ElasticInferenceServiceRerankRequestEntity(
+                query,
+                documents,
+                model.getServiceSettings().modelId(),
+                model.getTaskSettings().getTopNDocumentsOnly()
+            )
+        );
+
+        ByteArrayEntity byteEntity = new ByteArrayEntity(requestEntity.getBytes(StandardCharsets.UTF_8));
+        httpPost.setEntity(byteEntity);
+
+        traceContextHandler.propagateTraceContext(httpPost);
+        httpPost.setHeader(new BasicHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType()));
+
+        return httpPost;
+    }
+
+    public TraceContext getTraceContext() {
+        return traceContextHandler.traceContext();
+    }
+
+    @Override
+    public String getInferenceEntityId() {
+        return model.getInferenceEntityId();
+    }
+
+    @Override
+    public URI getURI() {
+        return uri;
+    }
+
+    @Override
+    public Request truncate() {
+        // no truncation
+        return this;
+    }
+
+    @Override
+    public boolean[] getTruncationInfo() {
+        // no truncation
+        return null;
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/elastic/rerank/ElasticInferenceServiceRerankRequestEntity.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.external.request.elastic.rerank;
+
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+public record ElasticInferenceServiceRerankRequestEntity(
+    String query,
+    List<String> documents,
+    String modelId,
+    @Nullable Integer topNDocumentsOnly
+) implements ToXContentObject {
+
+    private static final String QUERY_FIELD = "query";
+    private static final String MODEL_FIELD = "model";
+    private static final String TOP_N_DOCUMENTS_ONLY_FIELD = "top_n";
+    private static final String DOCUMENTS_FIELD = "documents";
+
+    public ElasticInferenceServiceRerankRequestEntity {
+        Objects.requireNonNull(query);
+        Objects.requireNonNull(documents);
+        Objects.requireNonNull(modelId);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        builder.startObject();
+
+        builder.field(QUERY_FIELD, query);
+
+        builder.field(MODEL_FIELD, modelId);
+
+        if (Objects.nonNull(topNDocumentsOnly)) {
+            builder.field(TOP_N_DOCUMENTS_ONLY_FIELD, topNDocumentsOnly);
+        }
+
+        builder.startArray(DOCUMENTS_FIELD);
+        for (String document : documents) {
+            builder.value(document);
+        }
+
+        builder.endArray();
+
+        builder.endObject();
+
+        return builder;
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/elastic/ElasticInferenceServiceRerankResponseEntity.java

@@ -0,0 +1,70 @@
+/*
+ * 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.external.response.elastic;
+
+import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.xcontent.ConstructingObjectParser;
+import org.elasticsearch.xcontent.ParseField;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentParser;
+import org.elasticsearch.xcontent.XContentParserConfiguration;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
+import org.elasticsearch.xpack.inference.external.http.HttpResult;
+
+import java.io.IOException;
+import java.util.List;
+
+import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
+
+public class ElasticInferenceServiceRerankResponseEntity {
+
+    record RerankResult(List<RerankResultEntry> entries) {
+
+        @SuppressWarnings("unchecked")
+        public static final ConstructingObjectParser<RerankResult, Void> PARSER = new ConstructingObjectParser<>(
+            RerankResult.class.getSimpleName(),
+            true,
+            args -> new RerankResult((List<RerankResultEntry>) args[0])
+        );
+
+        static {
+            PARSER.declareObjectArray(constructorArg(), RerankResultEntry.PARSER::apply, new ParseField("results"));
+        }
+
+        record RerankResultEntry(Integer index, Float relevanceScore) {
+
+            public static final ConstructingObjectParser<RerankResultEntry, Void> PARSER = new ConstructingObjectParser<>(
+                RerankResultEntry.class.getSimpleName(),
+                args -> new RerankResultEntry((Integer) args[0], (Float) args[1])
+            );
+
+            static {
+                PARSER.declareInt(constructorArg(), new ParseField("index"));
+                PARSER.declareFloat(constructorArg(), new ParseField("relevance_score"));
+            }
+
+            public RankedDocsResults.RankedDoc toRankedDoc() {
+                return new RankedDocsResults.RankedDoc(index, relevanceScore, null);
+            }
+        }
+    }
+
+    public static InferenceServiceResults fromResponse(HttpResult response) throws IOException {
+        var parserConfig = XContentParserConfiguration.EMPTY.withDeprecationHandler(LoggingDeprecationHandler.INSTANCE);
+
+        try (XContentParser jsonParser = XContentFactory.xContent(XContentType.JSON).createParser(parserConfig, response.body())) {
+            var rerankResult = RerankResult.PARSER.apply(jsonParser, null);
+
+            return new RankedDocsResults(rerankResult.entries.stream().map(RerankResult.RerankResultEntry::toRankedDoc).toList());
+        }
+    }
+
+    private ElasticInferenceServiceRerankResponseEntity() {}
+}