Add pagination / document limit logic for reranking

src/main/java/org/opensearch/search/relevance/transformer/ResultTransformer.java

@@ -23,18 +23,29 @@ public interface ResultTransformer {
 
   /**
    * Decide whether to apply the transformer on the input request
-   * @param request input request
+   * @param request input Search Request
    * @param configuration Configuration parameters for the transformer
    * @return boolean decision on whether to apply the transformer
    */
   boolean shouldTransform(final SearchRequest request, final ResultTransformerConfiguration configuration);
 
+  /**
+   * Preprocess the incoming Search Request to support the requirements of the transformer
+   * @param request input Search Request
+   * @param configuration Configuration parameters for the transformer
+   * @return SearchRequest with updated attributes
+   */
+  SearchRequest preprocessRequest(final SearchRequest request,
+      final ResultTransformerConfiguration configuration);
+
   /**
    * Rank hits based on the provided query
    * @param hits hits to be re-ranked
    * @param request Search request
    * @param configuration Configuration parameters for the transformer
    * @return SearchHits ordered by score generated by ranker
    */
-  SearchHits transform(final SearchHits hits, final SearchRequest request, final ResultTransformerConfiguration configuration);
+  SearchHits transform(final SearchHits hits,
+      final SearchRequest request,
+      final ResultTransformerConfiguration configuration);
 }

src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/Constants.java

@@ -17,6 +17,7 @@ public class Constants {
   // Transformer properties
   public static final String BODY_FIELD = "body_field";
   public static final String TITLE_FIELD = "title_field";
+  public static final String DOC_LIMIT = "doc_limit";
 
   public static final String KENDRA_SETTINGS_PREFIX =
       String.join(".", RESULT_TRANSFORMER_SETTING_PREFIX, KENDRA_INTELLIGENT_RANKING);
@@ -27,5 +28,8 @@ public class Constants {
       String.join(".", KENDRA_SETTINGS_PREFIX, PROPERTIES, BODY_FIELD);
   public static final String TITLE_FIELD_SETTING_NAME =
       String.join(".", KENDRA_SETTINGS_PREFIX, PROPERTIES, TITLE_FIELD);
+  public static final String DOC_LIMIT_SETTING_NAME =
+      String.join(".", KENDRA_SETTINGS_PREFIX, PROPERTIES, DOC_LIMIT);
 
+  public static final int KENDRA_DEFAULT_DOC_LIMIT = 25;
 }

src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankerSettings.java

@@ -15,7 +15,6 @@
 import org.opensearch.common.settings.SecureString;
 import org.opensearch.common.settings.Setting;
 import org.opensearch.common.settings.Setting.Property;
-import org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants;
 
 public class KendraIntelligentRankerSettings {
 
@@ -26,24 +25,13 @@ public class KendraIntelligentRankerSettings {
       Property.Dynamic, Property.IndexScope);
 
   /**
-   * Document field to be considered as "body" when invoking Kendra.
-   */
-  public static final Setting<List<String>> KENDRA_BODY_FIELD_SETTING = Setting.listSetting(Constants.BODY_FIELD_SETTING_NAME, Collections.emptyList(),
-      Function.identity(), new FieldSettingValidator(Constants.BODY_FIELD_SETTING_NAME),
-      Property.Dynamic, Property.IndexScope);
-
-  /**
-   * Document field to be considered as "title" when invoking Kendra.
+   * Validator for body and title field settings
    */
-  public static final Setting<List<String>> KENDRA_TITLE_FIELD_SETTING = Setting.listSetting(Constants.TITLE_FIELD_SETTING_NAME, Collections.emptyList(),
-      Function.identity(), new FieldSettingValidator(Constants.TITLE_FIELD_SETTING_NAME),
-      Property.Dynamic, Property.IndexScope);
-
-  static final class FieldSettingValidator implements Setting.Validator<List<String>> {
+  static final class FieldValidator implements Setting.Validator<List<String>> {
 
     private String settingName;
 
-    public FieldSettingValidator(final String name) {
+    public FieldValidator(final String name) {
       this.settingName = name;
     }
 
@@ -55,6 +43,55 @@ public void validate(List<String> value) {
     }
   }
 
+  /**
+   * Validator for doc limit setting
+   */
+  static final class DocLimitValidator implements Setting.Validator<Integer> {
+
+    private String settingName;
+
+    public DocLimitValidator(final String name) {
+      this.settingName = name;
+    }
+
+    @Override
+    public void validate(Integer value) {
+      if (value != null && value < Constants.KENDRA_DEFAULT_DOC_LIMIT) {
+        throw new IllegalArgumentException("Setting the value of [" + this.settingName + "] below "
+            + Constants.KENDRA_DEFAULT_DOC_LIMIT + " will affect ranking accuracy");
+      }
+    }
+  }
+
+  /**
+   * Validator objects
+   */
+  public static final FieldValidator BODY_FIELD_VALIDATOR = new FieldValidator(Constants.BODY_FIELD);
+  public static final FieldValidator TITLE_FIELD_VALIDATOR = new FieldValidator(Constants.TITLE_FIELD);
+  public static final DocLimitValidator DOC_LIMIT_VALIDATOR = new DocLimitValidator(Constants.DOC_LIMIT);
+
+  /**
+   * Document field to be considered as "body" when invoking Kendra.
+   */
+  public static final Setting<List<String>> KENDRA_BODY_FIELD_SETTING = Setting.listSetting(Constants.BODY_FIELD_SETTING_NAME,
+      Collections.emptyList(), Function.identity(), BODY_FIELD_VALIDATOR,
+      Property.Dynamic, Property.IndexScope);
+
+  /**
+   * Document field to be considered as "title" when invoking Kendra.
+   */
+  public static final Setting<List<String>> KENDRA_TITLE_FIELD_SETTING = Setting.listSetting(Constants.TITLE_FIELD_SETTING_NAME,
+      Collections.emptyList(), Function.identity(), TITLE_FIELD_VALIDATOR,
+      Property.Dynamic, Property.IndexScope);
+
+
+
+  public static final Setting<Integer> KENDRA_DOC_LIMIT_SETTING = Setting.intSetting(
+      Constants.DOC_LIMIT_SETTING_NAME, Constants.KENDRA_DEFAULT_DOC_LIMIT, 1,
+      DOC_LIMIT_VALIDATOR, Property.Dynamic, Property.IndexScope);
+
+
+
   /**
    * The access key (ie login id) for connecting to Kendra.
    */
@@ -83,6 +120,7 @@ public static final List<Setting<?>> getAllSettings() {
       KENDRA_ORDER_SETTING,
       KENDRA_BODY_FIELD_SETTING,
       KENDRA_TITLE_FIELD_SETTING,
+      KENDRA_DOC_LIMIT_SETTING,
       ACCESS_KEY_SETTING,
       SECRET_KEY_SETTING,
       SESSION_TOKEN_SETTING,

src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/configuration/KendraIntelligentRankingConfiguration.java

@@ -8,8 +8,14 @@
 package org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration;
 
 import static org.opensearch.search.relevance.configuration.Constants.ORDER;
+import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants.DOC_LIMIT_SETTING_NAME;
+import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.Constants.KENDRA_DEFAULT_DOC_LIMIT;
+import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankerSettings.BODY_FIELD_VALIDATOR;
+import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankerSettings.DOC_LIMIT_VALIDATOR;
+import static org.opensearch.search.relevance.transformer.kendraintelligentranking.configuration.KendraIntelligentRankerSettings.TITLE_FIELD_VALIDATOR;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
 import org.opensearch.common.ParseField;
@@ -55,7 +61,8 @@ public KendraIntelligentRankingConfiguration(Settings settings) {
     this.order = settings.getAsInt(ORDER, 0);
     this.properties = new KendraIntelligentRankingProperties(
         settings.getAsList("properties.body_field"),
-        settings.getAsList("properties.title_field"));
+        settings.getAsList("properties.title_field"),
+        settings.getAsInt("properties.doc_limit", KENDRA_DEFAULT_DOC_LIMIT));
   }
 
   @Override
@@ -117,48 +124,54 @@ public KendraIntelligentRankingProperties getProperties() {
   public static class KendraIntelligentRankingProperties implements Writeable, ToXContentObject {
     protected static final ParseField BODY_FIELD = new ParseField(Constants.BODY_FIELD);
     protected static final ParseField TITLE_FIELD = new ParseField(Constants.TITLE_FIELD);
+    protected static final ParseField DOC_LIMIT = new ParseField(Constants.DOC_LIMIT);
 
     private static final ObjectParser<KendraIntelligentRankingProperties, Void> PARSER;
 
     static {
       PARSER = new ObjectParser<>("kendra_intelligent_ranking_configuration", KendraIntelligentRankingProperties::new);
       PARSER.declareStringArray(KendraIntelligentRankingProperties::setBodyFields, BODY_FIELD);
       PARSER.declareStringArray(KendraIntelligentRankingProperties::setTitleFields, TITLE_FIELD);
+      PARSER.declareInt(KendraIntelligentRankingProperties::setDocLimit, DOC_LIMIT);
     }
 
     private List<String> bodyFields;
     private List<String> titleFields;
+    private int docLimit;
 
-    public KendraIntelligentRankingProperties() {}
+    public KendraIntelligentRankingProperties() {
+      bodyFields = Collections.emptyList();
+      titleFields = Collections.emptyList();
+      docLimit = KENDRA_DEFAULT_DOC_LIMIT;
+    }
 
-    public KendraIntelligentRankingProperties(final List<String> bodyFields, final List<String> titleFields) {
+    public KendraIntelligentRankingProperties(final List<String> bodyFields,
+        final List<String> titleFields, final int docLimit) {
       this.bodyFields = bodyFields;
       this.titleFields = titleFields;
+      this.docLimit = docLimit;
     }
 
     public KendraIntelligentRankingProperties(StreamInput input) throws IOException {
       this.bodyFields = input.readStringList();
       this.bodyFields = input.readStringList();
+      this.docLimit = input.readInt();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
       out.writeStringCollection(this.bodyFields);
       out.writeStringCollection(this.titleFields);
+      out.writeInt(this.docLimit);
     }
 
     public static KendraIntelligentRankingProperties parse(XContentParser parser, Void context) throws IOException {
       try {
         KendraIntelligentRankingProperties properties = PARSER.parse(parser, null);
         if (properties != null) {
-          if (properties.getBodyFields() != null && !properties.getBodyFields().isEmpty() && properties.getBodyFields().size() > 1) {
-            throw new ParsingException(parser.getTokenLocation(),
-                "[" + BODY_FIELD + "] can have at most 1 element");
-          }
-          if (properties.getTitleFields() != null && !properties.getTitleFields().isEmpty() && properties.getTitleFields().size() > 1) {
-            throw new ParsingException(parser.getTokenLocation(),
-                "[" + TITLE_FIELD + "] can have at most 1 element");
-          }
+          BODY_FIELD_VALIDATOR.validate(properties.getBodyFields());
+          TITLE_FIELD_VALIDATOR.validate(properties.getTitleFields());
+          DOC_LIMIT_VALIDATOR.validate(properties.getDocLimit());
         }
         return properties;
       } catch (IllegalArgumentException iae) {
@@ -171,6 +184,7 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
       builder.startObject();
       builder.field(BODY_FIELD.getPreferredName(), this.bodyFields);
       builder.field(TITLE_FIELD.getPreferredName(), this.titleFields);
+      builder.field(DOC_LIMIT.getPreferredName(), this.docLimit);
       return builder.endObject();
     }
 
@@ -181,8 +195,8 @@ public boolean equals(Object o) {
 
       KendraIntelligentRankingProperties properties = (KendraIntelligentRankingProperties) o;
 
-      if (bodyFields != properties.bodyFields) return false;
-      return (titleFields == properties.titleFields);
+      return (bodyFields == properties.bodyFields) && (titleFields == properties.titleFields) &&
+          (docLimit == properties.docLimit);
     }
 
     @Override
@@ -205,5 +219,13 @@ public List<String> getTitleFields() {
     public void setTitleFields(final List<String> titleFields) {
       this.titleFields = titleFields;
     }
+
+    public int getDocLimit() {
+      return this.docLimit;
+    }
+
+    public void setDocLimit(final int docLimit) {
+      this.docLimit = docLimit;
+    }
   }
 }

src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/KendraIntelligentRankingException.java

@@ -0,0 +1,30 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.search.relevance.transformer.kendraintelligentranking.model;
+
+import java.io.IOException;
+import org.opensearch.OpenSearchException;
+import org.opensearch.common.io.stream.StreamInput;
+
+public class KendraIntelligentRankingException extends OpenSearchException {
+  public KendraIntelligentRankingException(StreamInput in) throws IOException {
+    super(in);
+  }
+
+  public KendraIntelligentRankingException(String message) {
+    super(message);
+  }
+
+  public KendraIntelligentRankingException(String message, Throwable cause) {
+    super(message, cause);
+  }
+
+  public KendraIntelligentRankingException(String message, Throwable cause, Object... args) {
+    super(message, cause, args);
+  }
+}

src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResult.java

@@ -24,4 +24,21 @@ public String getRescoreId() {
   public List<RescoreResultItem> getResultItems() {
     return resultItems;
   }
+
+  /**
+   * Setter used for unit tests.
+   * @param rescoreId The identifier associated with the scores that Amazon Kendra Intelligent Ranking
+   *                  gives to the results.
+   */
+  public void setRescoreId(String rescoreId) {
+    this.rescoreId = rescoreId;
+  }
+
+  /**
+   * Setter used for unit tests.
+   * @param resultItems A list of result items for documents with new relevancy scores.
+   */
+  public void setResultItems(List<RescoreResultItem> resultItems) {
+    this.resultItems = resultItems;
+  }
 }

src/main/java/org/opensearch/search/relevance/transformer/kendraintelligentranking/model/dto/RescoreResultItem.java

@@ -22,4 +22,20 @@ public String getDocumentId() {
   public Float getScore() {
     return score;
   }
+
+  /**
+   * Setter for unit tests.
+   * @param documentId the ID of the rescored document.
+   */
+  public void setDocumentId(String documentId) {
+    this.documentId = documentId;
+  }
+
+  /**
+   * Setter for unit tests.
+   * @param score the updated score of the document.
+   */
+  public void setScore(Float score) {
+    this.score = score;
+  }
 }

src/test/java/org/opensearch/search/relevance/SearchRelevancePluginIT.java

@@ -7,33 +7,21 @@
  */
 package org.opensearch.search.relevance;
 
-import com.carrotsearch.randomizedtesting.annotations.ThreadLeakScope;
 import org.apache.http.util.EntityUtils;
 import org.opensearch.client.Request;
 import org.opensearch.client.Response;
-import org.opensearch.plugins.Plugin;
-import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.test.rest.OpenSearchRestTestCase;
 
 import java.io.IOException;
-import java.util.Collection;
-import java.util.Collections;
 
-import static org.hamcrest.Matchers.containsString;
-
-@ThreadLeakScope(ThreadLeakScope.Scope.NONE)
-@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE)
-public class SearchRelevancePluginIT extends OpenSearchIntegTestCase {
-
-    @Override
-    protected Collection<Class<? extends Plugin>> nodePlugins() {
-        return Collections.singletonList(SearchRelevancePlugin.class);
-    }
+public class SearchRelevancePluginIT extends OpenSearchRestTestCase {
 
     public void testPluginInstalled() throws IOException {
-        Response response = createRestClient().performRequest(new Request("GET", "/_cat/plugins"));
+        Response response = client().performRequest(new Request("GET", "/_cat/plugins"));
         String body = EntityUtils.toString(response.getEntity());
 
         logger.info("response body: {}", body);
-        assertThat(body, containsString("search-processor"));
+        assertNotNull(body);
+        assertTrue(body.contains("search-processor"));
     }
 }