Add rerank to the unit tests

# Conflicts:
#	x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/sagemaker/SageMakerServiceTests.java

Author: davidkyle

server/src/main/java/org/elasticsearch/inference/RerankingInferenceService.java

@@ -11,10 +11,17 @@
 
 public interface RerankingInferenceService {
 
+    /**
+     * The default window size for small reranking models.
+     */
+    int CONSERVATIVE_DEFAULT_WINDOW_SIZE = 250;
+    int LARGE_WINDOW_SIZE = 500;
+
     /**
      * The reranking model's max window or an approximation of
      * measured in the number of words.
+     * @param modelId The model ID
      * @return Window size in words
      */
-    int rerankerWindowSize();
+    int rerankerWindowSize(String modelId);
 }

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

@@ -26,6 +26,7 @@
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.RerankingInferenceService;
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
@@ -72,7 +73,7 @@
 import static org.elasticsearch.xpack.inference.services.azureaistudio.completion.AzureAiStudioChatCompletionTaskSettings.DEFAULT_MAX_NEW_TOKENS;
 import static org.elasticsearch.xpack.inference.services.openai.OpenAiServiceFields.EMBEDDING_MAX_BATCH_SIZE;
 
-public class AzureAiStudioService extends SenderService {
+public class AzureAiStudioService extends SenderService implements RerankingInferenceService {
 
     public static final String NAME = "azureaistudio";
 
@@ -400,6 +401,11 @@ private static void checkProviderAndEndpointTypeForTask(
         }
     }
 
+    @Override
+    public int rerankerWindowSize(String modelId) {
+        return RerankingInferenceService.LARGE_WINDOW_SIZE;
+    }
+
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.RerankingInferenceService;
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
@@ -66,7 +67,7 @@
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwUnsupportedUnifiedCompletionOperation;
 import static org.elasticsearch.xpack.inference.services.cohere.CohereServiceFields.EMBEDDING_MAX_BATCH_SIZE;
 
-public class CohereService extends SenderService {
+public class CohereService extends SenderService implements RerankingInferenceService {
     public static final String NAME = "cohere";
 
     private static final String SERVICE_NAME = "Cohere";
@@ -361,6 +362,12 @@ public Set<TaskType> supportedStreamingTasks() {
         return COMPLETION_ONLY;
     }
 
+    @Override
+    public int rerankerWindowSize(String modelId) {
+        // Cohere rerank model truncates at 4096 tokens https://docs.cohere.com/reference/rerank
+        return RerankingInferenceService.LARGE_WINDOW_SIZE;
+    }
+
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();

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

@@ -31,6 +31,7 @@
 import org.elasticsearch.inference.MinimalServiceSettings;
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.RerankingInferenceService;
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.inference.UnifiedCompletionRequest;
@@ -84,7 +85,7 @@
 import static org.elasticsearch.xpack.inference.services.elasticsearch.ElserModels.ELSER_V2_MODEL;
 import static org.elasticsearch.xpack.inference.services.elasticsearch.ElserModels.ELSER_V2_MODEL_LINUX_X86;
 
-public class ElasticsearchInternalService extends BaseElasticsearchInternalService {
+public class ElasticsearchInternalService extends BaseElasticsearchInternalService implements RerankingInferenceService {
 
     public static final String NAME = "elasticsearch";
     public static final String OLD_ELSER_SERVICE_NAME = "elser";
@@ -1060,6 +1061,12 @@ static TaskType inferenceConfigToTaskType(InferenceConfig config) {
         }
     }
 
+    @Override
+    public int rerankerWindowSize(String modelId) {
+        // TODO rerank chunking should use the same value
+        return RerankingInferenceService.CONSERVATIVE_DEFAULT_WINDOW_SIZE;
+    }
+
     /**
      * Iterates over the batch executing a limited number requests at a time to avoid
      * filling the ML node inference queue.

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/googlevertexai/GoogleVertexAiService.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.RerankingInferenceService;
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
@@ -69,7 +70,7 @@
 import static org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiServiceFields.PROJECT_ID;
 import static org.elasticsearch.xpack.inference.services.googlevertexai.action.GoogleVertexAiActionCreator.COMPLETION_ERROR_PREFIX;
 
-public class GoogleVertexAiService extends SenderService {
+public class GoogleVertexAiService extends SenderService implements RerankingInferenceService {
 
     public static final String NAME = "googlevertexai";
 
@@ -383,6 +384,11 @@ private static GoogleVertexAiModel createModel(
         };
     }
 
+    @Override
+    public int rerankerWindowSize(String modelId) {
+        return 0; // TODO
+    }
+
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/voyageai/VoyageAIService.java

@@ -25,6 +25,7 @@
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
+import org.elasticsearch.inference.RerankingInferenceService;
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
@@ -61,7 +62,7 @@
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwIfNotEmptyMap;
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.throwUnsupportedUnifiedCompletionOperation;
 
-public class VoyageAIService extends SenderService {
+public class VoyageAIService extends SenderService implements RerankingInferenceService {
     public static final String NAME = "voyageai";
 
     private static final String SERVICE_NAME = "Voyage AI";
@@ -369,6 +370,12 @@ public TransportVersion getMinimalSupportedVersion() {
         return TransportVersions.VOYAGE_AI_INTEGRATION_ADDED;
     }
 
+    @Override
+    public int rerankerWindowSize(String modelId) {
+        https://docs.voyageai.com/reference/reranker-api
+        return RerankingInferenceService.LARGE_WINDOW_SIZE;
+    }
+
     public static class Configuration {
         public static InferenceServiceConfiguration get() {
             return configuration.getOrCompute();

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/AbstractInferenceServiceTests.java

@@ -17,7 +17,6 @@
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
-import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.http.MockWebServer;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
@@ -52,7 +51,7 @@
  * To use this class, extend it and pass the constructor a configuration.
  * </p>
  */
-public abstract class AbstractInferenceServiceTests extends ESTestCase {
+public abstract class AbstractInferenceServiceTests extends InferenceServiceTestCase {
 
     protected final MockWebServer webServer = new MockWebServer();
     protected ThreadPool threadPool;

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/InferenceServiceTestCase.java

@@ -8,9 +8,47 @@
 package org.elasticsearch.xpack.inference.services;
 
 import org.elasticsearch.inference.InferenceService;
+import org.elasticsearch.inference.RerankingInferenceService;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.test.ESTestCase;
 
+import java.io.IOException;
+
 public abstract class InferenceServiceTestCase extends ESTestCase {
 
     public abstract InferenceService createInferenceService();
+
+    public void testRerankersImplementRerankInterface() throws IOException {
+        try (InferenceService inferenceService = createInferenceService()) {
+            boolean implementsReranking = inferenceService instanceof RerankingInferenceService;
+            boolean hasRerankTaskType = inferenceService.supportedTaskTypes().contains(TaskType.RERANK);
+            if (implementsReranking != hasRerankTaskType) {
+                fail(
+                    "Reranking inference services should implement RerankingInferenceService and support the RERANK task type. "
+                        + "Service ["
+                        + inferenceService.name()
+                        + "] supports task type: ["
+                        + hasRerankTaskType
+                        + "] and implements"
+                        + " RerankingInferenceService: ["
+                        + implementsReranking
+                        + "]"
+                );
+            }
+        }
+    }
+
+    public void testRerankersHaveWindowSize() throws IOException {
+        try (InferenceService inferenceService = createInferenceService()) {
+            if (inferenceService instanceof RerankingInferenceService rerankingInferenceService) {
+                assertRerankerWindowSize(rerankingInferenceService);
+            }
+        }
+    }
+
+    protected void assertRerankerWindowSize(
+        RerankingInferenceService rerankingInferenceService
+    ) {
+        fail("Reranking services should override this test method to verify window size");
+    }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/alibabacloudsearch/AlibabaCloudSearchServiceTests.java

@@ -19,14 +19,14 @@
 import org.elasticsearch.inference.ChunkInferenceInput;
 import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.inference.ChunkingSettings;
+import org.elasticsearch.inference.InferenceService;
 import org.elasticsearch.inference.InferenceServiceConfiguration;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.InputType;
 import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
-import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentType;
@@ -42,6 +42,7 @@
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests;
 import org.elasticsearch.xpack.inference.logging.ThrottlerManager;
+import org.elasticsearch.xpack.inference.services.InferenceServiceTestCase;
 import org.elasticsearch.xpack.inference.services.ServiceFields;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.action.AlibabaCloudSearchActionVisitor;
 import org.elasticsearch.xpack.inference.services.alibabacloudsearch.embeddings.AlibabaCloudSearchEmbeddingsModel;
@@ -73,7 +74,7 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.mockito.Mockito.mock;
 
-public class AlibabaCloudSearchServiceTests extends ESTestCase {
+public class AlibabaCloudSearchServiceTests extends InferenceServiceTestCase {
     private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS);
     private ThreadPool threadPool;
     private HttpClientManager clientManager;
@@ -710,4 +711,9 @@ private Map<String, Object> getRequestConfigMap(
             Map.of(ModelConfigurations.SERVICE_SETTINGS, builtServiceSettings, ModelConfigurations.TASK_SETTINGS, taskSettings)
         );
     }
+
+    @Override
+    public InferenceService createInferenceService() {
+        return new AlibabaCloudSearchService(mock(HttpRequestSender.Factory.class), createWithEmptySettings(threadPool));
+    }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/amazonbedrock/AmazonBedrockServiceTests.java

@@ -23,6 +23,7 @@
 import org.elasticsearch.inference.ChunkInferenceInput;
 import org.elasticsearch.inference.ChunkedInference;
 import org.elasticsearch.inference.ChunkingSettings;
+import org.elasticsearch.inference.InferenceService;
 import org.elasticsearch.inference.InferenceServiceConfiguration;
 import org.elasticsearch.inference.InferenceServiceResults;
 import org.elasticsearch.inference.InputType;
@@ -31,7 +32,6 @@
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.SimilarityMeasure;
 import org.elasticsearch.inference.TaskType;
-import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xcontent.ToXContent;
 import org.elasticsearch.xcontent.XContentType;
@@ -43,6 +43,7 @@
 import org.elasticsearch.xpack.inference.common.amazon.AwsSecretSettings;
 import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.services.InferenceServiceTestCase;
 import org.elasticsearch.xpack.inference.services.ServiceComponentsTests;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.client.AmazonBedrockMockRequestSender;
 import org.elasticsearch.xpack.inference.services.amazonbedrock.completion.AmazonBedrockChatCompletionModel;
@@ -92,7 +93,7 @@
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
-public class AmazonBedrockServiceTests extends ESTestCase {
+public class AmazonBedrockServiceTests extends InferenceServiceTestCase {
     private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS);
     private ThreadPool threadPool;
 
@@ -1405,6 +1406,11 @@ private AmazonBedrockService createAmazonBedrockService() {
         );
     }
 
+    @Override
+    public InferenceService createInferenceService() {
+        return createAmazonBedrockService();
+    }
+
     private Map<String, Object> getRequestConfigMap(
         Map<String, Object> serviceSettings,
         Map<String, Object> taskSettings,