Adding rerank common options

Author: ymao1

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/voyageai/VoyageAIActionCreator.java

@@ -75,7 +75,13 @@ public ExecutableAction create(VoyageAIRerankModel model, Map<String, Object> ta
             serviceComponents.threadPool(),
             overriddenModel,
             RERANK_HANDLER,
-            (rerankInput) -> new VoyageAIRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model),
+            (rerankInput) -> new VoyageAIRerankRequest(
+                rerankInput.getQuery(),
+                rerankInput.getChunks(),
+                rerankInput.getReturnDocuments(),
+                rerankInput.getTopN(),
+                model
+            ),
             QueryAndDocsInputs.class
         );
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/AlibabaCloudSearchRerankRequestManager.java

@@ -69,6 +69,8 @@ public void execute(
             account,
             rerankInput.getQuery(),
             rerankInput.getChunks(),
+            rerankInput.getReturnDocuments(),
+            rerankInput.getTopN(),
             model
         );
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CohereRerankRequestManager.java

@@ -49,7 +49,13 @@ public void execute(
         ActionListener<InferenceServiceResults> listener
     ) {
         var rerankInput = QueryAndDocsInputs.of(inferenceInputs);
-        CohereRerankRequest request = new CohereRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model);
+        CohereRerankRequest request = new CohereRerankRequest(
+            rerankInput.getQuery(),
+            rerankInput.getChunks(),
+            rerankInput.getReturnDocuments(),
+            rerankInput.getTopN(),
+            model
+        );
 
         execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/GoogleVertexAiRerankRequestManager.java

@@ -62,7 +62,13 @@ public void execute(
         ActionListener<InferenceServiceResults> listener
     ) {
         var rerankInput = QueryAndDocsInputs.of(inferenceInputs);
-        GoogleVertexAiRerankRequest request = new GoogleVertexAiRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model);
+        GoogleVertexAiRerankRequest request = new GoogleVertexAiRerankRequest(
+            rerankInput.getQuery(),
+            rerankInput.getChunks(),
+            rerankInput.getReturnDocuments(),
+            rerankInput.getTopN(),
+            model
+        );
 
         execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/JinaAIRerankRequestManager.java

@@ -49,7 +49,13 @@ public void execute(
         ActionListener<InferenceServiceResults> listener
     ) {
         var rerankInput = QueryAndDocsInputs.of(inferenceInputs);
-        JinaAIRerankRequest request = new JinaAIRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model);
+        JinaAIRerankRequest request = new JinaAIRerankRequest(
+            rerankInput.getQuery(),
+            rerankInput.getChunks(),
+            rerankInput.getReturnDocuments(),
+            rerankInput.getTopN(),
+            model
+        );
 
         execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/QueryAndDocsInputs.java

@@ -7,6 +7,8 @@
 
 package org.elasticsearch.xpack.inference.external.http.sender;
 
+import org.elasticsearch.core.Nullable;
+
 import java.util.List;
 import java.util.Objects;
 
@@ -22,15 +24,25 @@ public static QueryAndDocsInputs of(InferenceInputs inferenceInputs) {
 
     private final String query;
     private final List<String> chunks;
+    private final Boolean returnDocuments;
+    private final Integer topN;
 
     public QueryAndDocsInputs(String query, List<String> chunks) {
-        this(query, chunks, false);
+        this(query, chunks, null, null, false);
     }
 
-    public QueryAndDocsInputs(String query, List<String> chunks, boolean stream) {
+    public QueryAndDocsInputs(
+        String query,
+        List<String> chunks,
+        @Nullable Boolean returnDocuments,
+        @Nullable Integer topN,
+        boolean stream
+    ) {
         super(stream);
         this.query = Objects.requireNonNull(query);
         this.chunks = Objects.requireNonNull(chunks);
+        this.returnDocuments = returnDocuments;
+        this.topN = topN;
     }
 
     public String getQuery() {
@@ -41,6 +53,14 @@ public List<String> getChunks() {
         return chunks;
     }
 
+    public Boolean getReturnDocuments() {
+        return returnDocuments;
+    }
+
+    public Integer getTopN() {
+        return topN;
+    }
+
     public int inputSize() {
         return chunks.size();
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/alibabacloudsearch/AlibabaCloudSearchRerankRequest.java

@@ -12,6 +12,7 @@
 import org.apache.http.client.utils.URIBuilder;
 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.alibabacloudsearch.AlibabaCloudSearchAccount;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
@@ -32,6 +33,8 @@ public class AlibabaCloudSearchRerankRequest implements Request {
     private final AlibabaCloudSearchAccount account;
     private final String query;
     private final List<String> input;
+    private final Boolean returnDocuments;
+    private final Integer topN;
     private final URI uri;
     private final AlibabaCloudSearchRerankTaskSettings taskSettings;
     private final String model;
@@ -44,13 +47,17 @@ public AlibabaCloudSearchRerankRequest(
         AlibabaCloudSearchAccount account,
         String query,
         List<String> input,
+        @Nullable Boolean returnDocuments,
+        @Nullable Integer topN,
         AlibabaCloudSearchRerankModel rerankModel
     ) {
         Objects.requireNonNull(rerankModel);
 
         this.account = Objects.requireNonNull(account);
         this.query = Objects.requireNonNull(query);
         this.input = Objects.requireNonNull(input);
+        this.returnDocuments = returnDocuments;
+        this.topN = topN;
         taskSettings = rerankModel.getTaskSettings();
         model = rerankModel.getServiceSettings().getCommonSettings().modelId();
         host = rerankModel.getServiceSettings().getCommonSettings().getHost();
@@ -67,7 +74,8 @@ public HttpRequest createHttpRequest() {
         HttpPost httpPost = new HttpPost(uri);
 
         ByteArrayEntity byteEntity = new ByteArrayEntity(
-            Strings.toString(new AlibabaCloudSearchRerankRequestEntity(query, input, taskSettings)).getBytes(StandardCharsets.UTF_8)
+            Strings.toString(new AlibabaCloudSearchRerankRequestEntity(query, input, returnDocuments, topN, taskSettings))
+                .getBytes(StandardCharsets.UTF_8)
         );
         httpPost.setEntity(byteEntity);
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/alibabacloudsearch/AlibabaCloudSearchRerankRequestEntity.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.request.alibabacloudsearch;
 
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.rerank.AlibabaCloudSearchRerankTaskSettings;
@@ -15,9 +16,13 @@
 import java.util.List;
 import java.util.Objects;
 
-public record AlibabaCloudSearchRerankRequestEntity(String query, List<String> input, AlibabaCloudSearchRerankTaskSettings taskSettings)
-    implements
-        ToXContentObject {
+public record AlibabaCloudSearchRerankRequestEntity(
+    String query,
+    List<String> input,
+    @Nullable Boolean returnDocuments,
+    @Nullable Integer topN,
+    AlibabaCloudSearchRerankTaskSettings taskSettings
+) implements ToXContentObject {
 
     private static final String SEARCH_QUERY = "query";
     private static final String TEXTS_FIELD = "docs";

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/cohere/CohereRerankRequest.java

@@ -11,6 +11,7 @@
 import org.apache.http.client.utils.URIBuilder;
 import org.apache.http.entity.ByteArrayEntity;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xpack.inference.external.cohere.CohereAccount;
 import org.elasticsearch.xpack.inference.external.request.HttpRequest;
 import org.elasticsearch.xpack.inference.external.request.Request;
@@ -28,16 +29,26 @@ public class CohereRerankRequest extends CohereRequest {
     private final CohereAccount account;
     private final String query;
     private final List<String> input;
+    private final Boolean returnDocuments;
+    private final Integer topN;
     private final CohereRerankTaskSettings taskSettings;
     private final String model;
     private final String inferenceEntityId;
 
-    public CohereRerankRequest(String query, List<String> input, CohereRerankModel model) {
+    public CohereRerankRequest(
+        String query,
+        List<String> input,
+        @Nullable Boolean returnDocuments,
+        @Nullable Integer topN,
+        CohereRerankModel model
+    ) {
         Objects.requireNonNull(model);
 
         this.account = CohereAccount.of(model, CohereRerankRequest::buildDefaultUri);
         this.input = Objects.requireNonNull(input);
         this.query = Objects.requireNonNull(query);
+        this.returnDocuments = returnDocuments;
+        this.topN = topN;
         taskSettings = model.getTaskSettings();
         this.model = model.getServiceSettings().modelId();
         inferenceEntityId = model.getInferenceEntityId();
@@ -48,7 +59,8 @@ public HttpRequest createHttpRequest() {
         HttpPost httpPost = new HttpPost(account.uri());
 
         ByteArrayEntity byteEntity = new ByteArrayEntity(
-            Strings.toString(new CohereRerankRequestEntity(query, input, taskSettings, model)).getBytes(StandardCharsets.UTF_8)
+            Strings.toString(new CohereRerankRequestEntity(query, input, returnDocuments, topN, taskSettings, model))
+                .getBytes(StandardCharsets.UTF_8)
         );
         httpPost.setEntity(byteEntity);
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/cohere/CohereRerankRequestEntity.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.inference.external.request.cohere;
 
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xpack.inference.services.cohere.rerank.CohereRerankTaskSettings;
@@ -15,9 +16,14 @@
 import java.util.List;
 import java.util.Objects;
 
-public record CohereRerankRequestEntity(String model, String query, List<String> documents, CohereRerankTaskSettings taskSettings)
-    implements
-        ToXContentObject {
+public record CohereRerankRequestEntity(
+    String model,
+    String query,
+    List<String> documents,
+    @Nullable Boolean returnDocuments,
+    @Nullable Integer topN,
+    CohereRerankTaskSettings taskSettings
+) implements ToXContentObject {
 
     private static final String DOCUMENTS_FIELD = "documents";
     private static final String QUERY_FIELD = "query";
@@ -29,8 +35,15 @@ public record CohereRerankRequestEntity(String model, String query, List<String>
         Objects.requireNonNull(taskSettings);
     }
 
-    public CohereRerankRequestEntity(String query, List<String> input, CohereRerankTaskSettings taskSettings, String model) {
-        this(model, query, input, taskSettings);
+    public CohereRerankRequestEntity(
+        String query,
+        List<String> input,
+        @Nullable Boolean returnDocuments,
+        @Nullable Integer topN,
+        CohereRerankTaskSettings taskSettings,
+        String model
+    ) {
+        this(model, query, input, returnDocuments, topN, taskSettings);
     }
 
     @Override
@@ -41,11 +54,17 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         builder.field(QUERY_FIELD, query);
         builder.field(DOCUMENTS_FIELD, documents);
 
-        if (taskSettings.getDoesReturnDocuments() != null) {
+        // prefer the root level return_documents over task settings
+        if (returnDocuments != null) {
+            builder.field(CohereRerankTaskSettings.RETURN_DOCUMENTS, returnDocuments);
+        } else if (taskSettings.getDoesReturnDocuments() != null) {
             builder.field(CohereRerankTaskSettings.RETURN_DOCUMENTS, taskSettings.getDoesReturnDocuments());
         }
 
-        if (taskSettings.getTopNDocumentsOnly() != null) {
+        // prefer the root level top_n over task settings
+        if (topN != null) {
+            builder.field(CohereRerankTaskSettings.TOP_N_DOCS_ONLY, topN);
+        } else if (taskSettings.getTopNDocumentsOnly() != null) {
             builder.field(CohereRerankTaskSettings.TOP_N_DOCS_ONLY, taskSettings.getTopNDocumentsOnly());
         }